Treap map

avltree2.ml

type ('a, 'b) tree =
  | Empty
  | Node of
      { left : ('a, 'b) tree
      ; right : ('a, 'b) tree
      ; key : 'a
      ; value : 'b
      ; height : int
      }

let[@inline] height = function
  | Empty -> 0
  | Node { height } -> height
;;

let mknode l k v r =
  let hl = height l in
  let hr = height r in
  Node
    { left = l
    ; key = k
    ; value = v
    ; right = r
    ; height = (if hl >= hr then hl + 1 else hr + 1)
    }
;;

let balance l k v r =
  let hl = height l in
  let hr = height r in
  if hl > hr + 1
  then (
    let[@warning "-8"] (Node ln) = l in
    if height ln.left >= height ln.right
    then mknode ln.left ln.key ln.value (mknode ln.right k v r)
    else (
      let[@warning "-8"] (Node lrn) = ln.right in
      mknode
        (mknode ln.left ln.key ln.value lrn.left)
        lrn.key
        lrn.value
        (mknode lrn.right k v r)))
  else if hr > hl + 1
  then (
    let[@warning "-8"] (Node rn) = r in
    if height rn.right >= height rn.left
    then mknode (mknode l k v rn.left) rn.key rn.value rn.right
    else (
      let[@warning "-8"] (Node rln) = rn.left in
      mknode
        (mknode l k v rln.left)
        rln.key
        rln.value
        (mknode rln.right rn.key rn.value rn.right)))
  else mknode l k v r
;;

let rec add ~cmp t k v =
  match t with
  | Empty -> Node { left = Empty; right = Empty; key = k; value = v; height = 1 }
  | Node n ->
    let c = cmp k n.key in
    if c = 0
    then Node { n with value = v }
    else if c < 0
    then balance (add ~cmp n.left k v) n.key n.value n.right
    else balance n.left n.key n.value (add ~cmp n.right k v)
;;

let rec find ~cmp t k =
  match t with
  | Empty -> None
  | Node n ->
    let cr = cmp k n.key in
    if cr < 0
    then find ~cmp n.left k
    else if cr > 0
    then find ~cmp n.right k
    else Some n.value
;;

bench.ml

open Core_bench
module IntMap = Map.Make (Int)

module Int_table = Stdlib.Hashtbl.Make (struct
    include Stdlib.Int

    let hash = Hashtbl.hash
  end)

let rec mem_all n ~f =
  match n with
  | -1 -> ()
  | n ->
    ignore (f n);
    mem_all (pred n) ~f
;;

let () =
  let open Core in
  let seeds = Array.create ~len:100_000 0 in
  for i = 0 to 100_000 - 1 do
    seeds.(i) <- Random.int Int.max_value
  done;
  let args = [ 100_00 ] in
  let tests =
    [ (fun args ->
        [ Bench.Test.create_indexed ~name:"Avltree.mem" ~args (fun arg ->
            let map = ref Avltree2.Empty in
            for i = 0 to arg - 1 do
              map := Avltree2.add ~cmp:Int.compare !map seeds.(i) i
            done;
            Printf.printf "height: %d" (Avltree2.height !map);
            let f i = Avltree2.find ~cmp:Int.compare !map i in
            Staged.stage @@ fun () -> mem_all arg ~f)
        ; Bench.Test.create_indexed ~name:"Treap.mem" ~args (fun arg ->
            let map = ref Treap.Empty in
            for i = 0 to arg - 1 do
              map := Treap.add ~cmp:Int.compare !map seeds.(i) i
            done;
            Printf.printf "height: %d" (Treap.height !map);
            let f i = Treap.find ~cmp:Int.compare !map i in
            Staged.stage @@ fun () -> mem_all arg ~f)
        ])
    ]
  in
  List.iter
    ~f:(fun tests ->
      List.iter
        ~f:(fun arg ->
          let args = [ arg ] in
          Command_unix.run @@ Bench.make_command (tests args))
        args)
    tests
;;

treap.ml

type ('a, 'b) treap =
  | Empty
  | Node of
      { left : ('a, 'b) treap
      ; right : ('a, 'b) treap
      ; priority : int
      ; key : 'a
      ; value : 'b
      }

let rec add ~cmp t k v =
  match t with
  | Empty ->
    Node { left = Empty; right = Empty; priority = Random.bits (); key = k; value = v }
  | Node n ->
    let cr = cmp k n.key in
    if cr < 0
    then (
      let nl = add ~cmp n.left k v in
      let[@warning "-8"] (Node nln) = nl in
      if nln.priority < n.priority
      then Node { nln with right = Node { n with left = nln.right } }
      else Node { n with left = nl })
    else if cr > 0
    then (
      let nr = add ~cmp n.right k v in
      let[@warning "-8"] (Node nrn) = nr in
      if nrn.priority < n.priority
      then Node { nrn with left = Node { n with right = nrn.left } }
      else Node { n with right = nr })
    else Node { n with value = v }
;;

let rec find ~cmp t k =
  match t with
  | Empty -> None
  | Node n ->
    let cr = cmp k n.key in
    if cr < 0
    then find ~cmp n.left k
    else if cr > 0
    then find ~cmp n.right k
    else Some n.value
;;

let rec height t =
  match t with
  | Empty -> 0
  | Node { left; right } -> Int.max (height left) (height right) + 1
;;