From 2f1cb79013bc246a178d73e8ef22234cb9b15859 Mon Sep 17 00:00:00 2001
From: Danila Fedorin <danila.fedorin@gmail.com>
Date: Sun, 26 Nov 2023 12:27:44 -0800
Subject: [PATCH] Clean up the search and proving code somewhat

---
 src/Bergamot/Rules.elm  | 46 +++++++++++++++++++++++++++--------------
 src/Bergamot/Search.elm |  6 +++---
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/src/Bergamot/Rules.elm b/src/Bergamot/Rules.elm
index 928e400..90df3d8 100644
--- a/src/Bergamot/Rules.elm
+++ b/src/Bergamot/Rules.elm
@@ -33,6 +33,9 @@ andThen f p env ps =
     p env ps
     |> Search.andThen (\(a, psp) -> (f a) env psp)
 
+apply : Prover a -> Prover (a -> b) -> Prover b
+apply pa pf = pf |> andThen (\f -> map f pa)
+
 map : (a -> b) -> Prover a -> Prover b
 map f p env ps =
     p env ps
@@ -51,28 +54,41 @@ fail env ps = Search.fail
 getEnv : Prover RuleEnv
 getEnv env ps = Search.pure (env, ps)
 
-rule : Term UnificationVar -> Rule -> Prover ProofTree
-rule t r env ps =
+liftInstantiation : (a -> InstantiationState -> (b, InstantiationState)) -> a -> Prover b
+liftInstantiation f a env ps =
     let
-        (conc, is) = instantiate r.conclusion ps.instantiationState
-        (prems, isp) = instantiateList r.premises is
+        (b, is) = f a ps.instantiationState
     in
-        case unify t conc ps.unificationState of
-            Nothing -> Search.fail
-            Just (tp, usp) ->
-                let
-                    psp = { ps | instantiationState = resetVars isp
-                               , unificationState = usp }
-                in
-                    provePremises prems env psp
-                    |> Search.map (Tuple.mapFirst (\trees -> MkProofTree { name = r.name, conclusion = tp, premises = trees }))
+        Search.pure (b, { ps | instantiationState = is })
+
+liftUnification : (a -> a -> UnificationState -> Maybe (a, UnificationState)) -> a -> a -> Prover a
+liftUnification f a1 a2 env ps =
+    case f a1 a2 ps.unificationState of
+        Just (a, us) -> Search.pure (a, { ps | unificationState = us })
+        Nothing -> Search.fail
+
+withVars : Prover a -> Prover a
+withVars p env ps =
+    p env ps
+    |> Search.map (Tuple.mapSecond (\psp -> { psp | instantiationState = resetVars psp.instantiationState }))
+
+rule : Term UnificationVar -> Rule -> Prover ProofTree
+rule t r =
+        withVars (pure Tuple.pair
+            |> apply (liftInstantiation instantiate r.conclusion)
+            |> apply (liftInstantiation instantiateList r.premises))
+        |> andThen (\(conc, prems) ->
+            pure (\tp trees -> MkProofTree { name = r.name, conclusion = tp, premises = trees })
+            |> apply (liftUnification unify t conc)
+            |> apply (provePremises prems))
 
 provePremises : List (Term UnificationVar) -> Prover (List ProofTree)
 provePremises l =
     case l of
         t :: ts ->
-            prove t
-            |> andThen (\tree -> map (\trees -> tree :: trees) (provePremises ts))
+            pure (::)
+            |> apply (prove t)
+            |> apply (provePremises ts)
         [] -> pure []
 
 prove : Term UnificationVar -> Prover ProofTree
diff --git a/src/Bergamot/Search.elm b/src/Bergamot/Search.elm
index fe0007c..4338b72 100644
--- a/src/Bergamot/Search.elm
+++ b/src/Bergamot/Search.elm
@@ -12,11 +12,11 @@ map f s () =
         Empty -> Empty
         Found a sp -> Found (f a) (map f sp)
 
-apply : Search (a -> b) -> Search a -> Search b
-apply sf sa () =
+apply : Search a -> Search (a -> b) -> Search b
+apply sa sf () =
     case sf () of
         Empty -> Empty
-        Found f sfp -> interleave (map f sa) (apply sfp sa) ()
+        Found f sfp -> interleave (map f sa) (apply sa sfp) ()
 
 andThen : (a -> Search b) -> Search a -> Search b
 andThen f sa () =