@electric-sql/d2ts/src/operators/filter.ts

D2TS is a TypeScript implementation of Differential Dataflow.

import { IStreamBuilder, PipedOperator } from '../types.js'
import { DifferenceStreamReader, DifferenceStreamWriter } from '../graph.js'
import { StreamBuilder } from '../d2.js'
import { LinearUnaryOperator } from './base.js'
import { MultiSet } from '../multiset.js'
import { Antichain } from '../order.js'

/**
 * Operator that filters elements from the input stream
 */
export class FilterOperator<T> extends LinearUnaryOperator<T, T> {
  #f: (data: T) => boolean

  constructor(
    id: number,
    inputA: DifferenceStreamReader<T>,
    output: DifferenceStreamWriter<T>,
    f: (data: T) => boolean,
    initialFrontier: Antichain,
  ) {
    super(id, inputA, output, initialFrontier)
    this.#f = f
  }

  inner(collection: MultiSet<T>): MultiSet<T> {
    return collection.filter(this.#f)
  }
}

/**
 * Filters elements from the input stream
 * @param f - The predicate to filter elements
 */
export function filter<T>(f: (data: T) => boolean): PipedOperator<T, T> {
  return (stream: IStreamBuilder<T>): IStreamBuilder<T> => {
    const output = new StreamBuilder<T>(
      stream.graph,
      new DifferenceStreamWriter<T>(),
    )
    const operator = new FilterOperator<T>(
      stream.graph.getNextOperatorId(),
      stream.connectReader(),
      output.writer,
      f,
      stream.graph.frontier(),
    )
    stream.graph.addOperator(operator)
    stream.graph.addStream(output.connectReader())
    return output
  }
}