feat: allow metas in intermediate tc answers

library/Init/Lean/TypeClass/Context.lean

@@ -250,66 +250,101 @@ partial def eInstantiate (ctx : Context) : Expr → Expr
 
 -- AlphaNormalization
 
-structure AlphaNormData : Type :=
+structure MetaNormData (α : Type) : Type :=
+(ctx        : α)
 (eRenameMap : RBMap Nat Nat (λ n₁ n₂ => n₁ < n₂) := mkRBMap _ _ _)
 (uRenameMap : RBMap Nat Nat (λ n₁ n₂ => n₁ < n₂) := mkRBMap _ _ _)
 
-partial def uAlphaNormalizeCore : Level → StateM AlphaNormData Level
+structure MetaNormFuncs (α : Type) : Type :=
+(uNewMeta   : Nat → StateM (MetaNormData α) Level)
+(uMkMeta    : Nat → StateM (MetaNormData α) Level)
+(eNewMeta   : Nat → StateM (MetaNormData α) Expr)
+(eMkMeta    : Nat → StateM (MetaNormData α) Expr)
+
+partial def uMetaNormalizeCore {α : Type} (fs : MetaNormFuncs α) : Level → StateM (MetaNormData α) Level
 | l =>
   if !l.hasMVar then pure l else
     match l with
     | Level.zero        => pure l
     | Level.param _     => pure l
-    | Level.succ l      => do l ← uAlphaNormalizeCore l;
+    | Level.succ l      => do l ← uMetaNormalizeCore l;
                              pure $ Level.succ l
-    | Level.max l₁ l₂   => do l₁ ← uAlphaNormalizeCore l₁;
-                             l₂ ← uAlphaNormalizeCore l₂;
+    | Level.max l₁ l₂   => do l₁ ← uMetaNormalizeCore l₁;
+                             l₂ ← uMetaNormalizeCore l₂;
                              pure $ Level.max l₁ l₂
-    | Level.imax l₁ l₂  => do l₁ ← uAlphaNormalizeCore l₁;
-                              l₂ ← uAlphaNormalizeCore l₂;
+    | Level.imax l₁ l₂  => do l₁ ← uMetaNormalizeCore l₁;
+                              l₂ ← uMetaNormalizeCore l₂;
                               pure $ Level.imax l₁ l₂
     | Level.mvar m      =>
       match uMetaIdx l with
       | none     => pure l
-      | some idx => do
+      | some idx =>  do
         lookupStatus ← get >>= λ ϕ => pure $ ϕ.uRenameMap.find idx;
         match lookupStatus with
-        | none => do
-          l ← get >>= λ ϕ => pure $ Level.mvar (mkNumName alphaMetaPrefix ϕ.uRenameMap.size);
-          modify $ λ ϕ => { uRenameMap := ϕ.uRenameMap.insert idx ϕ.uRenameMap.size .. ϕ };
-          pure l
-        | some alphaIdx => pure $ Level.mvar (mkNumName alphaMetaPrefix alphaIdx)
+        | none => fs.uNewMeta idx
+        | some idx => fs.uMkMeta idx
 
-partial def eAlphaNormalizeCore : Expr → StateM AlphaNormData Expr
+partial def eMetaNormalizeCore {α : Type} (fs : MetaNormFuncs α) : Expr → StateM (MetaNormData α) Expr
 | e =>
   if e.isConst then do
-    ls ← e.constLevels!.mapM uAlphaNormalizeCore;
+    ls ← e.constLevels!.mapM (uMetaNormalizeCore fs);
     pure $ Expr.updateConst! e ls
   else if e.isFVar then pure e
   else if !e.hasMVar then pure e
   else match e with
        | Expr.app f a => do
-         f ← eAlphaNormalizeCore f;
-         a ← eAlphaNormalizeCore a;
+         f ← eMetaNormalizeCore f;
+         a ← eMetaNormalizeCore a;
          pure $ Expr.app f a
        | Expr.forallE n i d b => do
-         d ← eAlphaNormalizeCore d;
-         b ← eAlphaNormalizeCore b;
+         d ← eMetaNormalizeCore d;
+         b ← eMetaNormalizeCore b;
          pure $ Expr.forallE n i d b
        | _ =>
          match eMetaIdx e with
          | none     => pure e
          | some idx => do
            lookupStatus ← get >>= λ ϕ => pure $ ϕ.eRenameMap.find idx;
            match lookupStatus with
-           | none => do
-             e ← get >>= λ ϕ => pure $ Expr.mvar (mkNumName alphaMetaPrefix ϕ.eRenameMap.size);
-             modify $ λ ϕ => { eRenameMap := ϕ.eRenameMap.insert idx ϕ.eRenameMap.size .. ϕ };
-             pure e
-           | some alphaIdx => pure $ Expr.mvar (mkNumName alphaMetaPrefix alphaIdx)
+           | none => fs.eNewMeta idx
+           | some idx => fs.eMkMeta idx
 
 def αNorm (e : Expr) : Expr :=
-(eAlphaNormalizeCore e).run' {}
+let fs : MetaNormFuncs Unit := {
+  uNewMeta := λ idx => do {
+    l ← get >>= λ ϕ => pure $ Level.mvar (mkNumName alphaMetaPrefix ϕ.uRenameMap.size);
+    modify $ λ ϕ => { uRenameMap := ϕ.uRenameMap.insert idx ϕ.uRenameMap.size .. ϕ };
+    pure l },
+  uMkMeta := λ idx => pure $ Level.mvar (mkNumName alphaMetaPrefix idx),
+  eNewMeta := λ idx => do {
+    e ← get >>= λ ϕ => pure $ Expr.mvar (mkNumName alphaMetaPrefix ϕ.eRenameMap.size);
+    modify $ λ ϕ => { eRenameMap := ϕ.eRenameMap.insert idx ϕ.eRenameMap.size .. ϕ };
+    pure e },
+  eMkMeta := λ idx => pure $ Expr.mvar (mkNumName alphaMetaPrefix idx)
+};
+(eMetaNormalizeCore fs e).run' { ctx := () }
+
+def internalize (oldCtx : Context) (val type : Expr) (newCtx : Context) : Expr × Expr × Context :=
+let fs : MetaNormFuncs (Context × Context) := {
+  uNewMeta := λ idx => do {
+    (oldCtx, newCtx) ← get >>= λ ϕ => pure ϕ.ctx;
+    (l, newCtx) ← pure $ StateT.run Context.uNewMeta newCtx;
+    match Context.uMetaIdx l with
+    | some newIdx => modify $ λ ϕ => { ctx := (oldCtx, newCtx), uRenameMap := ϕ.uRenameMap.insert idx newIdx, .. ϕ }
+    | none => panic "unreachable";
+    pure l },
+  uMkMeta := λ idx => pure $ Level.mvar (mkNumName metaPrefix idx),
+  eNewMeta := λ idx => do {
+    (oldCtx, newCtx) ← get >>= λ ϕ => pure ϕ.ctx;
+    (e, newCtx) ← pure $ StateT.run (Context.eNewMeta $ oldCtx.eTypes.get! idx) newCtx;
+    match Context.eMetaIdx e with
+    | some newIdx => modify $ λ ϕ => { ctx := (oldCtx, newCtx), eRenameMap := ϕ.eRenameMap.insert idx newIdx, .. ϕ }
+    | none => panic "unreachable";
+    pure e },
+  eMkMeta := λ idx => pure $ Expr.mvar (mkNumName metaPrefix idx)
+};
+match (do newType ← eMetaNormalizeCore fs type; newVal ← eMetaNormalizeCore fs val; pure (newVal, newType)).run { ctx := (oldCtx, newCtx) } with
+| ((newVal, newType), ϕ) => (newVal, newType, ϕ.ctx.2)
 
 end Context
 end TypeClass

library/Init/Lean/TypeClass/Synth.lean

@@ -29,6 +29,15 @@ instance TypedExpr.HasToString : HasToString TypedExpr :=
 instance TypedExpr.Inhabited : Inhabited TypedExpr :=
 ⟨⟨arbitrary _, arbitrary _⟩⟩
 
+structure Answer : Type :=
+(ctx : Context) (typedExpr : TypedExpr)
+
+instance Answer.HasToString : HasToString Answer :=
+⟨λ ⟨_, ⟨val, type⟩⟩ => "Answer(" ++ toString val ++ ", " ++ toString type ++ ")"⟩
+
+instance Answer.Inhabited : Inhabited TypedExpr :=
+⟨⟨arbitrary _, arbitrary _⟩⟩
+
 structure Node : Type :=
 (anormSubgoal : Expr)
 (ctx          : Context)
@@ -47,6 +56,10 @@ inductive Waiter : Type
 | consumerNode : ConsumerNode → Waiter
 | root         : Waiter
 
+def Waiter.isRoot : Waiter → Bool
+| Waiter.consumerNode _ => false
+| root => true
+
 -- TODO(dselsam): support local instances once elaborator is in Lean
 inductive Instance : Type
 | lDecl : LocalDecl → Instance
@@ -60,15 +73,15 @@ instance GeneratorNode.Inhabited : Inhabited GeneratorNode :=
 
 structure TableEntry : Type :=
 (waiters    : Array Waiter)
-(answers    : Array TypedExpr := #[])
+(answers    : Array Answer := #[])
 
 structure TCState : Type :=
 (env            : Environment)
 (finalAnswer    : Option TypedExpr                  := none)
 (mainMVar       : Expr                              := arbitrary _)
 (generatorStack : Stack GeneratorNode               := Stack.empty)
 (consumerStack  : Stack ConsumerNode                := Stack.empty)
-(resumeQueue    : Queue (ConsumerNode × TypedExpr)  := Queue.empty)
+(resumeQueue    : Queue (ConsumerNode × Answer)     := Queue.empty)
 (tableEntries   : PersistentHashMap Expr TableEntry := PersistentHashMap.empty)
 
 abbrev TCMethod : Type → Type := EStateM String TCState
@@ -91,17 +104,18 @@ def quickIsClass (env : Environment) : Expr → Option (Option Name)
 | _            => some none
 
 def newSubgoal (waiter : Waiter) (ctx : Context) (anormSubgoal mvar : Expr) : TCMethod Unit :=
-do let mvarType := ctx.eInfer mvar;
+do let mvarType := ctx.eInstantiate (ctx.eInfer mvar);
    isClassStatus ← get >>= λ ϕ => pure $ quickIsClass ϕ.env mvarType;
    match isClassStatus with
    | none      => throw $ "quickIsClass not sufficient to show `" ++ toString mvarType ++ "` is a class"
    | some none => throw $ "found non-class goal `" ++ toString mvarType ++ "`"
    | some (some n) => do
+     let ⟨newVal, newType, newCtx⟩ := Context.internalize ctx mvar mvarType {};
      gNode ← get >>= λ ϕ => pure {
        GeneratorNode .
-       ctx                := ctx,
+       ctx                := newCtx,
        anormSubgoal       := anormSubgoal,
-       futureAnswer       := ⟨mvar, mvarType⟩,
+       futureAnswer       := ⟨newVal, newType⟩,
        remainingInstances := (getClassInstances ϕ.env n).map Instance.const
      };
      let tableEntry : TableEntry := { waiters := #[waiter] };
@@ -152,23 +166,28 @@ do ((cNode, answer), resumeQueue) ← get >>= λ ϕ =>
    match cNode.remainingSubgoals with
    | []           => throw "resume found no remaining subgoals"
    | (mvar::rest) => do
+     let newCtx : Context := cNode.ctx;
+     let ⟨newVal, newType, newCtx⟩ : Expr × Expr × Context := Context.internalize answer.ctx answer.typedExpr.val answer.typedExpr.type newCtx;
      result : Option (Context × List Expr) ←
-       tryResolve cNode.ctx ⟨mvar, cNode.ctx.eInfer mvar⟩ answer;
+       tryResolve newCtx ⟨mvar, newCtx.eInfer mvar⟩ ⟨newVal, newType⟩;
      modify $ λ ϕ => { resumeQueue := resumeQueue, .. ϕ };
      match result with
      | none => pure ()
-     | some (ctx, newMVars) => newConsumerNode cNode.toNode ctx (newMVars ++ rest)
+     | some (newCtx, newMVars) => newConsumerNode cNode.toNode newCtx (newMVars ++ rest)
 
-def wakeUp (answer : TypedExpr) : Waiter → TCMethod Unit
-| Waiter.root               => modify $ λ ϕ => { finalAnswer := some answer .. ϕ }
+def wakeUp (answer : Answer) : Waiter → TCMethod Unit
+| Waiter.root               => modify $ λ ϕ => { finalAnswer := some answer.typedExpr .. ϕ }
 | Waiter.consumerNode cNode => modify $ λ ϕ => { resumeQueue := ϕ.resumeQueue.enqueue (cNode, answer), .. ϕ }
 
-def newAnswer (anormSubgoal : Expr) (answer : TypedExpr) : TCMethod Unit :=
+def newAnswer (anormSubgoal : Expr) (answer : Answer) : TCMethod Unit :=
 do lookupStatus ← get >>= λ ϕ => pure $ ϕ.tableEntries.find anormSubgoal;
    match lookupStatus with
    | none       => throw $ "[newAnswer]: " ++ toString anormSubgoal ++ " not found in table!"
    | some entry => do
-     if entry.answers.any (λ answer₁ => answer₁.type == answer.type) then pure() else do
+     if entry.answers.any (λ answer₁ => Context.αNorm (answer₁.typedExpr.type) == Context.αNorm (answer.typedExpr.type)) then pure ()
+     else if entry.waiters.any Waiter.isRoot
+             && (Context.eHasTmpMVar answer.typedExpr.type || Context.eHasTmpMVar answer.typedExpr.val) then pure()
+     else do
        let newEntry : TableEntry := { answers := entry.answers.push answer .. entry };
        modify $ λ ϕ => { tableEntries := ϕ.tableEntries.insert anormSubgoal newEntry .. ϕ };
        entry.waiters.forM (wakeUp answer)
@@ -178,12 +197,12 @@ do cNode ← get >>= λ ϕ => pure ϕ.consumerStack.peek!;
    modify $ λ ϕ => { consumerStack := ϕ.consumerStack.pop .. ϕ };
    match cNode.remainingSubgoals with
    | [] => do
-     let answer : TypedExpr := {
-       val := cNode.ctx.eInstantiate cNode.futureAnswer.val,
-       type := cNode.ctx.eInstantiate cNode.futureAnswer.type
-     };
-     when (Context.eHasTmpMVar answer.val || Context.eHasTmpMVar answer.type) $
-       throw $ "answer " ++ toString answer ++ " not fully instantiated";
+     let answer : Answer := {
+       ctx := cNode.ctx,
+       typedExpr := {
+         val  := cNode.ctx.eInstantiate cNode.futureAnswer.val,
+         type := cNode.ctx.eInstantiate cNode.futureAnswer.type
+     }};
      newAnswer cNode.anormSubgoal answer
 
    | mvar::rest => do

tests/lean/run/typeclass_metas_internal_goals.lean

@@ -0,0 +1,70 @@
+namespace T1
+
+class Foo (α : Type) : Type := (u : Unit := ())
+class Bar (α : Type) : Type := (u : Unit := ())
+class Top : Type := (u : Unit := ())
+
+instance FooAll (α : Type) : Foo α := {u:=()}
+instance BarNat : Bar Nat := {u:=()}
+
+instance FooBarToTop (α : Type) [Foo α] [Bar α] : Top := {u:=()}
+
+#synth Top
+
+end T1
+
+namespace T2
+
+class Foo (α β : Type) : Type := (u : Unit := ())
+class Bar (α β : Type) : Type := (u : Unit := ())
+class Top : Type := (u : Unit := ())
+
+instance FooNatA (β : Type) : Foo Nat β := {u:=()}
+instance BarANat (α : Type) : Bar α Nat := {u:=()}
+
+instance FooBarToTop (α β : Type) [Foo α β] [Bar α β] : Top := {u:=()}
+
+#synth Top
+
+end T2
+
+namespace T3
+
+class Base (α : Type) := (u:Unit)
+class Depends (α : Type) [Base α] := (u:Unit)
+class Top := (u:Unit)
+
+instance AllBase {α : Type} : Base α := {u:=()}
+instance DependsNotConstrainingImplicit {α : Type} /- [Base α] -/ {_:Base α} : Depends α := {u:=()}
+
+instance BaseAsImplicit₁ {α : Type} {_:Base α} [Depends α] : Top := {u:=()}
+instance BaseAsInstImplicit {α : Type} [Base α] [Depends α] : Top := {u:=()}
+instance BaseAsImplicit₂ {α : Type} {_:Base α} [Depends α] : Top := {u:=()}
+
+axiom K : Type
+instance BaseK : Base K := {u:=()}
+
+#synth Top
+
+end T3
+
+namespace T4
+
+class Foo (α β γ : Type) := (u:Unit)
+class Bar (α β γ : Type) := (u:Unit)
+class Top := (u:Unit)
+instance FooBarToTop (α β γ : Type) [Foo α β γ] [Bar α β γ] : Top := {u:=()}
+
+instance Foo₁ (β γ : Type) : Foo Unit β γ := {u:=()}
+instance Foo₂ (α γ : Type) : Foo α Unit γ := {u:=()}
+instance Foo₃ (α β : Type) : Foo α β Unit := {u:=()}
+
+instance Foo₁₂ (γ : Type) : Foo Unit Nat γ := {u:=()}
+instance Foo₂₃ (α : Type) : Foo α Unit Nat := {u:=()}
+instance Foo₃₁ (β : Type) : Foo Nat β Unit := {u:=()}
+
+instance Bar0 : Bar Unit Int (List Int) := {u:=()}
+
+#synth Top
+
+end T4