Source

ocaml-logicm / logicM_stream.ml

let impl_name = "stream"

    open ExtStream

    type 'a m = 'a Stream.t
    type 'a mplus_sig = ?chunk:int -> 'a m -> 'a m -> 'a m

    let mzero = Stream.sempty

    let return = Stream.ising

    let get = Stream.get_opt

    let rec take n m =
      if n = 0
      then []
      else
        match get m with
        | None -> ( (* Printf.printf "take finished\n%!"; *) [] )
        | Some x -> x :: take (n-1) m

    let take_all m =
      let rec inner acc =
        match get m with
        | None -> acc
        | Some x -> inner (x :: acc)
      in
        inner []

    let runL optcount m =
      match optcount with
      | None -> take_all m
      | Some n -> take n m

    let from_stream s = s

    let rec interleave ?(chunk=100) a b =
      let cont = ref `Empty in
      let do_cont () =
        match !cont with
        | `Empty -> assert false
        | `Interleave (a, b) -> interleave ~chunk a b
        | `One a -> a
      in
      let rec inner ~left ~acc a b =
        if left = 0
        then
          (cont := `Interleave (a, b); acc)
        else
          match get a with
          | None -> (cont := `One b; acc)
          | Some x -> inner ~left:(left-1) ~acc:(Stream.icons x acc) b a
      in
        inner ~left:chunk ~acc:(Stream.slazy do_cont) a b

    let ( <+> ) ?(chunk=100) a b = interleave ~chunk a b

    let limit n m =
      let rec inner acc n m =
        if n <= 0
        then acc
        else
          match get m with
          | None -> mzero
          | Some x -> inner (Stream.icons x acc) (n-1) m
      in
        inner mzero n m

    let once m =
        match get m with
        | None -> mzero
        | Some x -> return x

    let join_list_step ~chunk lst : 'a m * 'a m list * (unit -> 'a m) ref =
      let cont = ref (fun () -> assert false) in
      let do_cont () = !cont () in
      let rec join_m ~left ~sacc m =
        if left = 0
        then (sacc, Some m)
        else
          match get m with
          | None -> (sacc, None)
          | Some x ->
              join_m m ~left:(left-1)
                ~sacc:(Stream.icons x sacc)
      in
      let join_list accs lst =
        List.fold_left
          (fun (sacc, macc) m ->
             let (sacc, mres) = join_m ~left:chunk ~sacc m in
             (sacc, match mres with None -> macc | Some m -> m :: macc)
          )
          accs
          lst
      in
        let (sacc, macc) = join_list ((Stream.slazy do_cont), []) lst in
        (sacc, macc, cont)

    let map ?(chunk=100) f m =
      let cont = ref (fun () -> assert false) in
      let do_cont () = !cont () in
      let rec inner acc ~left =
        if left = 0
        then acc
        else
          match get m with
          | None -> acc
          | Some x -> inner (Stream.icons (f x) acc) ~left:(left-1)
      in
        inner (Stream.slazy do_cont) ~left:chunk

    let rec bind_diag_inner ~chunk ~blist f m =
      let (stream, blist, cont) = join_list_step ~chunk blist in
      let mtook = take chunk m in
      let new_b's = List.rev_map f mtook in
      let blist = List.rev_append new_b's blist in
      ( cont :=
          (match (blist, mtook) with
           | [], [] -> fun () -> mzero
           | _ -> fun () -> bind_diag_inner ~chunk ~blist f m
          );
        stream
      )

    let bind_diag ~chunk f m =
      bind_diag_inner ~chunk ~blist:[] f m

    let map_all_rev f m =
      let rec inner acc =
        match get m with
        | None -> acc
        | Some x -> inner ((f x) :: acc)
      in
        inner []

    let bind_wide ~chunk f m =
      let all_b = map_all_rev f m in
      bind_diag_inner ~chunk ~blist:all_b f mzero

    let bind_deep ~chunk f m =
      let rec inner ~chunk ~blist m =
        let (stream, blist, cont) =
          join_list_step ~chunk blist in
        ( cont :=
            (if blist = []
             then
               match get m with
                 None -> fun () -> mzero
               | Some a -> fun () -> inner ~chunk ~blist:[f a] m
             else
               fun () -> inner ~chunk ~blist m
            );
          stream
        )
      in
        inner ~chunk ~blist:[] m


    let default_diag_n = 20
    and default_wide_n = 10
    and default_deep_n = 100

    type alg =
      [ `Diag
      | `Diagn of int
      | `Wide
      | `Widen of int
      | `Deep
      | `Deepn of int
      ]

    let ( >>= ) m f = bind_diag ~chunk:default_diag_n f m

    let bind ?(alg=`Diag) f m =
      match alg with
      | `Diag -> bind_diag ~chunk:default_diag_n f m
      | `Diagn n -> bind_diag ~chunk:n f m

      | `Wide -> bind_wide ~chunk:default_wide_n f m
      | `Widen n -> bind_wide ~chunk:n f m

      | `Deep -> bind_deep ~chunk:default_deep_n f m
      | `Deepn n -> bind_deep ~chunk:n f m

    external identity : 'a -> 'a = "%identity"

    let join ?(alg=`Diag) m =
      match alg with
      | `Diag -> bind_diag ~chunk:default_diag_n identity m
      | `Diagn n -> bind_diag ~chunk:n identity m

      | `Wide -> bind_wide ~chunk:default_wide_n identity m
      | `Widen n -> bind_wide ~chunk:n identity m

      | `Deep -> bind_deep ~chunk:default_deep_n identity m
      | `Deepn n -> bind_deep ~chunk:n identity m


    let filter f m =  (* TODO: more optimal impl. *)
      m >>= (fun x -> if f x then return x else mzero)

    let ifteu ?(alg=`Diag) m th el =
      match get m with
      | None -> el ()
      | Some h -> bind ~alg th (Stream.icons h m)

    let ifte ?(alg=`Diag) m th el =
      ifteu ~alg m th (fun () -> el)

    let ifm m ~th ~el =
      match get m with
      | None -> el ()
      | Some h -> th (Stream.icons h m)

    let guard cond =
      if cond then return () else mzero

    let (stream_assign_data : 'a Stream.t -> 'a Stream.t -> unit) =
    fun sfrom sto ->
      Obj.set_field (Obj.repr sto) 1 (Obj.field (Obj.repr sfrom) 1)

    let (obj_dup : 'a -> 'a) = fun x -> Obj.obj (Obj.dup (Obj.repr x))

    let msplit m =
      match Stream.peek m with
      | None -> None
      | Some hd ->
          ( Stream.junk m;
            let tl = obj_dup m in
            stream_assign_data (Stream.icons hd tl) m;
            Some (hd, tl)
          )

    let destruct_bind m th el =
      match msplit m with
      | None -> el
      | Some (h, t) -> th h t

    let destruct_bindu m th el =
      match msplit m with
      | None -> el ()
      | Some (h, t) -> th h t

    let iter = Stream.iter