use Effect hashmap for current active segment tracking

Author: yujonglee

apps/desktop/src/utils/segment.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "vitest";
-import { buildSegments } from "./segment";
+import { buildSegments, SegmentKey } from "./segment";
 
 describe("buildSegments", () => {
   const testCases = [
@@ -21,14 +21,14 @@ describe("buildSegments", () => {
       ],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [
             expect.objectContaining({ text: "1", isFinal: true }),
             expect.objectContaining({ text: "2", isFinal: false }),
           ],
         }),
         expect.objectContaining({
-          key: { channel: 1 },
+          key: SegmentKey.Channel({ channel: 1 }),
           words: [
             expect.objectContaining({ text: "3", isFinal: false }),
             expect.objectContaining({ text: "4", isFinal: false }),
@@ -47,14 +47,14 @@ describe("buildSegments", () => {
       ],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [
             expect.objectContaining({ text: "1" }),
             expect.objectContaining({ text: "2" }),
           ],
         }),
         expect.objectContaining({
-          key: { channel: 1 },
+          key: SegmentKey.Channel({ channel: 1 }),
           words: [
             expect.objectContaining({ text: "3" }),
           ],
@@ -72,7 +72,7 @@ describe("buildSegments", () => {
       ],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [
             expect.objectContaining({ text: "1" }),
             expect.objectContaining({ text: "2" }),
@@ -91,15 +91,15 @@ describe("buildSegments", () => {
       partialWords: [],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [expect.objectContaining({ text: "first" })],
         }),
         expect.objectContaining({
-          key: { channel: 1 },
+          key: SegmentKey.Channel({ channel: 1 }),
           words: [expect.objectContaining({ text: "other" })],
         }),
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [expect.objectContaining({ text: "return" })],
         }),
       ],
@@ -113,7 +113,7 @@ describe("buildSegments", () => {
       partialWords: [],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [
             expect.objectContaining({ text: "first" }),
             expect.objectContaining({ text: "second" }),
@@ -130,9 +130,18 @@ describe("buildSegments", () => {
       ],
       partialWords: [],
       expected: [
-        expect.objectContaining({ key: { channel: 0 }, words: [expect.objectContaining({ text: "a" })] }),
-        expect.objectContaining({ key: { channel: 1 }, words: [expect.objectContaining({ text: "b" })] }),
-        expect.objectContaining({ key: { channel: 2 }, words: [expect.objectContaining({ text: "c" })] }),
+        expect.objectContaining({
+          key: SegmentKey.Channel({ channel: 0 }),
+          words: [expect.objectContaining({ text: "a" })],
+        }),
+        expect.objectContaining({
+          key: SegmentKey.Channel({ channel: 1 }),
+          words: [expect.objectContaining({ text: "b" })],
+        }),
+        expect.objectContaining({
+          key: SegmentKey.Channel({ channel: 2 }),
+          words: [expect.objectContaining({ text: "c" })],
+        }),
       ],
     },
     {
@@ -143,7 +152,7 @@ describe("buildSegments", () => {
       partialWords: [],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [expect.objectContaining({ text: "only", isFinal: true })],
         }),
       ],
@@ -160,15 +169,15 @@ describe("buildSegments", () => {
       partialWords: [],
       expected: [
         expect.objectContaining({
-          key: { channel: 0 },
+          key: SegmentKey.Channel({ channel: 0 }),
           words: [
             expect.objectContaining({ text: "a1" }),
             expect.objectContaining({ text: "a2" }),
             expect.objectContaining({ text: "a3" }),
           ],
         }),
         expect.objectContaining({
-          key: { channel: 1 },
+          key: SegmentKey.Channel({ channel: 1 }),
           words: [
             expect.objectContaining({ text: "b1" }),
             expect.objectContaining({ text: "b2" }),

apps/desktop/src/utils/segment.ts

@@ -1,3 +1,5 @@
+import { Data, HashMap, Option } from "effect";
+
 export type WordLike = {
   text: string;
   start_ms: number;
@@ -14,7 +16,11 @@ export type Segment<TWord extends SegmentWord = SegmentWord> = {
   words: TWord[];
 };
 
-type SegmentKey = { channel: number } | { speaker_index: number; channel: number };
+export type SegmentKey = Data.TaggedEnum<{
+  Channel: { channel: number };
+}>;
+
+export const SegmentKey = Data.taggedEnum<SegmentKey>();
 
 export function buildSegments<
   TFinal extends WordLike,
@@ -38,41 +44,41 @@ export function buildSegments<
       channel: word.channel,
       isFinal: false,
     })),
-  ];
+  ].sort((a, b) => a.start_ms - b.start_ms);
 
   return createSpeakerTurns(allWords);
 }
 
-function createSpeakerTurns<TWord extends SegmentWord>(
-  words: TWord[],
-  maxGapMs = 2000,
-): Segment<TWord>[] {
+function createSpeakerTurns<TWord extends SegmentWord>(words: TWord[]): Segment<TWord>[] {
+  const MAX_GAP_MS = 2000;
+
   if (words.length === 0) {
     return [];
   }
 
-  const sortedWords = [...words].sort((a, b) => a.start_ms - b.start_ms);
   const segments: Segment<TWord>[] = [];
-  const currentByChannel = new Map<number, Segment<TWord>>();
+  let currentActiveSegment = HashMap.empty<SegmentKey, Segment<TWord>>();
 
-  for (const word of sortedWords) {
-    const current = currentByChannel.get(word.channel);
+  for (const word of words) {
+    const key = SegmentKey.Channel({ channel: word.channel });
+    const currentOption = HashMap.get(currentActiveSegment, key);
 
-    if (!current) {
-      const newSegment = { key: { channel: word.channel }, words: [word] };
-      currentByChannel.set(word.channel, newSegment);
+    if (Option.isNone(currentOption)) {
+      const newSegment = { key, words: [word] };
+      currentActiveSegment = HashMap.set(currentActiveSegment, key, newSegment);
       segments.push(newSegment);
       continue;
     }
 
+    const current = currentOption.value;
     const lastWord = current.words[current.words.length - 1];
     const gap = word.start_ms - lastWord.end_ms;
 
-    if (gap <= maxGapMs) {
+    if (gap <= MAX_GAP_MS) {
       current.words.push(word);
     } else {
-      const newSegment = { key: { channel: word.channel }, words: [word] };
-      currentByChannel.set(word.channel, newSegment);
+      const newSegment = { key, words: [word] };
+      currentActiveSegment = HashMap.set(currentActiveSegment, key, newSegment);
       segments.push(newSegment);
     }
   }