diff --git a/src/data/CCIntMap.ml b/src/data/CCIntMap.ml index 63a16c3a..227ea9a4 100644 --- a/src/data/CCIntMap.ml +++ b/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 diff --git a/src/data/CCIntMap.mli b/src/data/CCIntMap.mli index d3622db9..c4a7901d 100644 --- a/src/data/CCIntMap.mli +++ b/src/data/CCIntMap.mli @@ -66,6 +66,10 @@ 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 + +val map : ('a -> 'b) -> 'a t -> 'b t + val choose : 'a t -> (int * 'a) option val choose_exn : 'a t -> int * 'a diff --git a/src/data/CCTrie.ml b/src/data/CCTrie.ml index 36b0ddf3..c161c7d7 100644 --- a/src/data/CCTrie.ml +++ b/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 *) @@ -356,6 +362,33 @@ module Make(W : WORD) = struct T.fold (fun acc k v -> (k,v) :: acc) [] t1 \ |> 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 + + 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 let iter f t = _fold diff --git a/src/data/CCTrie.mli b/src/data/CCTrie.mli index cc0c7505..5c49c1a5 100644 --- a/src/data/CCTrie.mli +++ b/src/data/CCTrie.mli @@ -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, giving both key and value. Will use {!WORD.of_list} to rebuild keys. *) + + val map : ('a -> 'b) -> 'a t -> 'b t + (** Map values, giving only the value. *) + val iter : (key -> 'a -> unit) -> 'a t -> unit (** Same as {!fold}, but for effectful functions *) diff --git a/src/data/CCWBTree.ml b/src/data/CCWBTree.ml index d7c0b895..bdff0d45 100644 --- a/src/data/CCWBTree.ml +++ b/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, _) -> diff --git a/src/data/CCWBTree.mli b/src/data/CCWBTree.mli index f1f89065..a419de1b 100644 --- a/src/data/CCWBTree.mli +++ b/src/data/CCWBTree.mli @@ -62,6 +62,12 @@ 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. *) + + val map : f:('a -> 'b) -> 'a t -> 'b t + (** Map values, giving only the value. *) + val iter : f:(key -> 'a -> unit) -> 'a t -> unit val split : key -> 'a t -> 'a t * 'a option * 'a t