Source file dmap.ml

(**************************************************************************)
(*                                                                        *)
(*   Authors:                                                             *)
(*             Xavier Leroy, projet Cristal, INRIA Rocquencourt           *)
(*   Benoît Montagu, projet Celtique, INRIA Rennes Bretagne Atlantique    *)
(*                                                                        *)
(*   Copyright 1996 Institut National de Recherche en Informatique et     *)
(*     en Automatique.                                                    *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

type (_, _) cmp = Lt : (_, _) cmp | Eq : ('a, 'a) cmp | Gt : (_, _) cmp

module type DORDERED = sig
  type _ t

  val compare : 'a t -> 'b t -> ('a, 'b) cmp
end

module type S_WITH_VALUE = sig
  type 'a key
  type 'a value
  type t

  val empty : t
  val is_empty : t -> bool
  val mem : 'a key -> t -> bool
  val add : 'a key -> 'a value -> t -> t
  val update : 'a key -> ('a value option -> 'a value option) -> t -> t
  val singleton : 'a key -> 'a value -> t
  val remove : 'a key -> t -> t

  type poly_merge = {
    merge_fun :
      'a. 'a key -> 'a value option -> 'a value option -> 'a value option;
  }
  [@@unboxed]

  val merge : poly_merge -> t -> t -> t

  type poly_union = {
    union_fun : 'a. 'a key -> 'a value -> 'a value -> 'a value option;
  }
  [@@unboxed]

  val union : poly_union -> t -> t -> t

  type 'b poly1 = { fun1 : 'a. 'a key -> 'a value -> 'b } [@@unboxed]

  type 'b poly2 = { fun2 : 'a. 'a key -> 'a value -> 'a value -> 'b }
  [@@unboxed]

  val compare : int poly2 -> t -> t -> int
  val equal : bool poly2 -> t -> t -> bool
  val iter : unit poly1 -> t -> unit
  val fold : ('b -> 'b) poly1 -> t -> 'b -> 'b
  val for_all : bool poly1 -> t -> bool
  val exists : bool poly1 -> t -> bool
  val filter : bool poly1 -> t -> t

  type poly_filter_map = {
    filter_map_fun : 'a. 'a key -> 'a value -> 'a value option;
  }
  [@@unboxed]

  val filter_map : poly_filter_map -> t -> t
  val partition : bool poly1 -> t -> t * t
  val cardinal : t -> int

  type binding = Binding : 'a key * 'a value -> binding

  val bindings : t -> binding list
  val min_binding : t -> binding
  val min_binding_opt : t -> binding option
  val max_binding : t -> binding
  val max_binding_opt : t -> binding option
  val choose : t -> binding
  val choose_opt : t -> binding option
  val split : 'a key -> t -> t * 'a value option * t
  val find : 'a key -> t -> 'a value
  val find_opt : 'a key -> t -> 'a value option

  type 'b poly0 = { fun0 : 'a. 'a key -> 'b } [@@unboxed]

  val find_first : bool poly0 -> t -> binding
  val find_first_opt : bool poly0 -> t -> binding option
  val find_last : bool poly0 -> t -> binding
  val find_last_opt : bool poly0 -> t -> binding option

  type poly_mapi = { mapi_fun : 'a. 'a key -> 'a value -> 'a value } [@@unboxed]

  val mapi : poly_mapi -> t -> t
  val to_seq : t -> binding Seq.t
  val to_rev_seq : t -> binding Seq.t
  val to_seq_from : 'a key -> t -> binding Seq.t
  val add_seq : binding Seq.t -> t -> t
  val of_seq : binding Seq.t -> t
end

module type S = S_WITH_VALUE with type 'a value := 'a

module type TYPE1 = sig
  type 'a t
end

module MakeWithValue (Ord : DORDERED) (Val : TYPE1) = struct
  type 'a key = 'a Ord.t
  type 'a value = 'a Val.t

  type t =
    | Empty
    | Node : { l : t; v : 'a key; d : 'a value; r : t; h : int } -> t

  let height = function Empty -> 0 | Node { h; _ } -> h

  let create l x d r =
    let hl = height l and hr = height r in
    Node { l; v = x; d; r; h = (if hl >= hr then hl + 1 else hr + 1) }

  let singleton x d = Node { l = Empty; v = x; d; r = Empty; h = 1 }

  let bal l x d r =
    let hl = match l with Empty -> 0 | Node { h; _ } -> h in
    let hr = match r with Empty -> 0 | Node { h; _ } -> h in
    if hl > hr + 2 then
      match l with
      | Empty -> invalid_arg "Dmap.bal"
      | Node { l = ll; v = lv; d = ld; r = lr; h = _ } -> (
          if height ll >= height lr then create ll lv ld (create lr x d r)
          else
            match lr with
            | Empty -> invalid_arg "Dmap.bal"
            | Node { l = lrl; v = lrv; d = lrd; r = lrr; h = _ } ->
                create (create ll lv ld lrl) lrv lrd (create lrr x d r))
    else if hr > hl + 2 then
      match r with
      | Empty -> invalid_arg "Dmap.bal"
      | Node { l = rl; v = rv; d = rd; r = rr; h = _ } -> (
          if height rr >= height rl then create (create l x d rl) rv rd rr
          else
            match rl with
            | Empty -> invalid_arg "Dmap.bal"
            | Node { l = rll; v = rlv; d = rld; r = rlr; h = _ } ->
                create (create l x d rll) rlv rld (create rlr rv rd rr))
    else Node { l; v = x; d; r; h = (if hl >= hr then hl + 1 else hr + 1) }

  let empty = Empty
  let is_empty = function Empty -> true | _ -> false

  let rec add : type a. a key -> a value -> t -> t =
   fun x data -> function
    | Empty -> Node { l = Empty; v = x; d = data; r = Empty; h = 1 }
    | Node { l; v; d; r; h } as m -> (
        match Ord.compare x v with
        | Eq -> if d == data then m else Node { l; v = x; d = data; r; h }
        | Lt ->
            let ll = add x data l in
            if l == ll then m else bal ll v d r
        | Gt ->
            let rr = add x data r in
            if r == rr then m else bal l v d rr)

  type binding = Binding : 'a key * 'a value -> binding

  let rec find : type a. a key -> t -> a value =
   fun x -> function
    | Empty -> raise Not_found
    | Node { l; v; d; r; h = _ } -> (
        match Ord.compare x v with Eq -> d | Lt -> find x l | Gt -> find x r)

  type 'b poly0 = { fun0 : 'a. 'a key -> 'b } [@@unboxed]
  type 'b poly1 = { fun1 : 'a. 'a key -> 'a value -> 'b } [@@unboxed]

  type 'b poly2 = { fun2 : 'a. 'a key -> 'a value -> 'a value -> 'b }
  [@@unboxed]

  let rec find_first_aux :
      type a. a key -> a value -> bool poly0 -> t -> binding =
   fun v0 d0 f -> function
    | Empty -> Binding (v0, d0)
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_first_aux v d f l else find_first_aux v0 d0 f r

  let rec find_first f = function
    | Empty -> raise Not_found
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_first_aux v d f l else find_first f r

  let rec find_first_opt_aux :
      type a. a key -> a value -> bool poly0 -> t -> binding option =
   fun v0 d0 f -> function
    | Empty -> Some (Binding (v0, d0))
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_first_opt_aux v d f l
        else find_first_opt_aux v0 d0 f r

  let rec find_first_opt f = function
    | Empty -> None
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_first_opt_aux v d f l else find_first_opt f r

  let rec find_last_aux : type a. a key -> a value -> bool poly0 -> t -> binding
      =
   fun v0 d0 f -> function
    | Empty -> Binding (v0, d0)
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_last_aux v d f r else find_last_aux v0 d0 f l

  let rec find_last f = function
    | Empty -> raise Not_found
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_last_aux v d f r else find_last f l

  let rec find_last_opt_aux :
      type a. a key -> a value -> bool poly0 -> t -> binding option =
   fun v0 d0 f -> function
    | Empty -> Some (Binding (v0, d0))
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_last_opt_aux v d f r
        else find_last_opt_aux v0 d0 f l

  let rec find_last_opt f = function
    | Empty -> None
    | Node { l; v; d; r; h = _ } ->
        if f.fun0 v then find_last_opt_aux v d f r else find_last_opt f l

  let rec find_opt : type a. a key -> t -> a value option =
   fun x -> function
    | Empty -> None
    | Node { l; v; d; r; h = _ } -> (
        match Ord.compare x v with
        | Eq -> Some d
        | Lt -> find_opt x l
        | Gt -> find_opt x r)

  let rec mem : type a. a key -> t -> bool =
   fun x -> function
    | Empty -> false
    | Node { l; v; r; d = _; h = _ } -> (
        match Ord.compare x v with Eq -> true | Lt -> mem x l | Gt -> mem x r)

  let rec min_binding = function
    | Empty -> raise Not_found
    | Node { l = Empty; v; d; r = _; h = _ } -> Binding (v, d)
    | Node { l; v = _; d = _; r = _; h = _ } -> min_binding l

  let rec min_binding_opt = function
    | Empty -> None
    | Node { l = Empty; v; d; r = _; h = _ } -> Some (Binding (v, d))
    | Node { l; v = _; d = _; r = _; h = _ } -> min_binding_opt l

  let rec max_binding = function
    | Empty -> raise Not_found
    | Node { v; d; r = Empty; l = _; h = _ } -> Binding (v, d)
    | Node { r; v = _; d = _; l = _; h = _ } -> max_binding r

  let rec max_binding_opt = function
    | Empty -> None
    | Node { v; d; r = Empty; l = _; h = _ } -> Some (Binding (v, d))
    | Node { r; v = _; d = _; l = _; h = _ } -> max_binding_opt r

  let rec remove_min_binding = function
    | Empty -> invalid_arg "Map.remove_min_elt"
    | Node { l = Empty; r; v = _; d = _; h = _ } -> r
    | Node { l; v; d; r; h = _ } -> bal (remove_min_binding l) v d r

  let merge t1 t2 =
    match (t1, t2) with
    | Empty, t -> t
    | t, Empty -> t
    | _, _ ->
        let (Binding (x, d)) = min_binding t2 in
        bal t1 x d (remove_min_binding t2)

  let rec remove : type a. a key -> t -> t =
   fun x -> function
    | Empty -> Empty
    | Node { l; v; d; r; h = _ } as m -> (
        match Ord.compare x v with
        | Eq -> merge l r
        | Lt ->
            let ll = remove x l in
            if l == ll then m else bal ll v d r
        | Gt ->
            let rr = remove x r in
            if r == rr then m else bal l v d rr)

  let rec update : type a. a key -> (a value option -> a value option) -> t -> t
      =
   fun x f -> function
    | Empty -> (
        match f None with
        | None -> Empty
        | Some data -> Node { l = Empty; v = x; d = data; r = Empty; h = 1 })
    | Node { l; v; d; r; h } as m -> (
        match Ord.compare x v with
        | Eq -> (
            match f (Some d) with
            | None -> merge l r
            | Some data ->
                if d == data then m else Node { l; v = x; d = data; r; h })
        | Lt ->
            let ll = update x f l in
            if l == ll then m else bal ll v d r
        | Gt ->
            let rr = update x f r in
            if r == rr then m else bal l v d rr)

  let rec iter f = function
    | Empty -> ()
    | Node { l; v; d; r; h = _ } ->
        iter f l;
        f.fun1 v d;
        iter f r

  type poly_mapi = { mapi_fun : 'a. 'a key -> 'a value -> 'a value } [@@unboxed]

  let rec mapi f = function
    | Empty -> Empty
    | Node { l; v; d; r; h } ->
        let l' = mapi f l in
        let d' = f.mapi_fun v d in
        let r' = mapi f r in
        Node { l = l'; v; d = d'; r = r'; h }

  let rec fold f m accu =
    match m with
    | Empty -> accu
    | Node { l; v; d; r; h = _ } -> fold f r (f.fun1 v d (fold f l accu))

  let rec for_all p = function
    | Empty -> true
    | Node { l; v; d; r; h = _ } -> p.fun1 v d && for_all p l && for_all p r

  let rec exists p = function
    | Empty -> false
    | Node { l; v; d; r; h = _ } -> p.fun1 v d || exists p l || exists p r

  (* Beware: those two functions assume that the added k is *strictly*
     smaller (or bigger) than all the present keys in the tree; it
     does not test for equality with the current min (or max) key.

     Indeed, they are only used during the "join" operation which
     respects this precondition.
  *)

  let rec add_min_binding k x = function
    | Empty -> singleton k x
    | Node { l; v; d; r; h = _ } -> bal (add_min_binding k x l) v d r

  let rec add_max_binding k x = function
    | Empty -> singleton k x
    | Node { l; v; d; r; h = _ } -> bal l v d (add_max_binding k x r)

  (* Same as create and bal, but no assumptions are made on the
     relative heights of l and r. *)

  let rec join l v d r =
    match (l, r) with
    | Empty, _ -> add_min_binding v d r
    | _, Empty -> add_max_binding v d l
    | ( Node { l = ll; v = lv; d = ld; r = lr; h = lh },
        Node { l = rl; v = rv; d = rd; r = rr; h = rh } ) ->
        if lh > rh + 2 then bal ll lv ld (join lr v d r)
        else if rh > lh + 2 then bal (join l v d rl) rv rd rr
        else create l v d r

  (* Merge two trees l and r into one.
     All elements of l must precede the elements of r.
     No assumption on the heights of l and r. *)

  let concat t1 t2 =
    match (t1, t2) with
    | Empty, t -> t
    | t, Empty -> t
    | _, _ ->
        let (Binding (x, d)) = min_binding t2 in
        join t1 x d (remove_min_binding t2)

  let concat_or_join t1 v d t2 =
    match d with Some d -> join t1 v d t2 | None -> concat t1 t2

  let rec split : type a. a key -> t -> t * a value option * t =
   fun x -> function
    | Empty -> (Empty, None, Empty)
    | Node { l; v; d; r; h = _ } -> (
        match Ord.compare x v with
        | Eq -> (l, Some d, r)
        | Lt ->
            let ll, pres, rl = split x l in
            (ll, pres, join rl v d r)
        | Gt ->
            let lr, pres, rr = split x r in
            (join l v d lr, pres, rr))

  type poly_merge = {
    merge_fun :
      'a. 'a key -> 'a value option -> 'a value option -> 'a value option;
  }
  [@@unboxed]

  let rec merge f s1 s2 =
    match (s1, s2) with
    | Empty, Empty -> Empty
    | Node { l = l1; v = v1; d = d1; r = r1; h = h1 }, _ when h1 >= height s2 ->
        let l2, d2, r2 = split v1 s2 in
        concat_or_join (merge f l1 l2) v1
          (f.merge_fun v1 (Some d1) d2)
          (merge f r1 r2)
    | _, Node { l = l2; v = v2; d = d2; r = r2; h = _ } ->
        let l1, d1, r1 = split v2 s1 in
        concat_or_join (merge f l1 l2) v2
          (f.merge_fun v2 d1 (Some d2))
          (merge f r1 r2)
    | _ -> assert false

  type poly_union = {
    union_fun : 'a. 'a key -> 'a value -> 'a value -> 'a value option;
  }
  [@@unboxed]

  let rec union f s1 s2 =
    match (s1, s2) with
    | Empty, s | s, Empty -> s
    | ( Node { l = l1; v = v1; d = d1; r = r1; h = h1 },
        Node { l = l2; v = v2; d = d2; r = r2; h = h2 } ) -> (
        if h1 >= h2 then
          let l2, d2, r2 = split v1 s2 in
          let l = union f l1 l2 and r = union f r1 r2 in
          match d2 with
          | None -> join l v1 d1 r
          | Some d2 -> concat_or_join l v1 (f.union_fun v1 d1 d2) r
        else
          let l1, d1, r1 = split v2 s1 in
          let l = union f l1 l2 and r = union f r1 r2 in
          match d1 with
          | None -> join l v2 d2 r
          | Some d1 -> concat_or_join l v2 (f.union_fun v2 d1 d2) r)

  let rec filter p = function
    | Empty -> Empty
    | Node { l; v; d; r; h = _ } as m ->
        (* call [p] in the expected left-to-right order *)
        let l' = filter p l in
        let pvd = p.fun1 v d in
        let r' = filter p r in
        if pvd then if l == l' && r == r' then m else join l' v d r'
        else concat l' r'

  type poly_filter_map = {
    filter_map_fun : 'a. 'a key -> 'a value -> 'a value option;
  }
  [@@unboxed]

  let rec filter_map f = function
    | Empty -> Empty
    | Node { l; v; d; r; h = _ } -> (
        (* call [f] in the expected left-to-right order *)
        let l' = filter_map f l in
        let fvd = f.filter_map_fun v d in
        let r' = filter_map f r in
        match fvd with Some d' -> join l' v d' r' | None -> concat l' r')

  let rec partition p = function
    | Empty -> (Empty, Empty)
    | Node { l; v; d; r; h = _ } ->
        (* call [p] in the expected left-to-right order *)
        let lt, lf = partition p l in
        let pvd = p.fun1 v d in
        let rt, rf = partition p r in
        if pvd then (join lt v d rt, concat lf rf)
        else (concat lt rt, join lf v d rf)

  type enumeration =
    | End
    | More : 'a key * 'a value * t * enumeration -> enumeration

  let rec cons_enum m e =
    match m with
    | Empty -> e
    | Node { l; v; d; r; h = _ } -> cons_enum l (More (v, d, r, e))

  let compare cmp m1 m2 =
    let rec compare_aux e1 e2 =
      match (e1, e2) with
      | End, End -> 0
      | End, _ -> -1
      | _, End -> 1
      | More (v1, d1, r1, e1), More (v2, d2, r2, e2) -> (
          match Ord.compare v1 v2 with
          | Lt -> -1
          | Gt -> 1
          | Eq ->
              let c = cmp.fun2 v1 d1 d2 in
              if c <> 0 then c
              else compare_aux (cons_enum r1 e1) (cons_enum r2 e2))
    in
    compare_aux (cons_enum m1 End) (cons_enum m2 End)

  let equal cmp m1 m2 =
    let rec equal_aux e1 e2 =
      match (e1, e2) with
      | End, End -> true
      | End, _ -> false
      | _, End -> false
      | More (v1, d1, r1, e1), More (v2, d2, r2, e2) -> (
          match Ord.compare v1 v2 with
          | Eq ->
              cmp.fun2 v1 d1 d2 && equal_aux (cons_enum r1 e1) (cons_enum r2 e2)
          | _ -> false)
    in
    equal_aux (cons_enum m1 End) (cons_enum m2 End)

  let rec cardinal = function
    | Empty -> 0
    | Node { l; r; v = _; d = _; h = _ } -> cardinal l + 1 + cardinal r

  let rec bindings_aux accu = function
    | Empty -> accu
    | Node { l; v; d; r; h = _ } ->
        bindings_aux (Binding (v, d) :: bindings_aux accu r) l

  let bindings s = bindings_aux [] s
  let choose = min_binding
  let choose_opt = min_binding_opt
  let add_seq i m = Seq.fold_left (fun m (Binding (k, v)) -> add k v m) m i
  let of_seq i = add_seq i empty

  let rec seq_of_enum_ c () =
    match c with
    | End -> Seq.Nil
    | More (k, v, t, rest) ->
        Seq.Cons (Binding (k, v), seq_of_enum_ (cons_enum t rest))

  let to_seq m = seq_of_enum_ (cons_enum m End)

  let rec snoc_enum s e =
    match s with
    | Empty -> e
    | Node { l; v; d; r; h = _ } -> snoc_enum r (More (v, d, l, e))

  let rec rev_seq_of_enum_ c () =
    match c with
    | End -> Seq.Nil
    | More (k, v, t, rest) ->
        Seq.Cons (Binding (k, v), rev_seq_of_enum_ (snoc_enum t rest))

  let to_rev_seq c = rev_seq_of_enum_ (snoc_enum c End)

  let to_seq_from low m =
    let rec aux : type a. a key -> t -> enumeration -> enumeration =
     fun low m c ->
      match m with
      | Empty -> c
      | Node { l; v; d; r; _ } -> (
          match Ord.compare v low with
          | Eq -> More (v, d, r, c)
          | Lt -> aux low r c
          | Gt -> aux low l (More (v, d, r, c)))
    in
    seq_of_enum_ (aux low m End)
end

module IdVal = struct
  type 'a t = 'a
end

module Make (Ord : DORDERED) : S with type 'a key = 'a Ord.t =
  MakeWithValue (Ord) (IdVal)

(** Functor that converts a [DORDERED] into an ordered type *)
module ToOrdered (X : DORDERED) : sig
  type t = Hide : 'a X.t -> t

  val compare : t -> t -> int
end = struct
  type t = Hide : 'a X.t -> t

  let compare (Hide x1) (Hide x2) =
    match X.compare x1 x2 with Lt -> -1 | Eq -> 0 | Gt -> 1
end

(** Functor that creates a non-dependendent map for keys of a [DORDERED] type *)
module MakeMap (X : DORDERED) :
  Map.S
    with type key = ToOrdered(X).t
     and type 'a t = 'a Map.Make(ToOrdered(X)).t = struct
  include Map.Make (ToOrdered (X))
end

(** Functor that creates a set of values of a [DORDERED] type *)
module MakeSet (X : DORDERED) : Set.S with type t = Set.Make(ToOrdered(X)).t =
struct
  include Set.Make (ToOrdered (X))
end

(** Functor that extends the indices of a [DORDERED] type, given some
   type-level function *)
module Extend
    (X : DORDERED) (F : sig
      type !'a t
    end) : sig
  type _ t = Extend : 'a X.t -> 'a F.t t

  val compare : 'a t -> 'b t -> ('a, 'b) cmp
end = struct
  type _ t = Extend : 'a X.t -> 'a F.t t

  let compare (type a b) (Extend x1 : a t) (Extend x2 : b t) : (a, b) cmp =
    match X.compare x1 x2 with Lt -> Lt | Eq -> Eq | Gt -> Gt
end

(** Functor that extends the indices of a [DORDERED] type by adding
   some data to its left-hand side *)
module ExtendL
    (X : DORDERED) (Y : sig
      type t
    end) : sig
  type _ t = Extend : 'a X.t -> (Y.t * 'a) t

  val compare : 'a t -> 'b t -> ('a, 'b) cmp
end =
  Extend
    (X)
    (struct
      type 'a t = Y.t * 'a
    end)

(** Functor that extends the indices of a [DORDERED] type by adding
   some data to its right-hand side *)
module ExtendR
    (X : DORDERED) (Y : sig
      type t
    end) : sig
  type _ t = Extend : 'a X.t -> ('a * Y.t) t

  val compare : 'a t -> 'b t -> ('a, 'b) cmp
end =
  Extend
    (X)
    (struct
      type 'a t = 'a * Y.t
    end)