tiny hashconsing for ocaml

dune

(executable
  (name hashcons_test)
  (libraries iter containers ppx_deriving.std)
  (flags :standard -warn-error -a+8)
  (preprocess (pps ppx_deriving.std)))

hashcons_test.ml

(* vendored *)
module Hashcons = struct

  module type ARG = sig
    type t
    val equal : t -> t -> bool
    val hash : t -> int
    val set_id : t -> int -> unit
  end

  module Make(A : ARG): sig
    val hashcons : A.t -> A.t
    val to_seq : A.t Iter.t
    val shrink : unit -> unit
  end = struct
    module W = Weak.Make(A)

    let tbl_ = ref (W.create 1024)
    let n_ = ref 0

    (* hashcons terms *)
    let hashcons t =
      let t' = W.merge !tbl_ t in
      if t == t' then (
        incr n_;
        A.set_id t' !n_;
      );
      t'

    let to_seq yield =
      W.iter yield !tbl_

    let shrink () =
      let l = ref [] in
      to_seq (fun x -> l := x :: !l);
      tbl_ := W.create 1024;
      List.iter (fun t -> ignore @@ W.merge !tbl_ t) !l;
      ()
  end
end

type 'a view =
  | App of 'a * 'a
  | Lam of string * 'a
  | Var of string
[@@deriving show, eq]
;;

let hash_view h = function
  | App (a,b) -> CCHash.combine3 10 (h a) (h b)
  | Lam (x,bod) -> CCHash.combine3 20 (CCHash.string x) (h bod)
  | Var x -> CCHash.combine2 30 (CCHash.string x)

type t = {
  mutable id: int;
  view: t view;
}
let equal (x:t) y = x==y
let hash (x:t) = CCHash.int x.id

module H = Hashcons.Make(struct
    type nonrec t = t
    let hash t = hash_view hash t.view
    let equal a b = equal_view equal a.view b.view
    let set_id x id = assert (x.id < 0); x.id <- id
  end)

let var x : t = H.hashcons {id= -1; view=Var x}
let app x y : t = H.hashcons {id= -1; view=App (x,y)}
let lam x body : t = H.hashcons {id= -1; view=Lam (x,body)}

let rec of_n n : t =
  if n=0 then lam "z" @@ lam "s" @@ var "z"
  else (
    lam "z" @@ lam "s" @@ app (var "s") (of_n (n-1))
  )

let () =
  Printf.printf "### START\n%!";
  Gc.print_stat stdout;
  Printf.printf "### WITH TERM\n%!";
  begin
    let _t = of_n 100_000 in
    let n_t = H.to_seq |> Iter.length in
    Printf.printf "%d terms\n%!" n_t;
    Gc.print_stat stdout;
  end;
  Gc.major();
  Gc.compact();
  Printf.printf "### AFTER GC\n%!";
  let n_t = H.to_seq |> Iter.length in
  Printf.printf "%d terms\n%!" n_t;
  Gc.print_stat stdout;
  Printf.printf "### AFTER SHRINK\n%!";
  H.shrink();
  Gc.major();
  Gc.compact();
  let n_t = H.to_seq |> Iter.length in
  Printf.printf "%d terms\n%!" n_t;
  Gc.print_stat stdout;
  ()

output.txt

$ dune exec ./hashcons_test.exe
### START            
minor_collections: 0
major_collections: 0
compactions:       0

minor_words:    35273
promoted_words:     0
major_words:     3221

top_heap_words: 491520
heap_words:     491520
live_words:       3221
free_words:     488299
largest_free:   488299
fragments:           0

live_blocks: 6
free_blocks: 1
heap_chunks: 1
### WITH TERM
299261 terms
minor_collections: 121
major_collections: 10
compactions:       0

minor_words:    30543345
promoted_words:  4015466
major_words:     7048832

top_heap_words: 4612096
heap_words:     4612096
live_words:     3178211
free_words:     1429734
largest_free:    180202
fragments:         4151

live_blocks: 722958
free_blocks: 64815
heap_chunks: 17
### AFTER GC
0 terms
minor_collections: 123
major_collections: 13
compactions:       1

minor_words:    30543987
promoted_words:  4015466
major_words:     7048832

top_heap_words: 4612096
heap_words:     1975808
live_words:     1050697
free_words:      925111
largest_free:    455168
fragments:            0

live_blocks: 82903
free_blocks: 3
heap_chunks: 4
### AFTER SHRINK
0 terms
minor_collections: 125
major_collections: 16
compactions:       3

minor_words:    30544659
promoted_words:  4015472
major_words:     7050888

top_heap_words: 4612096
heap_words:       90624
live_words:       30347
free_words:       60277
largest_free:     60277
fragments:            0

live_blocks: 3807
free_blocks: 1
heap_chunks: 1