From 6f5f42b3e33f2dbe5d3aa3203f7dcda4698db13c Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Thu, 29 Oct 2020 13:26:59 -0700
Subject: [PATCH] [sledge] Add Trm.map

Reviewed By: ngorogiannis

Differential Revision: D24532358

fbshipit-source-id: a0e13284f
---
 sledge/src/trm.ml  | 20 +++++++++++++++++---
 sledge/src/trm.mli |  1 +
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/sledge/src/trm.ml b/sledge/src/trm.ml
index f68ed9355..ef9bf4204 100644
--- a/sledge/src/trm.ml
+++ b/sledge/src/trm.ml
@@ -446,9 +446,7 @@ let rec map_vars e ~f =
   match e with
   | Var _ as v -> (f (Var.of_ v) : Var.t :> t)
   | Z _ | Q _ -> e
-  | Arith a ->
-      let a' = Arith.map ~f:(map_vars ~f) a in
-      if a == a' then e else _Arith a'
+  | Arith a -> map1 (Arith.map ~f:(map_vars ~f)) e _Arith a
   | Splat x -> map1 (map_vars ~f) e _Splat x
   | Sized {seq; siz} -> map2 (map_vars ~f) e _Sized seq siz
   | Extract {seq; off; len} -> map3 (map_vars ~f) e _Extract seq off len
@@ -459,6 +457,22 @@ let rec map_vars e ~f =
   | Ancestor _ -> e
   | Apply (g, xs) -> mapN (map_vars ~f) e (_Apply g) xs
 
+let map e ~f =
+  match e with
+  | Var _ | Z _ | Q _ -> e
+  | Arith a ->
+      let a' = Arith.map ~f a in
+      if a == a' then e else _Arith a'
+  | Splat x -> map1 f e _Splat x
+  | Sized {seq; siz} -> map2 f e _Sized seq siz
+  | Extract {seq; off; len} -> map3 f e _Extract seq off len
+  | Concat xs -> mapN f e _Concat xs
+  | Select {idx; rcd} -> map1 f e (_Select idx) rcd
+  | Update {idx; rcd; elt} -> map2 f e (_Update idx) rcd elt
+  | Record xs -> mapN f e _Record xs
+  | Ancestor _ -> e
+  | Apply (g, xs) -> mapN f e (_Apply g) xs
+
 (** Traverse *)
 
 let iter_subtrms e ~f =
diff --git a/sledge/src/trm.mli b/sledge/src/trm.mli
index 745e53ea6..e178734ed 100644
--- a/sledge/src/trm.mli
+++ b/sledge/src/trm.mli
@@ -88,6 +88,7 @@ val apply : Ses.Funsym.t -> t array -> t
 (** Transform *)
 
 val map_vars : t -> f:(Var.t -> Var.t) -> t
+val map : t -> f:(t -> t) -> t
 
 (** Query *)