Adding a generic Int.Map using persistent arrays.

2 files changed, 241 insertions, 0 deletions

diff --git a/lib/int.ml b/lib/int.ml
index 952d04e9bd..43d6d037d2 100644
--- a/lib/int.ml
+++ b/lib/int.ml
@@ -47,3 +47,192 @@ module List = struct
   let mem_assoc = List.mem_assq
   let remove_assoc = List.remove_assq
 end
+
+let min (i : int) j = if i < j then i else j
+let max (i : int) j = if i > j then i else j
+
+(** Utility function *)
+let rec next from upto =
+  if from < upto then next (2 * from + 1) upto
+  else from
+
+
+module PArray =
+struct
+
+  type 'a t = 'a data ref
+  and 'a data =
+  | Root of 'a option array
+  | DSet of int * 'a option * 'a t
+
+  let empty n = ref (Root (Array.create n None))
+
+  let rec rerootk t k = match !t with
+  | Root _ -> k ()
+  | DSet (i, v, t') ->
+    let next () = match !t' with
+    | Root a as n ->
+      let v' = Array.unsafe_get a i in
+      let () = Array.unsafe_set a i v in
+      let () = t := n in
+      let () = t' := DSet (i, v', t) in
+      k ()
+    | DSet _ -> assert false
+    in
+    rerootk t' next
+
+  let reroot t = rerootk t (fun () -> ())
+
+  let rec get t i =
+  let () = assert (0 <= i) in
+  match !t with
+  | Root a ->
+    if Array.length a <= i then None
+    else Array.unsafe_get a i
+  | DSet _ ->
+    let () = reroot t in
+    match !t with
+    | Root a ->
+      if Array.length a <= i then None
+      else Array.unsafe_get a i
+    | DSet _ -> assert false
+
+  let set t i v =
+    let () = assert (0 <= i) in
+    let () = reroot t in
+    match !t with
+    | DSet _ -> assert false
+    | Root a as n ->
+      let len = Array.length a in
+      if i < len then
+        let old = Array.unsafe_get a i in
+        if old == v then t
+        else
+          let () = Array.unsafe_set a i v in
+          let res = ref n in
+          let () = t := DSet (i, old, res) in
+          res
+      else match v with
+      | None -> t (** Nothing to do! *)
+      | Some _ -> (** we must resize *)
+        let nlen = next len (succ i) in
+        let nlen = min nlen Sys.max_array_length in
+        let () = assert (i < nlen) in
+        let a' = Array.make nlen None in
+        let () = Array.blit a 0 a' 0 len in
+        let () = Array.unsafe_set a' i v in
+        let res = ref (Root a') in
+        let () = t := DSet (i, None, res) in
+        res
+
+end
+
+module PMap =
+struct
+
+  type key = int
+
+  (** Invariants:
+
+    1. an empty map is always [Empty].
+    2. the set of the [Map] constructor remembers the present keys.
+  *)
+  type 'a t = Empty | Map of Set.t * 'a PArray.t
+
+  let empty = Empty
+
+  let is_empty = function
+  | Empty -> true
+  | Map _ -> false
+
+  let singleton k x =
+    let len = next 19 (k + 1) in
+    let len = min Sys.max_array_length len in
+    let v = PArray.empty len in
+    let v = PArray.set v k (Some x) in
+    let s = Set.singleton k in
+    Map (s, v)
+
+  let add k x = function
+  | Empty -> singleton k x
+  | Map (s, v) ->
+    let s = match PArray.get v k with
+    | None -> Set.add k s
+    | Some _ -> s
+    in
+    let v = PArray.set v k (Some x) in
+    Map (s, v)
+
+  let remove k = function
+  | Empty -> Empty
+  | Map (s, v) ->
+    let s = Set.remove k s in
+    if Set.is_empty s then Empty
+    else
+      let v = PArray.set v k None in
+      Map (s, v)
+
+  let mem k = function
+  | Empty -> false
+  | Map (_, v) ->
+    match PArray.get v k with
+    | None -> false
+    | Some _ -> true
+
+  let find k = function
+  | Empty -> raise Not_found
+  | Map (_, v) ->
+    match PArray.get v k with
+    | None -> raise Not_found
+    | Some x -> x
+
+  let iter f = function
+  | Empty -> ()
+  | Map (s, v) ->
+    let iter k = match PArray.get v k with
+    | None -> ()
+    | Some x -> f k x
+    in
+    Set.iter iter s
+
+  let fold f m accu = match m with
+  | Empty -> accu
+  | Map (s, v) ->
+    let fold k accu = match PArray.get v k with
+    | None -> accu
+    | Some x -> f k x accu
+    in
+    Set.fold fold s accu
+
+  let exists f m = match m with
+  | Empty -> false
+  | Map (s, v) ->
+    let exists k = match PArray.get v k with
+    | None -> false
+    | Some x -> f k x
+    in
+    Set.exists exists s
+
+  let for_all f m = match m with
+  | Empty -> true
+  | Map (s, v) ->
+    let for_all k = match PArray.get v k with
+    | None -> true
+    | Some x -> f k x
+    in
+    Set.for_all for_all s
+
+  let cast = function
+  | Empty -> Map.empty
+  | Map (s, v) ->
+    let bind k = match PArray.get v k with
+    | None -> assert false
+    | Some x -> x
+    in
+    Map.bind bind s
+
+  let domain = function
+  | Empty -> Set.empty
+  | Map (s, _) -> s
+
+end
diff --git a/lib/int.mli b/lib/int.mli
index a84798950d..bc8896b71e 100644
--- a/lib/int.mli
+++ b/lib/int.mli
@@ -25,3 +25,55 @@ module List : sig
   val mem_assoc : int -> (int * 'a) list -> bool
   val remove_assoc : int -> (int * 'a) list -> (int * 'a) list
 end
+
+module PArray :
+sig
+  type 'a t
+  (** Persistent, auto-resizable arrays. The [get] and [set] functions never
+      fail whenever the index is between [0] and [Sys.max_array_length - 1]. *)
+  val empty : int -> 'a t
+  (** The empty array, with a given starting size. *)
+  val get : 'a t -> int -> 'a option
+  (** Get a value at the given index. Returns [None] if undefined. *)
+  val set : 'a t -> int -> 'a option -> 'a t
+  (** Set/unset a value at the given index. *)
+end
+
+module PMap :
+sig
+  type key = int
+  type 'a t
+  val empty : 'a t
+  val is_empty : 'a t -> bool
+  val mem : key -> 'a t -> bool
+  val add : key -> 'a -> 'a t -> 'a t
+  val singleton : key -> 'a -> 'a t
+  val remove : key -> 'a t -> 'a t
+(*   val merge : (key -> 'a option -> 'b option -> 'c option) -> 'a t -> 'b t -> 'c t *)
+(*   val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int *)
+(*   val equal : ('a -> 'a -> bool) -> 'a t -> 'a t -> bool *)
+  val iter : (key -> 'a -> unit) -> 'a t -> unit
+  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
+  val for_all : (key -> 'a -> bool) -> 'a t -> bool
+  val exists : (key -> 'a -> bool) -> 'a t -> bool
+(*   val filter : (key -> 'a -> bool) -> 'a t -> 'a t *)
+(*   val partition : (key -> 'a -> bool) -> 'a t -> 'a t * 'a t *)
+(*   val cardinal : 'a t -> int *)
+(*   val bindings : 'a t -> (key * 'a) list *)
+(*   val min_binding : 'a t -> key * 'a *)
+(*   val max_binding : 'a t -> key * 'a *)
+(*   val choose : 'a t -> key * 'a *)
+(*   val split : key -> 'a t -> 'a t * 'a option * 'a t *)
+  val find : key -> 'a t -> 'a
+(*   val map : ('a -> 'b) -> 'a t -> 'b t *)
+(*   val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t *)
+  val domain : 'a t -> Set.t
+  val cast : 'a t -> 'a Map.t
+end
+(** This is a (partial) implementation of a [Map] interface on integers, except
+    that it internally uses persistent arrays. This ensures O(1) accesses in
+    non-backtracking cases. It is thus better suited for zero-starting,
+    contiguous keys, or otherwise a lot of space will be empty. To keep track of
+    the present keys, a binary tree is also used, so that adding a key is
+    still logarithmic. It is therefore essential that most of the operations
+    are accesses and not add/removes. *)