Source

hasexp / zlist.ml

TYPE_CONV_PATH "Zlist"

module Monad = struct
  module Infix(M : sig 
    type 'a t
    val return : 'a -> 'a t
    val bind : 'a t -> ('a -> 'b t) -> 'b t
  end) = struct
    let (>>=) = M.bind
    let (>>|) t f = M.bind t (fun v -> M.return (f v))
  end
end

module Lazy = struct
  module Lazy = struct
    include Lazy

    let return : 'a -> 'a t = lazy_from_val
    let bind z f = lazy ( force( f (force z)) )
    let (!!) = force
  end

  include Lazy
  module Infix = Monad.Infix(Lazy)
end

module Z = Lazy
open Z.Infix
let (!!) = Z.(!!)

type 'a zlist =
    | Cons of 'a * 'a t
    | Nil

and 'a t = 'a zlist Z.t

let nil : 'a t = Z.return Nil
let cons v vs = Z.return (Cons (v, vs))
let (^^) = cons

let hd_tl t = 
  match !!t with
  | Nil -> None
  | Cons (v, vs) -> Some (v,vs)

let hd t =
  match hd_tl t with
  | None -> None
  | Some (v, _) -> Some v

let tl t =
  match hd_tl t with
  | None -> None
  | Some (_, tl) -> Some tl

let npeek t n = 
  let rec npeek acc n t =
    if n = 0 then acc
    else
      match Lazy.force t with
      | Nil -> acc
      | Cons (v, vs) -> npeek (v::acc) (n-1) vs
  in
  List.rev (npeek [] n t)

let rec force t = 
  match Lazy.force t with
  | Nil -> ()
  | Cons (_, vs) -> force vs

let rec iter ~f t =
  match Lazy.force t with
  | Nil -> ()
  | Cons (v, vs) -> f v; iter ~f vs

let to_list t =
  let rec force st t = 
    match Lazy.force t with
    | Nil -> List.rev st
    | Cons (v,vs) -> force (v::st) vs
  in
  force [] t

let create (f : unit -> 'a option) : 'a t =
  let rec zlist () =
    lazy (
      match f () with
      | None ->  Nil
      | Some v -> Cons (v, zlist ())
    )
  in
  zlist ()


let of_list l =
  let rec of_list = function
    | [] -> Z.return Nil
    | x::xs -> Z.return (Cons (x, of_list xs))
  in
  ref (of_list l)

let is_nil t = hd t = None

let rec map ~f t =
  lazy (
    match Lazy.force t with
    | Nil -> Nil
    | Cons (v, vs) -> Cons (f v, map ~f vs)
  )

let filter_map ~f t =
  let rec filter ~f t =
    lazy (
      match Lazy.force t with
      | Nil -> Nil
      | Cons (v, vs) -> 
	  match f v with
	  | Some v -> Cons (v, filter ~f vs)
	  | None -> Lazy.force (filter ~f vs)
    )
  in
  filter ~f t

let filter ~f t =
  filter_map t ~f:(fun v ->
    if f v then Some v else None)

let fold_right ~f t ~init = 
  let rec fold_right_ ~f t ~init =
    match Lazy.force t with
    | Nil -> init
    | Cons (v, vs) -> f v (fold_right_ ~f vs ~init)
  in
  lazy (fold_right_ ~f t ~init)

let rec of_list = function
  | [] -> Z.return Nil
  | x::xs -> Z.return (Cons (x, of_list xs))

let singleton v = of_list [v]

let is_singleton t = npeek t 2
  
let append t1 t2 =
  let rec zlist t =
    lazy (
      match Lazy.force t with
      | Cons (v,vs) -> Cons(v, zlist vs)
      | Nil -> Lazy.force t2
    )
  in
  zlist t1

let return = singleton

let zz z = lazy (Lazy.force (Lazy.force z))

let bind = fun t f ->
  zz (fold_right ~f:append (map ~f t) ~init:nil)

module Infix = struct
  let (^^) = (^^)
  let (>>=) = bind
  let (>>|) t f = bind t (fun v -> return (f v))
end

let test () =
  assert (to_list (append (of_list [1;2;3]) (of_list [4;5;6])) 
	   = [1;2;3;4;5;6]);
  assert (to_list 
	     (bind (of_list [1;2;3]) (fun n -> of_list [n*3; n*3+1; n*3+2]))
	   = [3;4;5;6;7;8;9;10;11])
    
type 'a forced = 'a list with sexp

let sexp_of_t sexp_of_a t =
  sexp_of_forced sexp_of_a (to_list t)
;;

let t_of_sexp a_of_sexp l =
  of_list (forced_of_sexp a_of_sexp l)
;;