From e61278cb9ab111c6114ff3440e129a7dfb4eaded Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 2 Feb 2021 04:36:03 -0800
Subject: [PATCH] [sledge] Add Trm.fold_map

Reviewed By: jvillard

Differential Revision: D25883707

fbshipit-source-id: 5c9b8a1cb
---
 sledge/nonstdlib/NS.mli |  7 +++++++
 sledge/nonstdlib/NS0.ml | 10 ++++++++++
 sledge/src/fol/trm.ml   |  2 ++
 sledge/src/fol/trm.mli  |  3 +++
 4 files changed, 22 insertions(+)

diff --git a/sledge/nonstdlib/NS.mli b/sledge/nonstdlib/NS.mli
index 08518cc2a..ab5bd3b67 100644
--- a/sledge/nonstdlib/NS.mli
+++ b/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. *)
diff --git a/sledge/nonstdlib/NS0.ml b/sledge/nonstdlib/NS0.ml
index 4d8887d29..8f388beb9 100644
--- a/sledge/nonstdlib/NS0.ml
+++ b/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 *)
diff --git a/sledge/src/fol/trm.ml b/sledge/src/fol/trm.ml
index ebcea49d7..849e6d3fa 100644
--- a/sledge/src/fol/trm.ml
+++ b/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
diff --git a/sledge/src/fol/trm.mli b/sledge/src/fol/trm.mli
index 7a91149b9..fa5dbfcc9 100644
--- a/sledge/src/fol/trm.mli
+++ b/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}. *)