From b9bb3ca220ff947805df99f0ff0a0da54afbbd06 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Wed, 2 Dec 2020 13:46:26 -0800
Subject: [PATCH] [sledge] Add Iter.fold_map and folding_map

Reviewed By: jvillard

Differential Revision: D25146151

fbshipit-source-id: 236b3bb57
---
 sledge/nonstdlib/iter.ml  | 18 ++++++++++++++++++
 sledge/nonstdlib/iter.mli |  3 +++
 2 files changed, 21 insertions(+)

diff --git a/sledge/nonstdlib/iter.ml b/sledge/nonstdlib/iter.ml
index 5b7adae07..bcc77cf40 100644
--- a/sledge/nonstdlib/iter.ml
+++ b/sledge/nonstdlib/iter.ml
@@ -66,6 +66,24 @@ let fold_opt seq s ~f =
     Some !state
   with Stop -> None
 
+let folding_map seq s ~f =
+  fold_map
+    ~f:(fun s x ->
+      let y, s = f x s in
+      (s, y) )
+    ~init:s seq
+
+let fold_map seq s ~f =
+  let r = ref s in
+  let seq' =
+    persistent (fun yield ->
+        seq (fun x ->
+            let y, s = f x !r in
+            r := s ;
+            yield y ) )
+  in
+  (!r, seq')
+
 let fold_until (type res) seq s ~f ~finish =
   let state = ref s in
   let exception Stop of res in
diff --git a/sledge/nonstdlib/iter.mli b/sledge/nonstdlib/iter.mli
index 06d866a74..6843e600d 100644
--- a/sledge/nonstdlib/iter.mli
+++ b/sledge/nonstdlib/iter.mli
@@ -74,6 +74,9 @@ val fold_opt : 'a t -> 's -> f:('a -> 's -> 's option) -> 's option
     the [Option] monad. If [f] returns [None], [None] is returned without
     any additional invocations of [f]. *)
 
+val fold_map : 'a t -> 's -> f:('a -> 's -> 'b * 's) -> 's * 'b t
+val folding_map : 'a t -> 's -> f:('a -> 's -> 'b * 's) -> 'b t
+
 val fold_until :
      'a t
   -> 's