From 63a8f22eeff53ad1e9ed36e3b9293bf400879e07 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 12 Jan 2021 04:26:43 -0800
Subject: [PATCH] [sledge] Optimize Context.Subst.compose1

Summary:
Trade a bit more code for lowering complexity from linear to
logarithmic.

Reviewed By: jvillard

Differential Revision: D25756569

fbshipit-source-id: 83ad575fe
---
 sledge/src/fol/context.ml | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/sledge/src/fol/context.ml b/sledge/src/fol/context.ml
index b363a6f65..a4ab7b8ee 100644
--- a/sledge/src/fol/context.ml
+++ b/sledge/src/fol/context.ml
@@ -118,16 +118,17 @@ end = struct
       assert (r' == r || not (equal r' r))]
 
   (** compose a substitution with a mapping *)
-  let compose1 ~key ~data s =
+  let compose1 ~key ~data r =
     match (key : Trm.t) with
-    | Z _ | Q _ -> s
+    | Z _ | Q _ -> r
+    | _ when Trm.equal key data -> r
     | _ ->
-        if Trm.equal key data then s
-        else (
-          assert (
-            (not (Trm.Map.mem key s))
-            || fail "domains intersect: %a" Trm.pp key () ) ;
-          compose s (Trm.Map.singleton key data) )
+        assert (
+          (not (Trm.Map.mem key r))
+          || fail "domains intersect: %a" Trm.pp key () ) ;
+        let s = Trm.Map.singleton key data in
+        let r' = Trm.Map.map_endo ~f:(norm s) r in
+        Trm.Map.add ~key ~data r'
 
   (** add an identity entry if the term is not already present *)
   let extend e s =