laziness.ml

(* An implementation of (part of) the lazy data structure of:

      First-Order Laziness
      Anton Lorenzen, Daan Leijen, Wouter Swierstra, Sam Lindley
      ICFP 2025

   Requires OCaml 5.2+ / OxCaml *)

module type LazyS = sig
  type _ repr
  type t
  val make : 'a repr -> t
  val make_forced : [`Forced] repr -> t
  val force : t -> [`Forced] repr
end

module MakeLazy (X : sig
  type _ t
  val force : 'a t -> [`Forced] t
end) : LazyS with type 'a repr := 'a X.t = struct
  type t = Obj.t

  let make x = Obj.with_tag Obj.lazy_tag (Obj.repr (ref x))

  let make_forced x =
    let x = Obj.repr x in
    let t = Obj.tag x in
    if t = Obj.lazy_tag || t = Obj.forward_tag || t = Obj.forcing_tag || t = Obj.double_tag
    then Obj.with_tag Obj.forward_tag (Obj.repr (ref x))
    else x

  external update_to_forcing : Obj.t -> int =
    "caml_lazy_update_to_forcing" [@@noalloc]

  external update_to_forward : Obj.t -> unit =
    "caml_lazy_update_to_forward" [@@noalloc]
  let force x =
    let x = Obj.repr (Sys.opaque_identity x) in
    match Obj.tag x with
    | tag when tag = Obj.forward_tag -> Obj.obj (Obj.field x 0)
    | tag when tag = Obj.forcing_tag -> raise Lazy.Undefined
    | tag when tag <> Obj.lazy_tag -> Obj.obj x
    | _ ->
       if update_to_forcing x <> 0 then raise Lazy.Undefined;
       let orig = Obj.field x 0 in
       let res = X.force (Obj.obj orig) in
       Obj.set_field x 0 (Obj.repr res);
       update_to_forward x;
       res
end

module rec LazyListTypes : sig
  type _ t =
    | Nil : [`Forced] t
    | Cons : int * LazyList.t  -> [`Forced] t
    | Append : LazyList.t * LazyList.t -> [`Lazy] t
end = LazyListTypes
and LazyList : LazyS with type 'a repr := 'a LazyListTypes.t =
MakeLazy (struct
  include LazyListTypes
  let force (type a) (x : a t) =
    match x with
    | Nil | Cons _ as x -> x
    | Append (xs, ys) ->
       match LazyList.force xs with
       | Nil -> LazyList.force ys
       | Cons (x, xs) -> Cons (x, LazyList.make (Append (xs, ys)))
end)


let () =
  let nil = LazyList.make_forced Nil in
  let cons x xs = LazyList.make_forced (Cons (x, xs)) in
  let singleton x = cons x nil in
  let append xs ys = LazyList.make (Append (xs, ys)) in
  let rec sum xs =
    match LazyList.force xs with
    | Nil -> 0
    | Cons (x, xs) -> x + sum xs
  in
  let len = 50 in
  let singletons = List.init len singleton in
  let appended = List.fold_left append nil singletons in
  let print_size ~phase =
    let size = Obj.reachable_words (Obj.repr appended) in
    Printf.printf "%12s: %4d words (%.1f/cons)\n" phase size (float_of_int size /. float_of_int len)
  in
  print_size ~phase:"lazy";
  let n = sum appended in
  assert (n = len * (len-1) / 2);
  print_size ~phase:"forced";
  Gc.minor ();
  print_size ~phase:"forced+GC"

output.txt

        lazy:  400 words (8.0/cons)
      forced:  250 words (5.0/cons)
   forced+GC:  150 words (3.0/cons)