Merge pull request #67 from johanneskloos/master

Provide `map` and `mapi` for some of the map types

src/data/CCIntMap.ml

@@ -287,6 +287,18 @@ let rec fold f t acc = match t with
 
 let cardinal t = fold (fun _ _ n -> n+1) t 0
 
+let rec mapi f t = match t with
+  | E -> E
+  | L (k, v) -> L (k, f k v)
+  | N (p, s, l, r) ->
+      N (p, s, mapi f l, mapi f r)
+
+let rec map f t = match t with
+  | E -> E
+  | L (k, v) -> L (k, f v)
+  | N (p, s, l, r) ->
+      N (p, s, map f l, map f r)
+
 let rec choose_exn = function
   | E -> raise Not_found
   | L (k, v) -> k, v

src/data/CCIntMap.mli

@@ -66,6 +66,12 @@ val iter : (int -> 'a -> unit) -> 'a t -> unit
 
 val fold : (int -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
 
+val mapi : (int -> 'a -> 'b) -> 'a t -> 'b t
+(** @since NEXT_RELEASE *)
+
+val map : ('a -> 'b) -> 'a t -> 'b t
+(** @since NEXT_RELEASE *)
+
 val choose : 'a t -> (int * 'a) option
 
 val choose_exn : 'a t -> int * 'a

src/data/CCTrie.ml

@@ -75,6 +75,12 @@ module type S = sig
   val fold : ('b -> key -> 'a -> 'b) -> 'b -> 'a t -> 'b
   (** Fold on key/value bindings. Will use {!WORD.of_list} to rebuild keys. *)
 
+  val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t
+  (** Map values in the try. Will use {!WORD.of_list} to rebuild keys. *)
+
+  val map : ('a -> 'b) -> 'a t -> 'b t
+  (** Map values in the try, not giving keys to the mapping function. *)
+
   val iter : (key -> 'a -> unit) -> 'a t -> unit
   (** Same as {!fold}, but for effectful functions *)
 
@@ -357,6 +363,45 @@ module Make(W : WORD) = struct
       |> List.sort Pervasives.compare = List.sort Pervasives.compare l1
   *)
 
+  let mapi f t =
+    let rec map_ prefix t = match t with
+      | Empty -> Empty
+      | Cons (c, t') -> Cons (c, map_ (_difflist_add prefix c) t')
+      | Node (v, map) ->
+          let v' = match v with
+            | None -> None
+            | Some v -> Some (f (W.of_list (prefix [])) v)
+          in let map' =
+            M.mapi (fun c t' ->
+                      let prefix' = _difflist_add prefix c in
+                        map_ prefix' t')
+              map
+          in Node (v', map')
+    in map_ _id t
+  (*$T
+    T.mapi (fun k v -> v ^ "!") t1 \
+    |> T.to_list |> List.sort Pervasives.compare =\
+       List.map (fun (k, v) -> (k, v ^ "!")) l1 |> List.sort Pervasives.compare
+  *)
+
+  let map f t =
+    let rec map_ = function
+      | Empty -> Empty
+      | Cons (c, t') -> Cons (c, map_ t')
+      | Node (v, map) ->
+          let v' = match v with
+            | None -> None
+            | Some v -> Some (f v)
+          in let map' = M.map map_ map
+          in Node (v', map')
+    in map_ t
+  (*$T
+    T.map (fun v -> v ^ "!") t1 \
+    |> T.to_list |> List.sort Pervasives.compare =\
+       List.map (fun (k, v) -> (k, v ^ "!")) l1 |> List.sort Pervasives.compare
+  *)
+
+
   let iter f t =
     _fold
       (fun () path y -> f (W.of_list (path [])) y)

src/data/CCTrie.mli

@@ -75,6 +75,16 @@ module type S = sig
   val fold : ('b -> key -> 'a -> 'b) -> 'b -> 'a t -> 'b
   (** Fold on key/value bindings. Will use {!WORD.of_list} to rebuild keys. *)
 
+  val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t
+  (** Map values, giving both key and value. Will use {!WORD.of_list} to rebuild keys.
+      @since NEXT_RELEASE
+   *)
+
+  val map : ('a -> 'b) -> 'a t -> 'b t
+  (** Map values, giving only the value.
+      @since NEXT_RELEASE
+    *)
+
   val iter : (key -> 'a -> unit) -> 'a t -> unit
   (** Same as {!fold}, but for effectful functions *)
 

src/data/CCWBTree.ml

@@ -97,6 +97,10 @@ module type S = sig
 
   val fold : f:('b -> key -> 'a -> 'b) -> x:'b -> 'a t -> 'b
 
+  val mapi : f:(key -> 'a -> 'b) -> 'a t -> 'b t
+
+  val map : f:('a -> 'b) -> 'a t -> 'b t
+
   val iter : f:(key -> 'a -> unit) -> 'a t -> unit
 
   val split : key -> 'a t -> 'a t * 'a option * 'a t
@@ -368,6 +372,16 @@ module MakeFull(K : KEY) : S with type key = K.t = struct
         let acc = f acc k v in
         fold ~f ~x:acc r
 
+  let rec mapi ~f = function
+    | E -> E
+    | N (k, v, l, r, w) ->
+        N (k, f k v, mapi ~f l, mapi ~f r, w)
+
+  let rec map ~f = function
+    | E -> E
+    | N (k, v, l, r, w) ->
+        N (k, f v, map ~f l, map ~f r, w)
+
   let rec iter ~f m = match m with
     | E -> ()
     | N (k, v, l, r, _) ->

src/data/CCWBTree.mli

@@ -62,6 +62,16 @@ module type S = sig
 
   val fold : f:('b -> key -> 'a -> 'b) -> x:'b -> 'a t -> 'b
 
+  val mapi : f:(key -> 'a -> 'b) -> 'a t -> 'b t
+  (** Map values, giving both key and value. Will use {!WORD.of_list} to rebuild keys.
+      @since NEXT_RELEASE
+    *)
+
+  val map : f:('a -> 'b) -> 'a t -> 'b t
+  (** Map values, giving only the value.
+      @since NEXT_RELEASE
+    *)
+
   val iter : f:(key -> 'a -> unit) -> 'a t -> unit
 
   val split : key -> 'a t -> 'a t * 'a option * 'a t