transducers.ml

open Elements

type ('a, 'state) step = 'state -> 'a -> [ `Continue of 'state | `Stop of 'state ]

type ('a, 'b, 'state) transducer =
  ('b, 'state) step -> ('a, 'state) step

let map f =
  fun step -> fun r x ->
        step r (f x)

let filter p =
  fun step -> fun r x ->
        if p x then step r x else `Continue r

let take n =
  fun step ->
    let count = ref 0 in
    fun r x ->
      if !count < n
      then (incr count; step r x)
      else `Stop r

let compose t1 t2 =
  fun xf -> (t1 (t2 xf))

let (>>) = compose

let transduce_list t xs =
  List.rev (List.fold_until
              ~f:(t (fun xs x -> `Continue (x :: xs)))
              ~init:[] xs)

let test () =
  assert (transduce_list (map ((+) 5) >> take 10) (List.range 0 5)
          = [5; 7; 9])