@typed/fx 1.12.2 → 1.12.3

package/src/keyed.ts

@@ -1,7 +1,7 @@
 import * as MutableHashMap from '@effect/data/MutableHashMap'
 import * as Option from '@effect/data/Option'
 import * as ReadonlyArray from '@effect/data/ReadonlyArray'
-import { Equivalence } from '@effect/data/typeclass/Equivalence'
+import * as Equivalence from '@effect/data/typeclass/Equivalence'
 import * as Effect from '@effect/io/Effect'
 import * as Fiber from '@effect/io/Fiber'
 import fastDeepEqual from 'fast-deep-equal'
@@ -12,18 +12,19 @@ import { Subject, makeHoldSubject } from './Subject.js'
 import { Cause } from './externals.js'
 import { ScopedFork, withScopedFork } from './helpers.js'
 
-export function keyed<R, E, A, R2, E2, B>(
+export function keyed<R, E, A, R2, E2, B, C>(
   fx: Fx<R, E, readonly A[]>,
-  f: (a: Fx<never, never, A>, initial: A) => Fx<R2, E2, B>,
-  eq: Equivalence<A> = fastDeepEqual,
+  f: (fx: RefSubject<never, A>) => Fx<R2, E2, B>,
+  getKey: (a: A) => C,
 ): Fx<R | R2, E | E2, readonly B[]> {
   return Fx(<R3>(sink: Sink<R3, E | E2, readonly B[]>) =>
     withScopedFork((fork) =>
       Effect.gen(function* ($) {
-        const state = createKeyedState<A, B>()
+        const state = createKeyedState<A, B, C>()
+        const eq = Equivalence.make((x: A, y: A) => fastDeepEqual(getKey(x), getKey(y)))
         const difference = ReadonlyArray.difference(eq)
         const intersection = ReadonlyArray.intersection(eq)
-        const emit = emitWhenReady(state)
+        const emit = emitWhenReady(state, getKey)
 
         // Let output emit to the sink
         const fiber = yield* $(fork(state.output.run(sink)))
@@ -36,6 +37,7 @@ export function keyed<R, E, A, R2, E2, B>(
                 updateState({
                   state,
                   updated: as,
+                  getKey,
                   f,
                   fork,
                   difference,
@@ -60,17 +62,17 @@ export function keyed<R, E, A, R2, E2, B>(
   )
 }
 
-type KeyedState<A, B> = {
+type KeyedState<A, B, C> = {
   previous: readonly A[]
   ended: boolean
 
-  readonly subjects: MutableHashMap.MutableHashMap<A, Subject<never, A>>
-  readonly fibers: MutableHashMap.MutableHashMap<A, Fiber.RuntimeFiber<never, void>>
-  readonly values: MutableHashMap.MutableHashMap<A, B>
+  readonly subjects: MutableHashMap.MutableHashMap<C, Subject<never, A>>
+  readonly fibers: MutableHashMap.MutableHashMap<C, Fiber.RuntimeFiber<never, void>>
+  readonly values: MutableHashMap.MutableHashMap<C, B>
   readonly output: Subject<never, readonly B[]>
 }
 
-function createKeyedState<A, B>(): KeyedState<A, B> {
+function createKeyedState<A, B, C>(): KeyedState<A, B, C> {
   return {
     previous: [],
     ended: false,
@@ -81,7 +83,7 @@ function createKeyedState<A, B>(): KeyedState<A, B> {
   }
 }
 
-function updateState<A, B, R2, E2, R3>({
+function updateState<A, B, C, R2, E2, R3>({
   state,
   updated,
   f,
@@ -90,15 +92,17 @@ function updateState<A, B, R2, E2, R3>({
   intersection,
   emit,
   error,
+  getKey,
 }: {
-  state: KeyedState<A, B>
+  state: KeyedState<A, B, C>
   updated: readonly A[]
-  f: (a: Fx<never, never, A>, initial: A) => Fx<R2, E2, B>
+  f: (fx: RefSubject<never, A>) => Fx<R2, E2, B>
   fork: ScopedFork
   difference: (self: Iterable<A>, that: Iterable<A>) => A[]
   intersection: (self: Iterable<A>, that: Iterable<A>) => A[]
   emit: Effect.Effect<never, never, void>
   error: (e: Cause.Cause<E2>) => Effect.Effect<R3, never, void>
+  getKey: (a: A) => C
 }) {
   return Effect.gen(function* ($) {
     const added = difference(updated, state.previous)
@@ -108,15 +112,17 @@ function updateState<A, B, R2, E2, R3>({
     state.previous = updated
 
     // Remove values that are no longer in the stream
-    yield* $(Effect.forEachParDiscard(removed, (value) => removeValue(state, value)))
+    yield* $(Effect.forEachParDiscard(removed, (value) => removeValue(state, value, getKey)))
 
     // Add values that are new to the stream
     yield* $(
-      Effect.forEachParDiscard(added, (value) => addValue({ state, value, f, fork, emit, error })),
+      Effect.forEachParDiscard(added, (value) =>
+        addValue({ state, value, f, fork, emit, error, getKey }),
+      ),
     )
 
     // Update values that are still in the stream
-    yield* $(Effect.forEachParDiscard(unchanged, (value) => updateValue(state, value)))
+    yield* $(Effect.forEachParDiscard(unchanged, (value) => updateValue(state, value, getKey)))
 
     // If nothing was added or removed, emit the current values
     if (added.length === 0 && removed.length === 0) {
@@ -125,47 +131,51 @@ function updateState<A, B, R2, E2, R3>({
   })
 }
 
-function removeValue<A, B>(state: KeyedState<A, B>, value: A) {
+function removeValue<A, B, C>(state: KeyedState<A, B, C>, value: A, getKey: (a: A) => C) {
   return Effect.gen(function* ($) {
-    const subject = MutableHashMap.get(state.subjects, value)
+    const key = getKey(value)
+    const subject = MutableHashMap.get(state.subjects, key)
 
     if (Option.isSome(subject)) yield* $(subject.value.end())
 
-    const fiber = MutableHashMap.get(state.fibers, value)
+    const fiber = MutableHashMap.get(state.fibers, key)
 
     if (Option.isSome(fiber)) yield* $(Fiber.interrupt(fiber.value))
 
-    MutableHashMap.remove(state.values, value)
-    MutableHashMap.remove(state.subjects, value)
-    MutableHashMap.remove(state.fibers, value)
+    MutableHashMap.remove(state.values, key)
+    MutableHashMap.remove(state.subjects, key)
+    MutableHashMap.remove(state.fibers, key)
   })
 }
 
-function addValue<A, B, R2, E2, R3>({
+function addValue<A, B, C, R2, E2, R3>({
   state,
   value,
   f,
   fork,
   emit,
   error,
+  getKey,
 }: {
-  state: KeyedState<A, B>
+  state: KeyedState<A, B, C>
   value: A
-  f: (a: Fx<never, never, A>, initial: A) => Fx<R2, E2, B>
+  f: (fx: RefSubject<never, A>) => Fx<R2, E2, B>
   fork: ScopedFork
   emit: Effect.Effect<never, never, void>
   error: (e: Cause.Cause<E2>) => Effect.Effect<R3, never, void>
+  getKey: (a: A) => C
 }) {
   return Effect.gen(function* ($) {
+    const key = getKey(value)
     const subject = RefSubject.unsafeMake<never, A>(Effect.succeed(value))
-    const fx = f(subject, value)
+    const fx = f(subject)
     const fiber = yield* $(
       fork(
         fx.run(
           Sink(
             (b: B) =>
               Effect.suspend(() => {
-                MutableHashMap.set(state.values, value, b)
+                MutableHashMap.set(state.values, key, b)
                 return emit
               }),
             error,
@@ -174,14 +184,14 @@ function addValue<A, B, R2, E2, R3>({
       ),
     )
 
-    MutableHashMap.set(state.subjects, value, subject)
-    MutableHashMap.set(state.fibers, value, fiber)
+    MutableHashMap.set(state.subjects, key, subject)
+    MutableHashMap.set(state.fibers, key, fiber)
   })
 }
 
-function updateValue<A, B>(state: KeyedState<A, B>, value: A) {
+function updateValue<A, B, C>(state: KeyedState<A, B, C>, value: A, getKey: (a: A) => C) {
   return Effect.gen(function* ($) {
-    const subject = MutableHashMap.get(state.subjects, value)
+    const subject = MutableHashMap.get(state.subjects, getKey(value))
 
     // Send the current value
     if (Option.isSome(subject)) {
@@ -190,10 +200,10 @@ function updateValue<A, B>(state: KeyedState<A, B>, value: A) {
   })
 }
 
-function emitWhenReady<A, B>(state: KeyedState<A, B>) {
+function emitWhenReady<A, B, C>(state: KeyedState<A, B, C>, getKey: (a: A) => C) {
   return Effect.suspend(() => {
     const values = ReadonlyArray.filterMap(state.previous, (value) =>
-      MutableHashMap.get(state.values, value),
+      MutableHashMap.get(state.values, getKey(value)),
     )
 
     // When all of the values have resolved at least once, emit the output
@@ -205,7 +215,7 @@ function emitWhenReady<A, B>(state: KeyedState<A, B>) {
   })
 }
 
-function endAll<A, B>(state: KeyedState<A, B>) {
+function endAll<A, B, C>(state: KeyedState<A, B, C>) {
   return Effect.gen(function* ($) {
     yield* $(Effect.forEachParDiscard(state.subjects, ([, subject]) => subject.end()))
   })