[sledge] Add Term.map_trms

Reviewed By: ngorogiannis

Differential Revision: D24532365

fbshipit-source-id: c5365b727

sledge/src/fol.ml

@@ -317,10 +317,28 @@ module Term = struct
         let t' = Trm.map_vars ~f t in
         if t' == t then c else `Trm t'
 
-  let map_vars ~f = function
+  let map_vars e ~f =
+    match e with
     | `Fml p -> `Fml (Fml.map_vars ~f p)
     | #cnd as c -> (map_vars_c ~f c :> exp)
 
+  let rec map_trms_c ~f c =
+    match c with
+    | `Ite (cnd, thn, els) ->
+        let cnd' = Fml.map_trms ~f cnd in
+        let thn' = map_trms_c ~f thn in
+        let els' = map_trms_c ~f els in
+        if cnd' == cnd && thn' == thn && els' == els then c
+        else _Ite cnd' thn' els'
+    | `Trm t ->
+        let t' = f t in
+        if t' == t then c else `Trm t'
+
+  let map_trms e ~f =
+    match e with
+    | `Fml p -> `Fml (Fml.map_trms ~f p)
+    | #cnd as c -> (map_trms_c ~f c :> exp)
+
   let fold_map_vars e s0 ~f =
     let s = ref s0 in
     let f x =

sledge/src/fol.mli

@@ -75,7 +75,8 @@ module rec Term : sig
 
   (** Transform *)
 
-  val map_vars : f:(Var.t -> Var.t) -> t -> t
+  val map_vars : t -> f:(Var.t -> Var.t) -> t
+  val map_trms : t -> f:(Trm.t -> Trm.t) -> t
   val fold_map_vars : t -> 's -> f:(Var.t -> 's -> Var.t * 's) -> t * 's
   val rename : Var.Subst.t -> t -> t
 end