From 7359f950886997e16fcba470f6fb23ea6aa64365 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Sun, 7 Aug 2022 18:04:38 -0700 Subject: [PATCH] refactor: treat `casesOn` and `matcher` applications uniformly --- src/Lean/Compiler/LCNF.lean | 58 +++++++++++++++---------------------- src/Lean/Compiler/Util.lean | 43 ++++++++++++++++++++------- 2 files changed, 55 insertions(+), 46 deletions(-) diff --git a/src/Lean/Compiler/LCNF.lean b/src/Lean/Compiler/LCNF.lean index db5d9d7d70..c44128fffd 100644 --- a/src/Lean/Compiler/LCNF.lean +++ b/src/Lean/Compiler/LCNF.lean @@ -78,20 +78,6 @@ where else k - /-- - Visit a `matcher`/`casesOn` alternative. - -/ - visitAlt (e : Expr) (numParams : Nat) : M Expr := do - withNewRootScope do - let mut (as, e) ← Compiler.visitBoundedLambda e numParams - if as.size < numParams then - e ← etaExpandN e (numParams - as.size) - let (as', e') ← Compiler.visitLambda e - as := as ++ as' - e := e' - e ← mkLetUsingScope (← visit e) - mkLambda as e - /-- If `args.size == arity`, then just return `app`. Otherwise return @@ -125,25 +111,29 @@ where else mkOverApplication (mkAppN f args[:arity]) args arity - visitMatcher (matcherInfo : Meta.MatcherInfo) (e : Expr) : M Expr := - etaIfUnderApplied e matcherInfo.arity do - let mut args := e.getAppArgs - for i in matcherInfo.getDiscrRange do - args ← args.modifyM i visitChild - for i in matcherInfo.getAltRange, altIdx in [: matcherInfo.numAlts] do - let alt ← visitAlt args[i]! matcherInfo.altNumParams[altIdx]! - args := args.set! i alt - mkAppWithArity e.getAppFn args matcherInfo.arity + /-- + Visit a `matcher`/`casesOn` alternative. + -/ + visitAlt (e : Expr) (numParams : Nat) : M Expr := do + withNewRootScope do + let mut (as, e) ← Compiler.visitBoundedLambda e numParams + if as.size < numParams then + e ← etaExpandN e (numParams - as.size) + let (as', e') ← Compiler.visitLambda e + as := as ++ as' + e := e' + e ← mkLetUsingScope (← visit e) + mkLambda as e - visitCasesOn (casesOnInfo : CasesOnInfo) (e : Expr) : M Expr := - etaIfUnderApplied e casesOnInfo.arity do + visitCases (casesInfo : CasesInfo) (e : Expr) : M Expr := + etaIfUnderApplied e casesInfo.arity do let mut args := e.getAppArgs - args ← args.modifyM casesOnInfo.majorPos visitChild - for i in casesOnInfo.minorsRange, ctor in casesOnInfo.ctors do - let .ctorInfo ctorVal ← getConstInfo ctor | unreachable! - let alt ← visitAlt args[i]! ctorVal.numFields + for i in casesInfo.discrsRange do + args ← args.modifyM i visitChild + for i in casesInfo.altsRange, numParams in casesInfo.altNumParams do + let alt ← visitAlt args[i]! numParams args := args.set! i alt - mkAppWithArity e.getAppFn args casesOnInfo.arity + mkAppWithArity e.getAppFn args casesInfo.arity visitCtor (arity : Nat) (e : Expr) : M Expr := etaIfUnderApplied e arity do @@ -239,9 +229,7 @@ where visitApp (e : Expr) : M Expr := do if let .const declName _ := e.getAppFn then - if let some matcherInfo ← Meta.getMatcherInfo? declName then - visitMatcher matcherInfo e - else if declName == ``Quot.lift then + if declName == ``Quot.lift then visitQuotLift e else if declName == ``Quot.mk then visitCtor 3 e @@ -251,8 +239,8 @@ where visitAndRec e else if declName == ``False.rec || declName == ``Empty.rec || declName == ``False.casesOn || declName == ``Empty.casesOn then visitFalseRec e - else if let some casesOnInfo ← getCasesOnInfo? declName then - visitCasesOn casesOnInfo e + else if let some casesInfo ← getCasesInfo? declName then + visitCases casesInfo e else if let some arity ← getCtorArity? declName then visitCtor arity e else if isNoConfusion (← getEnv) declName then diff --git a/src/Lean/Compiler/Util.lean b/src/Lean/Compiler/Util.lean index 5cfbe92dd7..d4ea84a8b7 100644 --- a/src/Lean/Compiler/Util.lean +++ b/src/Lean/Compiler/Util.lean @@ -4,6 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ import Lean.Compiler.CompilerM +import Lean.Meta.Match.MatcherInfo namespace Lean.Compiler @@ -41,23 +42,43 @@ def mkLcUnreachable (type : Expr) : CompilerM Expr := do let u ← Meta.getLevel type return .app (.const ``lcUnreachable [u]) type -structure CasesOnInfo where - arity : Nat - majorPos : Nat - minorsRange : Std.Range - ctors : List Name +/-- +Store information about `matcher` and `casesOn` declarations. -def getCasesOnInductiveVal? (declName : Name) : CoreM (Option InductiveVal) := do +We treat them uniformly in the code generator. +-/ +structure CasesInfo where + arity : Nat + discrsRange : Std.Range + altsRange : Std.Range + altNumParams : Array Nat + +private def getCasesOnInductiveVal? (declName : Name) : CoreM (Option InductiveVal) := do unless isCasesOnRecursor (← getEnv) declName do return none let .inductInfo val ← getConstInfo declName.getPrefix | return none return some val -def getCasesOnInfo? (declName : Name) : CoreM (Option CasesOnInfo) := do +private def getCasesOnInfo? (declName : Name) : CoreM (Option CasesInfo) := do let some val ← getCasesOnInductiveVal? declName | return none - let arity := val.numIndices + val.numParams + 1 /- motive -/ + 1 /- major -/ + val.numCtors - let majorPos := val.numIndices + val.numParams + 1 /- motive -/ - let minorsRange := { start := majorPos + 1, stop := arity } - return some { arity, majorPos, minorsRange, ctors := val.ctors } + let arity := val.numIndices + val.numParams + 1 /- motive -/ + 1 /- major -/ + val.numCtors + let majorPos := val.numIndices + val.numParams + 1 /- motive -/ + let discrsRange := { start := majorPos, stop := majorPos + 1 } + let altsRange := { start := majorPos + 1, stop := arity } + let altNumParams ← val.ctors.toArray.mapM fun ctor => do + let .ctorInfo ctorVal ← getConstInfo ctor | unreachable! + return ctorVal.numFields + return some { arity, discrsRange, altsRange, altNumParams } + +def getCasesInfo? (declName : Name) : CoreM (Option CasesInfo) := do + if let some matcherInfo ← Meta.getMatcherInfo? declName then + return some { + arity := matcherInfo.arity + discrsRange := matcherInfo.getDiscrRange + altsRange := matcherInfo.getAltRange + altNumParams := matcherInfo.altNumParams + } + else + getCasesOnInfo? declName def getCtorArity? (declName : Name) : CoreM (Option Nat) := do let .ctorInfo val ← getConstInfo declName | return none