feat: add isDefEqQuick

library/Init/Lean/Expr.lean

@@ -25,6 +25,13 @@ def Literal.hash : Literal → USize
 
 instance Literal.hashable : Hashable Literal := ⟨Literal.hash⟩
 
+def Literal.beq : Literal → Literal → Bool
+| Literal.natVal v₁, Literal.natVal v₂ => v₁ == v₂
+| Literal.strVal v₁, Literal.strVal v₂ => v₁ == v₂
+| _,                 _                 => false
+
+instance Literal.hasBeq : HasBeq Literal := ⟨Literal.beq⟩
+
 inductive BinderInfo
 | default | implicit | strictImplicit | instImplicit | auxDecl
 

library/Init/Lean/Meta/Basic.lean

@@ -66,6 +66,15 @@ structure Config :=
 (foApprox           : Bool    := false)
 (ctxApprox          : Bool    := false)
 (quasiPatternApprox : Bool    := false)
+/-
+  When the following flag is set,
+  `isDefEq` throws the exeption `Exeption.isDefEqStuck`
+  whenever it encounters a constraint `?m ... =?= t` where
+  `?m` is read only.
+  This feature is useful for type class resolution where
+  we may want to notify the caller that the TC problem may be solveable
+  later after it assigns `?m`. -/
+(isDefEqStuckEx     : Bool    := false)
 (debug              : Bool    := false)
 (transparency       : TransparencyMode := TransparencyMode.default)
 
@@ -183,6 +192,12 @@ adaptReader
     { config := { transparency := mode, .. ctx.config }, .. ctx })
   x
 
+def isSyntheticExprMVar (mvarId : Name) : MetaM Bool :=
+do mctx ← getMCtx;
+   match mctx.findDecl mvarId with
+   | some d => pure $ d.synthetic
+   | _      => throwEx $ Exception.unknownExprMVar mvarId
+
 def isReadOnlyOrSyntheticExprMVar (mvarId : Name) : MetaM Bool :=
 do mctx ← getMCtx;
    match mctx.findDecl mvarId with

library/Init/Lean/Meta/Exception.lean

@@ -28,6 +28,7 @@ inductive Exception
 | invalidProjection    (structName : Name) (idx : Nat) (s : Expr) (ctx : ExceptionContext)
 | revertFailure        (toRevert : Array Expr) (decl : LocalDecl) (ctx : ExceptionContext)
 | readOnlyMVar         (mvarId : Name) (ctx : ExceptionContext)
+| isDefEqStuck         (t s : Expr) (ctx : ExceptionContext)
 | letTypeMismatch      (fvarId : Name) (ctx : ExceptionContext)
 | appTypeMismatch      (f a : Expr) (ctx : ExceptionContext)
 | bug                  (b : Bug) (ctx : ExceptionContext)
@@ -49,6 +50,7 @@ def toStr : Exception → String
 | invalidProjection _ _ _ _     => "invalid projection"
 | revertFailure _ _ _           => "revert failure"
 | readOnlyMVar _ _              => "try to assign read only metavariable"
+| isDefEqStuck _ _ _            => "isDefEq is stuck"
 | letTypeMismatch _ _           => "type mismatch at let-expression"
 | appTypeMismatch _ _ _         => "application type mismatch"
 | bug _ _                       => "bug"

library/Init/Lean/Meta/ExprDefEq.lean

@@ -676,9 +676,10 @@ do arg ← if arg.getAppFn.hasExprMVar then instantiateMVars arg else pure arg;
     (whnf              : Expr → MetaM Expr)
     (isDefEq           : Expr → Expr → MetaM Bool)
     (synthesizePending : Expr → MetaM Bool)
-    (mvar : Expr) (args : Array Expr) (v : Expr) : MetaM Bool :=
-do mvarDecl ← getMVarDecl mvar.mvarId!;
-   processAssignmentAux whnf isDefEq synthesizePending mvar mvarDecl v 0 args
+    (mvarApp : Expr) (v : Expr) : MetaM Bool :=
+do let mvar := mvarApp.getAppFn;
+   mvarDecl ← getMVarDecl mvar.mvarId!;
+   processAssignmentAux whnf isDefEq synthesizePending mvar mvarDecl v 0 mvarApp.getAppArgs
 
 @[specialize] private def unfold {α}
     (whnf              : Expr → MetaM Expr)
@@ -716,7 +717,7 @@ do let isDefEqL (t s : Expr) : MetaM LBool := toLBoolM $ isDefEq t s;
      trace! `Meta.isDefEq.delta.unfoldLeftRight fn;
      toLBoolM $ isDefEq t s
    };
-   let isListLevelDefEqL (us vs : List Level) : MetaM LBool := toLBoolM $ isListLevelDefEq us vs;
+   let isListLevelDefEqL (us vs : List Level) : MetaM LBool := toLBoolM $ isListLevelDefEqAux us vs;
    let unfold (e failK successK) : MetaM LBool := unfoldDefinitionAux whnf (inferTypeAux whnf) isDefEq synthesizePending e failK successK;
    let tryHeuristic : MetaM Bool :=
      /- Try to solve `f a₁ ... aₙ =?= f b₁ ... bₙ` by solving `a₁ =?= b₁, ..., aₙ =?= bₙ` -/
@@ -726,7 +727,7 @@ do let isDefEqL (t s : Expr) : MetaM LBool := toLBoolM $ isDefEq t s;
        try $
          isDefEqArgs whnf isDefEq synthesizePending tFn t.getAppArgs s.getAppArgs
          <&&>
-         isListLevelDefEq tFn.constLevels! sFn.constLevels!;
+         isListLevelDefEqAux tFn.constLevels! sFn.constLevels!;
    tInfo? ← isDeltaCandidate t.getAppFn;
    sInfo? ← isDeltaCandidate s.getAppFn;
    match tInfo?, sInfo? with
@@ -791,5 +792,80 @@ do let isDefEqL (t s : Expr) : MetaM LBool := toLBoolM $ isDefEq t s;
            else
              kernelLikeUnfolding ())
 
+private def isAssigned : Expr → MetaM Bool
+| Expr.mvar mvarId _ => isExprMVarAssigned mvarId
+| _                  => pure false
+
+private def isSynthetic : Expr → MetaM Bool
+| Expr.mvar mvarId _ => isSyntheticExprMVar mvarId
+| _                  => pure false
+
+private def isAssignable : Expr → MetaM Bool
+| Expr.mvar mvarId _ => do b ← isReadOnlyOrSyntheticExprMVar mvarId; pure (!b)
+| _                  => pure false
+
+private def etaEq (t s : Expr) : Bool :=
+match t.etaExpanded? with
+| some t => t == s
+| none   => false
+
+private def isLetFVar (fvarId : Name) : MetaM Bool :=
+do decl ← getLocalDecl fvarId;
+   pure decl.isLet
+
+@[specialize] private partial def isDefEqQuick
+    (whnf              : Expr → MetaM Expr)
+    (isDefEq           : Expr → Expr → MetaM Bool)
+    (synthesizePending : Expr → MetaM Bool)
+    : Expr → Expr → MetaM LBool
+| Expr.lit  l₁ _,           Expr.lit l₂ _            => pure (l₁ == l₂).toLBool
+| Expr.sort u _,            Expr.sort v _            => toLBoolM $ isLevelDefEqAux u v
+| t@(Expr.lam _ _ _ _),     s@(Expr.lam _ _ _ _)     => toLBoolM $ isDefEqBinding whnf isDefEq t s
+| t@(Expr.forallE _ _ _ _), s@(Expr.forallE _ _ _ _) => toLBoolM $ isDefEqBinding whnf isDefEq t s
+| Expr.mdata _ t _,         s                        => isDefEqQuick t s
+| t,                        Expr.mdata _ s _         => isDefEqQuick t s
+| Expr.fvar fvarId₁ _,      Expr.fvar fvarId₂ _      =>
+  condM (isLetFVar fvarId₁ <||> isLetFVar fvarId₂)
+    (pure LBool.undef)
+    (pure (fvarId₁ == fvarId₂).toLBool)
+| t, s =>
+  cond (t == s) (pure LBool.true) $
+  cond (etaEq t s || etaEq s t) (pure LBool.true) $  -- t =?= (fun xs => t xs)
+  let tFn := t.getAppFn;
+  let sFn := s.getAppFn;
+  cond (!tFn.isMVar && !sFn.isMVar) (pure LBool.undef) $
+  condM (isAssigned tFn) (do t ← instantiateMVars t; isDefEqQuick t s) $
+  condM (isAssigned sFn) (do s ← instantiateMVars s; isDefEqQuick t s) $
+  condM (isSynthetic tFn <&&> synthesizePending tFn) (do t ← instantiateMVars t; isDefEqQuick t s) $
+  condM (isSynthetic sFn <&&> synthesizePending sFn) (do s ← instantiateMVars s; isDefEqQuick t s) $ do
+  tAssign? ← isAssignable tFn;
+  sAssign? ← isAssignable sFn;
+  let assign (t s : Expr) : MetaM LBool := toLBoolM $ processAssignment whnf isDefEq synthesizePending t s;
+  cond (tAssign? && !sAssign?)  (assign t s) $
+  cond (!tAssign? && sAssign?)  (assign s t) $
+  cond (!tAssign? && !sAssign?)
+    (if tFn.isMVar || sFn.isMVar then do
+       ctx ← read;
+       if ctx.config.isDefEqStuckEx then throwEx $ Exception.isDefEqStuck t s
+       else pure LBool.false
+     else pure LBool.undef) $ do
+  -- Both `t` and `s` are terms of the form `?m ...`
+  tMVarDecl ← getMVarDecl tFn.mvarId!;
+  sMVarDecl ← getMVarDecl sFn.mvarId!;
+  cond (!sMVarDecl.lctx.isSubPrefixOf tMVarDecl.lctx) (assign s t) $
+  /-
+    Local context for `s` is a sub prefix of the local context for `t`.
+
+    Remark:
+    It is easier to solve the assignment
+        ?m2 := ?m1 a_1 ... a_n
+    than
+        ?m1 a_1 ... a_n := ?m2
+    Reason: the first one has a precise solution. For example,
+    consider the constraint `?m1 ?m =?= ?m2` -/
+  cond (!t.isApp && s.isApp) (assign t s) $
+  cond (!s.isApp && t.isApp && tMVarDecl.lctx.isSubPrefixOf sMVarDecl.lctx) (assign s t) $
+  assign t s
+
 end Meta
 end Lean

library/Init/Lean/Meta/LevelDefEq.lean

@@ -53,7 +53,7 @@ match u with
 | _ =>
   pure LevelConstraintKind.other
 
-private partial def isLevelDefEqAux : Level → Level → MetaM Bool
+partial def isLevelDefEqAux : Level → Level → MetaM Bool
 | Level.succ lhs _, Level.succ rhs _ => isLevelDefEqAux lhs rhs
 | lhs, rhs =>
   if lhs == rhs then
@@ -89,6 +89,11 @@ private partial def isLevelDefEqAux : Level → Level → MetaM Bool
               | _,         _         => do postponeIsLevelDefEq lhs rhs; pure true
             else do postponeIsLevelDefEq lhs rhs; pure true
 
+def isListLevelDefEqAux : List Level → List Level → MetaM Bool
+| [],    []    => pure true
+| u::us, v::vs => isLevelDefEqAux u v <&&> isListLevelDefEqAux us vs
+| _,     _     => pure false
+
 private def getNumPostponed : MetaM Nat :=
 do s ← get;
    pure s.postponed.size
@@ -167,10 +172,5 @@ try $ do
   if !r then pure false
   else processPostponed false
 
-def isListLevelDefEq : List Level → List Level → MetaM Bool
-| [],    []    => pure true
-| u::us, v::vs => isLevelDefEq u v <&&> isListLevelDefEq us vs
-| _,     _     => pure false
-
 end Meta
 end Lean