Lean4
/-- `e.expand` eta-expands all expressions that have as head a constant `n` in `reorder`.
They are expanded until they are applied to one more argument than the maximum in `reorder.find n`.
It also expands all kernel projections that have as head a constant `n` in `reorder`. -/
def expand (e : Expr) : MetaM Expr := do
let env ← getEnv
let reorderFn : Name → List (List ℕ) := fun nm ↦ (reorderAttr.find? env nm |>.getD [])
let e₂ ←
Lean.Meta.transform (input := e) (skipConstInApp := true) (post := fun e => return .done e) fun e ↦
e.withApp fun f args ↦ do
match f with
| .proj n i s =>
let some info := getStructureInfo? (← getEnv) n |
return .continue
let some projName := info.getProjFn? i |
unreachable!
-- if `projName` is explicitly tagged with `@[to_additive]`,
-- replace `f` with the application `projName s` and then visit `projName s args` again.
if findTranslation? env projName |>.isNone then
return .continue
return
.visit <|
(← whnfD (← inferType s)).withApp fun sf sargs ↦
mkAppN (mkApp (mkAppN (.const projName sf.constLevels!) sargs) s) args
| .const c _ =>
let reorder := reorderFn c
if reorder.isEmpty then
-- no need to expand if nothing needs reordering
return .continue
let needed_n := reorder.flatten.foldr Nat.max 0 + 1
if needed_n ≤ args.size then
return .continue
else
-- in this case, we need to reorder arguments that are not yet
-- applied, so first η-expand the function.
let e' ← etaExpandN (needed_n - args.size) e
trace[to_additive_detail]"expanded {e } to {e'}"
return .continue e'
| _ =>
return .continue
if e != e₂ then
trace[to_additive_detail]"expand:\nBefore: {e }\nAfter: {e₂}"
return e₂
