From e91fc9d1a3a3efb1d0b121c2927e4984728a2fc9 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 2 Feb 2021 04:35:40 -0800
Subject: [PATCH] [sledge] Improve Context.extend with Theory.classify

Reviewed By: jvillard

Differential Revision: D25883714

fbshipit-source-id: 14c54a379
---
 sledge/src/fol/context.ml | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/sledge/src/fol/context.ml b/sledge/src/fol/context.ml
index cdf19f694..e6a2074a1 100644
--- a/sledge/src/fol/context.ml
+++ b/sledge/src/fol/context.ml
@@ -406,16 +406,18 @@ let rec extend_ a s =
     ~call:(fun {pf} -> pf "@ %a@ %a" Trm.pp a Subst.pp s)
     ~retn:(fun {pf} s' -> pf "%a" Subst.pp_diff (s, s'))
   @@ fun () ->
-  match (a : Trm.t) with
-  | Z _ | Q _ -> s
-  | _ -> (
-      if Theory.is_interpreted a then Iter.fold ~f:extend_ (Trm.trms a) s
-      else
-        (* add uninterpreted terms *)
-        match Subst.extend a s with
-        (* and their subterms if newly added *)
-        | Some s -> Iter.fold ~f:extend_ (Trm.trms a) s
-        | None -> s )
+  match Theory.classify a with
+  | InterpAtom -> s
+  | InterpApp -> Iter.fold ~f:extend_ (Trm.trms a) s
+  | NonInterpAtom ->
+      (* add uninterpreted terms if missing *)
+      Trm.Map.change a s ~f:(function None -> Some a | a' -> a')
+  | UninterpApp -> (
+    (* add uninterpreted terms if missing *)
+    match Subst.extend a s with
+    | None -> s
+    (* and their subterms if newly added *)
+    | Some s -> Iter.fold ~f:extend_ (Trm.trms a) s )
 
 (** add a term to the carrier *)
 let extend a x =