From 54f6b9e974561398dc6d899fd525f1f7ecdc7297 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 27 Oct 2020 12:17:18 -0700
Subject: [PATCH] [sledge] Add iterator of subterms of Trm.t

Reviewed By: jvillard

Differential Revision: D24532347

fbshipit-source-id: c0686ae92
---
 sledge/src/trm.ml  | 18 ++++++++++++++++++
 sledge/src/trm.mli |  1 +
 2 files changed, 19 insertions(+)

diff --git a/sledge/src/trm.ml b/sledge/src/trm.ml
index ac94572aa..baed849ca 100644
--- a/sledge/src/trm.ml
+++ b/sledge/src/trm.ml
@@ -447,3 +447,21 @@ let rec map_vars e ~f =
   | Record xs -> mapN (map_vars ~f) e _Record xs
   | Ancestor _ -> e
   | Apply (g, xs) -> mapN (map_vars ~f) e (_Apply g) xs
+
+(** Traverse *)
+
+let iter_subtrms e ~f =
+  match e with
+  | Var _ | Z _ | Q _ | Ancestor _ -> ()
+  | Arith a -> Iter.iter ~f (Arith.trms a)
+  | Splat x | Select {rcd= x} -> f x
+  | Sized {seq= x; siz= y} | Update {rcd= x; elt= y} ->
+      f x ;
+      f y
+  | Extract {seq= x; off= y; len= z} ->
+      f x ;
+      f y ;
+      f z
+  | Concat xs | Record xs | Apply (_, xs) -> Array.iter ~f xs
+
+let subtrms e = Iter.from_labelled_iter (iter_subtrms e)
diff --git a/sledge/src/trm.mli b/sledge/src/trm.mli
index 7410b5e43..b341b7d3c 100644
--- a/sledge/src/trm.mli
+++ b/sledge/src/trm.mli
@@ -89,3 +89,4 @@ val seq_size : t -> t option
 (** Traverse *)
 
 val vars : t -> Var.t iter
+val subtrms : t -> t iter