From 2cdeb2721c2f22732060c70c87dd4ed8ec83a71d Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Tue, 7 Jan 2025 20:59:59 -0800
Subject: [PATCH] Improve auto solving - 30% faster and hopefully proper errors
 if a type mismatch is blocking it.

---
 TODO.md                  |  4 +++-
 port/Lib/CompileExp.newt |  2 +-
 port/Lib/Elab.newt       | 28 +++++++++++++++++++++++-----
 port/Lib/TopContext.newt |  2 +-
 port/Lib/Types.newt      | 25 +++++++------------------
 src/Lib/CompileExp.idr   |  2 +-
 src/Lib/Elab.idr         | 23 +++++++++++++++++++++--
 src/Lib/TopContext.idr   |  2 +-
 src/Lib/Types.idr        |  4 +++-
 9 files changed, 61 insertions(+), 31 deletions(-)

diff --git a/TODO.md b/TODO.md
index 525d202..a5501ad 100644
--- a/TODO.md
+++ b/TODO.md
@@ -37,7 +37,9 @@ More comments in code! This is getting big enough that I need to re-find my bear
 - [x] SortedMap.newt issue in `where`
 - [x] fix "insufficient patterns", wire in M or Either String
 - [x] Matching _,_ when Maybe is expected should be an error
-- [ ] There are issues with matching inside do blocks, I think we need to guess scrutinee? I could imagine constraining metas too (e.g. if Just ... at ?m123 then ?m123 =?= Maybe ?m456)
+- [ ] There are issues with matching inside do blocks
+    - I may have handled a lot of this by reworking autos to check just one level of constraints (with a 20% speedup)
+    - Do we need to guess scrutinee? Or we could refine a meta scrutinee type to match the constructor being split. This happens during checking for pi, where we create ?m -> ?m and unifiy it with the meta. Could we do the same in a case, where `Just ...` says the meta must be `Maybe ?m`?
   - ~~`newt/Fixme.newt` - I can drop `do` _and_ put `bind {IO}` to push it along. It may need to try to solve autos earlier and better.~~
   - Also, the root cause is tough to track down if there is a type error (this happens with `do` in Idris, too).
   - Part of it is the auto solutions are getting discarded because of an unrelated unification failure. Either auto shouldn't go as deep or should run earlier. Forgetting that lookupMap returns a (k,v) pair is a good example.
diff --git a/port/Lib/CompileExp.newt b/port/Lib/CompileExp.newt
index b3ed461..af18403 100644
--- a/port/Lib/CompileExp.newt
+++ b/port/Lib/CompileExp.newt
@@ -115,7 +115,7 @@ compileTerm (Meta _ k) = pure $ CRef "meta$\{show k}" -- FIXME
 compileTerm (Lam _ nm _ _ t) = CLam nm <$> compileTerm t
 compileTerm tm@(App _ _ _) = case funArgs tm of
   (Meta _ k, args) => do
-        -- this will be undefined, should only happen for use metas
+        info (getFC tm) "Compiling an unsolved meta \{show tm}"
         pure $ CApp (CRef "Meta\{show k}") Nil 0
   (t@(Ref fc nm _), args) => do
         args' <- traverse compileTerm args
diff --git a/port/Lib/Elab.newt b/port/Lib/Elab.newt
index e6461e6..8ff9dd1 100644
--- a/port/Lib/Elab.newt
+++ b/port/Lib/Elab.newt
@@ -6,6 +6,7 @@ import Data.String
 import Data.IORef
 import Lib.Types
 import Lib.Eval
+import Lib.Util
 import Lib.TopContext
 import Lib.Syntax
 
@@ -137,7 +138,7 @@ findMatches ctx ty ((MkEntry _ name type def) :: xs) = do
   -- FIXME we're restoring state, but the INFO logs have already been emitted
   -- Also redo this whole thing to run during elab, recheck constraints, etc.
   mc <- readIORef top.metaCtx
-  catchError(do
+  catchError (do
       -- TODO sort out the FC here
       let fc = getFC ty
       debug $ \ _ => "TRY \{show name} : \{rpprint Nil type} for \{show ty}"
@@ -145,6 +146,9 @@ findMatches ctx ty ((MkEntry _ name type def) :: xs) = do
       -- tm <- check (mkCtx fc) (RVar fc name) ty
       -- FIXME RVar should optionally have qualified names
       let (QN ns nm) = name
+
+      let (cod, tele) = splitTele type
+      modifyIORef top.metaCtx $ \ mc => MC mc.metas mc.next CheckFirst
       tm <- check ctx (RVar fc nm) ty
       debug $ \ _ => "Found \{rpprint Nil tm} for \{show ty}"
       writeIORef top.metaCtx mc
@@ -276,7 +280,7 @@ updateMeta ix f = do
   top <- get
   mc <- readIORef top.metaCtx
   metas <- go mc.metas
-  writeIORef top.metaCtx $ MC metas mc.next
+  writeIORef top.metaCtx $ MC metas mc.next mc.mcmode
   where
     go : List MetaEntry -> M (List MetaEntry)
     go Nil = error' "Meta \{show ix} not found"
@@ -301,12 +305,14 @@ checkAutos ix (_ :: rest) = checkAutos ix rest
 
 addConstraint : Env -> Int -> SnocList Val -> Val -> M Unit
 addConstraint env ix sp tm = do
+  top <- get
+  mc <- readIORef top.metaCtx
+  let (CheckAll) = mc.mcmode | _ => pure MkUnit
   updateMeta ix $ \case
     (Unsolved pos k a b c cons) => do
       debug $ \ _ => "Add constraint m\{show ix} \{show sp} =?= \{show tm}"
       pure (Unsolved pos k a b c (MkMc (getFC tm) env sp tm :: cons))
     (Solved _ k tm) => error' "Meta \{show k} already solved (addConstraint :: Nil)"
-  top <- get
   mc <- readIORef top.metaCtx
   checkAutos ix mc.metas
   -- this loops too
@@ -401,6 +407,18 @@ unify : Env -> UnifyMode -> Val -> Val  -> M UnifyResult
 .boundNames : Context -> List String
 ctx.boundNames = map snd $ filter (\x => fst x == Bound) $ zip ctx.bds (map fst ctx.types)
 
+maybeCheck : M Unit -> M Unit
+maybeCheck action = do
+  top <- get
+  mc <- readIORef top.metaCtx
+  case mc.mcmode of
+    CheckAll => action
+    CheckFirst => do
+      modifyIORef top.metaCtx $ \ mc => MC mc.metas mc.next NoCheck
+      action
+      modifyIORef top.metaCtx $ \ mc => MC mc.metas mc.next CheckFirst
+    NoCheck => pure MkUnit
+
 solve env m sp t = do
   meta@(Unsolved metaFC ix ctx_ ty kind cons) <- lookupMeta m
     | _ => error (getFC t) "Meta \{show m} already solved! (solve :: Nil)"
@@ -429,7 +447,7 @@ solve env m sp t = do
     updateMeta m $ \case
       (Unsolved pos k _ _ _ cons) => pure $ Solved pos k soln
       (Solved _ k x) => error' "Meta \{show ix} already solved! (solve2 :: Nil)"
-    for cons $ \case
+    maybeCheck $ for_ cons $ \case
       MkMc fc env sp rhs => do
         val <- vappSpine soln sp
         debug $ \ _ => "discharge l=\{show $ length' env} \{(show val)} =?= \{(show rhs)}"
@@ -629,7 +647,7 @@ freshMeta ctx fc ty kind = do
   mc <- readIORef top.metaCtx
   debug $ \ _ => "fresh meta \{show mc.next} : \{show ty} (\{show kind})"
   let newmeta = Unsolved fc mc.next ctx ty kind Nil
-  writeIORef top.metaCtx $ MC (newmeta :: mc.metas) (1 + mc.next)
+  writeIORef top.metaCtx $ MC (newmeta :: mc.metas) (1 + mc.next) mc.mcmode
   -- infinite loop - keeps trying Ord a => Ord (a \x a)
   -- when (kind == AutoSolve) $ \ _ => ignore $ trySolveAuto newmeta
   pure $ applyBDs 0 (Meta fc mc.next) ctx.bds
diff --git a/port/Lib/TopContext.newt b/port/Lib/TopContext.newt
index 1dec1c7..8834135 100644
--- a/port/Lib/TopContext.newt
+++ b/port/Lib/TopContext.newt
@@ -30,7 +30,7 @@ instance Show TopContext where
 -- TODO need to get class dependencies working
 emptyTop : ∀ io. {{Monad io}} {{HasIO io}} -> io TopContext
 emptyTop = do
-  mcctx <- newIORef (MC Nil 0)
+  mcctx <- newIORef (MC Nil 0 CheckAll)
   errs <- newIORef $ the (List Error) Nil
   pure $ MkTop EmptyMap mcctx False errs Nil EmptyMap
 
diff --git a/port/Lib/Types.newt b/port/Lib/Types.newt
index 77f87b6..df7bb3a 100644
--- a/port/Lib/Types.newt
+++ b/port/Lib/Types.newt
@@ -307,10 +307,17 @@ data MConstraint = MkMc FC Env (SnocList Val) Val
 
 data MetaEntry = Unsolved FC Int Context Val MetaKind (List MConstraint) | Solved FC Int Val
 
+-- The purpose of this is to only check one level of constraints when trying an Auto solution
+-- The idea being we narrow it down to the likely solution, and let any consequent type error
+-- bubble up to the user, rather than have a type error wipe out all solutions.
+-- We also don't bother adding constraints if not in CheckAll mode
+data MetaMode = CheckAll | CheckFirst | NoCheck
+
 record MetaContext where
   constructor MC
   metas : List MetaEntry
   next : Int
+  mcmode : MetaMode
 
 data Def = Axiom | TCon (List QName) | DCon Int QName | Fn Tm | PrimTCon
          | PrimFn String (List String)
@@ -485,21 +492,6 @@ error fc msg = throwError $ E fc msg
 error' : ∀ a. String -> M a
 error' msg = throwError $ E emptyFC msg
 
--- freshMeta : Context -> FC -> Val -> MetaKind -> M Tm
--- freshMeta ctx fc ty kind = do
---   top <- get
---   mc <- readIORef top.metaCtx
---   debug $ \ _ => "fresh meta \{show mc.next} : \{show ty} (\{show kind})"
---   writeIORef top.metaCtx $ MC (Unsolved fc mc.next ctx ty kind Nil :: mc.metas) (1 + mc.next)
---   pure $ applyBDs 0 (Meta fc mc.next) ctx.bds
---   where
---     -- hope I got the right order here :)
---     applyBDs : Int -> Tm -> List BD -> Tm
---     applyBDs k t Nil = t
---     -- review the order here
---     applyBDs k t (Bound :: xs) = App emptyFC (applyBDs (1 + k) t xs) (Bnd emptyFC k)
---     applyBDs k t (Defined :: xs) = applyBDs (1 + k) t xs
-
 lookupMeta : Int -> M MetaEntry
 lookupMeta ix = do
   top <- get
@@ -511,8 +503,5 @@ lookupMeta ix = do
     go (meta@(Unsolved _ k ys _ _ _) :: xs) = if k == ix then pure meta else go xs
     go (meta@(Solved _ k x) :: xs) = if k == ix then pure meta else go xs
 
--- we need more of topcontext later - Maybe switch it up so we're not passing
--- around top
-
 mkCtx : FC -> Context
 mkCtx fc = MkCtx 0 Nil Nil Nil fc
diff --git a/src/Lib/CompileExp.idr b/src/Lib/CompileExp.idr
index f6e7334..b34b15c 100644
--- a/src/Lib/CompileExp.idr
+++ b/src/Lib/CompileExp.idr
@@ -116,7 +116,7 @@ compileTerm (Meta _ k) = pure $ CRef "meta$\{show k}" -- FIXME
 compileTerm (Lam _ nm _ _ t) = pure $ CLam nm !(compileTerm t)
 compileTerm tm@(App _ _ _) with (funArgs tm)
   _ | (Meta _ k, args) = do
-        -- this will be undefined, should only happen for use metas
+        info (getFC tm) "Compiling an unsolved meta \{showTm tm}"
         pure $ CApp (CRef "Meta\{show k}") [] Z
   _ | (t@(Ref fc nm _), args) = do
         args' <- traverse compileTerm args
diff --git a/src/Lib/Elab.idr b/src/Lib/Elab.idr
index 5611b4b..b45f824 100644
--- a/src/Lib/Elab.idr
+++ b/src/Lib/Elab.idr
@@ -7,6 +7,7 @@ import Data.Vect
 import Data.String
 import Data.IORef
 import Lib.Types
+import Lib.Util
 import Lib.Eval
 import Lib.TopContext
 import Lib.Syntax
@@ -131,6 +132,8 @@ findMatches ctx ty ((MkEntry _ name type def) :: xs) = do
       -- tm <- check (mkCtx fc) (RVar fc name) ty
       -- FIXME RVar should optionally have qualified names
       let (QN ns nm) = name
+      let (cod, tele) = splitTele type
+      modifyIORef top.metaCtx { mcmode := CheckFirst }
       tm <- check ctx (RVar fc nm) ty
       debug "Found \{pprint [] tm} for \{show ty}"
       writeIORef top.metaCtx mc
@@ -287,12 +290,14 @@ checkAutos ix (_ :: rest) = checkAutos ix rest
 export
 addConstraint : Env -> Nat -> SnocList Val -> Val -> M ()
 addConstraint env ix sp tm = do
+  top <- get
+  mc <- readIORef top.metaCtx
+  let (CheckAll) = mc.mcmode | _ => pure ()
   updateMeta ix $ \case
     (Unsolved pos k a b c cons) => do
       debug "Add constraint m\{show ix} \{show sp} =?= \{show tm}"
       pure (Unsolved pos k a b c (MkMc (getFC tm) env sp tm :: cons))
     (Solved _ k tm) => error' "Meta \{show k} already solved [addConstraint]"
-  top <- get
   mc <- readIORef top.metaCtx
   checkAutos ix mc.metas
   -- this loops too
@@ -375,6 +380,19 @@ unify : Env -> UnifyMode -> Val -> Val  -> M UnifyResult
 (.boundNames) : Context -> List String
 ctx.boundNames = map snd $ filter (\x => fst x == Bound) $ toList $ zip ctx.bds (map fst ctx.types)
 
+
+maybeCheck : M () -> M ()
+maybeCheck action = do
+  top <- get
+  mc <- readIORef top.metaCtx
+  case mc.mcmode of
+    CheckAll => action
+    CheckFirst => do
+      modifyIORef top.metaCtx $ { mcmode := NoCheck }
+      action
+      modifyIORef top.metaCtx $ { mcmode := CheckFirst }
+    NoCheck => pure ()
+
 solve env m sp t = do
   meta@(Unsolved metaFC ix ctx_ ty kind cons) <- lookupMeta m
     | _ => error (getFC t) "Meta \{show m} already solved! [solve]"
@@ -403,7 +421,8 @@ solve env m sp t = do
     updateMeta m $ \case
       (Unsolved pos k _ _ _ cons) => pure $ Solved pos k soln
       (Solved _ k x) => error' "Meta \{show ix} already solved! [solve2]"
-    for_ cons $ \case
+
+    maybeCheck $ for_ cons $ \case
       MkMc fc env sp rhs => do
         val <- vappSpine soln sp
         debug "discharge l=\{show $ length env} \{(show val)} =?= \{(show rhs)}"
diff --git a/src/Lib/TopContext.idr b/src/Lib/TopContext.idr
index df3aa5d..3bdb23a 100644
--- a/src/Lib/TopContext.idr
+++ b/src/Lib/TopContext.idr
@@ -29,7 +29,7 @@ Show TopContext where
 
 public export
 empty : HasIO m => m TopContext
-empty = pure $ MkTop empty !(newIORef (MC [] 0)) False !(newIORef []) [] empty
+empty = pure $ MkTop empty !(newIORef (MC [] 0 CheckAll)) False !(newIORef []) [] empty
 
 public export
 setDef : QName -> FC -> Tm -> Def -> M ()
diff --git a/src/Lib/Types.idr b/src/Lib/Types.idr
index 7d39b07..ac877c7 100644
--- a/src/Lib/Types.idr
+++ b/src/Lib/Types.idr
@@ -363,13 +363,15 @@ data MConstraint = MkMc FC Env (SnocList Val) Val
 public export
 data MetaEntry = Unsolved FC Nat Context Val MetaKind (List MConstraint) | Solved FC Nat Val
 
+public export
+data MetaMode = CheckAll | CheckFirst | NoCheck
 
 public export
 record MetaContext where
   constructor MC
   metas : List MetaEntry
   next : Nat
-
+  mcmode : MetaMode
 
 public export
 data Def = Axiom | TCon (List QName) | DCon Nat QName | Fn Tm | PrimTCon