From 2fc8bc1f842e7dbc2ce1deb58847f6080ea565d1 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Thu, 16 Apr 2020 03:39:36 -0700
Subject: [PATCH] [sledge] Export Term.map_rec_pre and add
 Term.fold_map_rec_pre

Summary:
These are map and folding map that perform a cycle-preserving
pre-order transformation.

Reviewed By: jvillard

Differential Revision: D20877974

fbshipit-source-id: 251288228
---
 sledge/lib/import/import.ml  |  9 +++++++++
 sledge/lib/import/import.mli |  3 +++
 sledge/lib/term.ml           | 33 ++++++++++++++++++++++++++-------
 sledge/lib/term.mli          | 13 +++++++++++++
 4 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/sledge/lib/import/import.ml b/sledge/lib/import/import.ml
index 517a1bd18..5d2371e5e 100644
--- a/sledge/lib/import/import.ml
+++ b/sledge/lib/import/import.ml
@@ -77,6 +77,15 @@ module Array = struct
   include Core.Array
 
   let pp sep pp_elt fs a = List.pp sep pp_elt fs (to_list a)
+
+  let fold_map_inplace a ~init ~f =
+    let s = ref init in
+    let f x =
+      let s', x' = f !s x in
+      s := s' ;
+      x'
+    in
+    map_inplace a ~f ; !s
 end
 
 module IArray = IArray
diff --git a/sledge/lib/import/import.mli b/sledge/lib/import/import.mli
index 547c444af..42b346f24 100644
--- a/sledge/lib/import/import.mli
+++ b/sledge/lib/import/import.mli
@@ -76,6 +76,9 @@ module Array : sig
   include module type of Array
 
   val pp : (unit, unit) fmt -> 'a pp -> 'a array pp
+
+  val fold_map_inplace :
+    'a array -> init:'s -> f:('s -> 'a -> 's * 'a) -> 's
 end
 
 module IArray = IArray
diff --git a/sledge/lib/term.ml b/sledge/lib/term.ml
index 0d81a2a33..26b90ffa4 100644
--- a/sledge/lib/term.ml
+++ b/sledge/lib/term.ml
@@ -1132,13 +1132,7 @@ let fold_map e ~init ~f =
   let e' = map e ~f in
   (!s, e')
 
-(** Pre-order transformation that preserves cycles. Each subterm [x] from
-    root to leaves is presented to [f]. If [f x = Some x'] then the subterms
-    of [x] are not traversed and [x] is transformed to [x']. Otherwise
-    traversal proceeds to the subterms of [x], followed by rebuilding the
-    term structure on the transformed subterms. Cycles (through terms
-    involving [RecN]) are preserved. *)
-let map_rec_pre ~f e =
+let map_rec_pre e ~f =
   let rec map_rec_pre_f memo e =
     match f e with
     | Some e' -> e'
@@ -1161,6 +1155,31 @@ let map_rec_pre ~f e =
   in
   map_rec_pre_f [] e
 
+let fold_map_rec_pre e ~init ~f =
+  let rec fold_map_rec_pre_f memo s e =
+    match f s e with
+    | Some (s, e') -> (s, e')
+    | None -> (
+      match e with
+      | RecN (op, xs) -> (
+        match List.Assoc.find ~equal:( == ) memo e with
+        | None ->
+            let xs' = IArray.to_array xs in
+            let e' = RecN (op, IArray.of_array xs') in
+            let memo = List.Assoc.add ~equal:( == ) memo e e' in
+            let changed = ref false in
+            let s =
+              Array.fold_map_inplace ~init:s xs' ~f:(fun s x ->
+                  let s, x' = fold_map_rec_pre_f memo s x in
+                  if x' != x then changed := true ;
+                  (s, x') )
+            in
+            if !changed then (s, e') else (s, e)
+        | Some e' -> (s, e') )
+      | _ -> fold_map ~f:(fold_map_rec_pre_f memo) ~init:s e )
+  in
+  fold_map_rec_pre_f [] init e
+
 let rename sub e =
   map_rec_pre e ~f:(function
     | Var _ as v -> Some (Var.Subst.apply sub v)
diff --git a/sledge/lib/term.mli b/sledge/lib/term.mli
index a7bd4ed01..afe6fd53e 100644
--- a/sledge/lib/term.mli
+++ b/sledge/lib/term.mli
@@ -235,7 +235,20 @@ val size_of : Typ.t -> t
 (** Transform *)
 
 val map : t -> f:(t -> t) -> t
+
+val map_rec_pre : t -> f:(t -> t option) -> t
+(** Pre-order transformation that preserves cycles. Each subterm [x] from
+    root to leaves is presented to [f]. If [f x = Some x'] then the subterms
+    of [x] are not traversed and [x] is transformed to [x']. Otherwise
+    traversal proceeds to the subterms of [x], followed by rebuilding the
+    term structure on the transformed subterms. Cycles (through terms
+    involving [RecN]) are preserved. *)
+
 val fold_map : t -> init:'a -> f:('a -> t -> 'a * t) -> 'a * t
+
+val fold_map_rec_pre :
+  t -> init:'a -> f:('a -> t -> ('a * t) option) -> 'a * t
+
 val rename : Var.Subst.t -> t -> t
 
 (** Traverse *)