[sledge] Add Trm.fold_map

Reviewed By: jvillard

Differential Revision: D25883707

fbshipit-source-id: 5c9b8a1cb

sledge/nonstdlib/NS.mli

@@ -92,6 +92,13 @@ val map3 :
 
 val mapN : ('a -> 'a) -> 'b -> ('a array -> 'b) -> 'a array -> 'b
 
+val fold_map_from_map :
+     ('a -> f:('b -> 'c) -> 'd)
+  -> 'a
+  -> 's
+  -> f:('b -> 's -> 'c * 's)
+  -> 'd * 's
+
 (** Pretty-printing *)
 
 (** Pretty-printer for argument type. *)

sledge/nonstdlib/NS0.ml

@@ -119,6 +119,16 @@ let map_endo map t ~f =
   in
   if !change then t' else t
 
+let fold_map_from_map map x s ~f =
+  let s = ref s in
+  let f y =
+    let y', s' = f y !s in
+    s := s' ;
+    y'
+  in
+  let x' = map x ~f in
+  (x', !s)
+
 (** Containers *)
 
 (* from upcoming Stdlib *)

sledge/src/fol/trm.ml

@@ -492,3 +492,5 @@ let map e ~f =
       map3 f e (fun seq off len -> _Extract ~seq ~off ~len) seq off len
   | Concat xs -> mapN f e _Concat xs
   | Apply (g, xs) -> mapN f e (_Apply g) xs
+
+let fold_map e = fold_map_from_map map e

sledge/src/fol/trm.mli

@@ -118,3 +118,6 @@ val map_vars : t -> f:(Var.t -> Var.t) -> t
 
 val map : t -> f:(t -> t) -> t
 (** Map over the {!trms}. *)
+
+val fold_map : t -> 's -> f:(t -> 's -> t * 's) -> t * 's
+(** Fold while mapping over the {!trms}. *)