diff --git a/src/Lean/Compiler/LCNF/ToLCNF.lean b/src/Lean/Compiler/LCNF/ToLCNF.lean index 84a7b29fca..60077c8e3b 100644 --- a/src/Lean/Compiler/LCNF/ToLCNF.lean +++ b/src/Lean/Compiler/LCNF/ToLCNF.lean @@ -8,6 +8,7 @@ import Lean.ProjFns import Lean.Meta.CtorRecognizer import Lean.Compiler.BorrowedAnnotation import Lean.Compiler.CSimpAttr +import Lean.Compiler.ImplementedByAttr import Lean.Compiler.LCNF.Types import Lean.Compiler.LCNF.Bind import Lean.Compiler.LCNF.InferType @@ -567,6 +568,33 @@ where let result := .fvar auxDecl.fvarId mkOverApplication result args casesInfo.arity + visitCasesImplementedBy (casesInfo : CasesInfo) (f : Expr) (args : Array Expr) : M Arg := do + let mut args := args + let discr := args[casesInfo.discrPos]! + if discr matches .fvar _ then + let typeName := casesInfo.declName.getPrefix + let .inductInfo indVal ← getConstInfo typeName | unreachable! + args ← args.mapIdxM fun i arg => do + unless casesInfo.altsRange.start <= i && i < casesInfo.altsRange.stop do return arg + let altIdx := i - casesInfo.altsRange.start + let numParams := casesInfo.altNumParams[altIdx]! + let ctorName := indVal.ctors[altIdx]! + + -- We simplify `casesOn` arguments that simply reconstruct the discriminant and replace + -- them with the actual discriminant. This is required for hash consing to work correctly, + -- and should eventually be fixed by changing the elaborated term to use the original + -- variable. + Meta.MetaM.run' <| Meta.lambdaBoundedTelescope arg numParams fun paramExprs body => do + let fn := body.getAppFn + let args := body.getAppArgs + let args := args.map fun arg => + if arg.getAppFn.constName? == some ctorName && arg.getAppArgs == paramExprs then + discr + else + arg + Meta.mkLambdaFVars paramExprs (mkAppN fn args) + visitAppDefaultConst f args + visitCtor (arity : Nat) (e : Expr) : M Arg := etaIfUnderApplied e arity do visitAppDefaultConst e.getAppFn e.getAppArgs @@ -671,7 +699,7 @@ where visitApp (e : Expr) : M Arg := do if let some (args, n, t, v, b) := e.letFunAppArgs? then visitCore <| mkAppN (.letE n t v b (nonDep := true)) args - else if let .const declName _ := CSimp.replaceConstants (← getEnv) e.getAppFn then + else if let .const declName us := CSimp.replaceConstants (← getEnv) e.getAppFn then if declName == ``Quot.lift then visitQuotLift e else if declName == ``Quot.mk then @@ -687,7 +715,10 @@ where else if declName == ``False.rec || declName == ``Empty.rec || declName == ``False.casesOn || declName == ``Empty.casesOn then visitFalseRec e else if let some casesInfo ← getCasesInfo? declName then - visitCases casesInfo e + if (getImplementedBy? (← getEnv) declName).isSome then + e.withApp (visitCasesImplementedBy casesInfo) + else + visitCases casesInfo e else if let some arity ← getCtorArity? declName then visitCtor arity e else if isNoConfusion (← getEnv) declName then