Update test

needed to float out of ite

src/Init/Conv.lean

@@ -325,4 +325,31 @@ syntax (name := conv) "conv" (" at " ident)? (" in " (occs)? term)? " => " convS
 /-- `norm_cast` tactic in `conv` mode. -/
 syntax (name := normCast) "norm_cast" : conv
 
+/--
+Lifts out `match` expressions, or, equivalently, pushes function applications into the
+branches of a match. For example it can rewrite
+```
+f (match o with | some x => x + 1 | none => 0)
+```
+to
+```
+match o with | some x => f (x + 1) | none => f 0
+```
+
+For the purposes of this tactic, `if-then-else` expressions are treated like `match` expressions.
+
+It can only lift matches with a non-dependent motive, no extra arguments and when the context
+(here `fun x => f x`) is type-correct.
+
+It lifts the first eligible match it finds to the top. To lift less far (e.g. into the
+left-hand side of an equality) focus on the desired position first (e.g. using `lhs` and `rhs`).
+
+Also see the `liftMatch` simproc for use in the simplifier.
+-/
+syntax (name := liftMatch) "lift_match" : conv
+
+@[inherit_doc liftMatch]
+macro (name := _root_.Lean.Parser.Tactic.liftMatch) "lift_match" : tactic =>
+  `(tactic| conv => lift_match)
+
 end Lean.Parser.Tactic.Conv

src/Lean/Meta/Match.lean

@@ -9,6 +9,7 @@ import Lean.Meta.Match.Match
 import Lean.Meta.Match.CaseValues
 import Lean.Meta.Match.CaseArraySizes
 import Lean.Meta.Match.MatchEqs
+import Lean.Meta.Match.Lift
 
 namespace Lean
 

src/Lean/Meta/Match/Lift.lean

@@ -0,0 +1,99 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+
+prelude
+import Lean.Meta.Match.MatcherApp.Basic
+import Lean.Meta.Match.MatchEqsExt
+import Lean.Elab.Tactic.Meta
+
+open Lean Meta Elab Term
+
+namespace Lean.Meta
+
+def deriveMatchLift (name : Name) : MetaM Unit := do
+  mapError (f := (m!"Cannot construct match lifter:{indentD ·}")) do
+    let some info ← getMatcherInfo? name | throwError "getMatcherInfo? failed"
+    -- Fail early if splitter cannot be generated
+    try
+      discard <| Match.getEquationsFor name
+    catch _ =>
+      throwError "Could not construct splitter for {name}"
+
+    let cinfo ← getConstInfo name
+    let (u, v, us, us', levelParams) := if let some upos := info.uElimPos? then
+      let u := mkLevelParam `u
+      let v := mkLevelParam `v
+      let lps := (List.range cinfo.levelParams.length).map (Name.mkSimple s!"u_{·+1}")
+      let us := lps.map mkLevelParam
+      let us := us.set upos u
+      let us' := us.set upos v
+      let lps := [`u, `v] ++ lps.eraseIdx upos
+      (u, v, us, us', lps)
+    else
+      let lps := cinfo.levelParams
+      let us := lps.map mkLevelParam
+      (0, 0, us, us, lps)
+    let matchf := .const name us
+    let matchType ← inferType matchf
+    let type ← forallBoundedTelescope matchType info.numParams fun params matchType => do
+      let matchType ← whnf matchType
+      unless matchType.isForall do throwError "resual type {matchType} of {.ofConstName name} not a forall"
+      withLocalDecl `α .implicit (.sort u) fun α => do
+      withLocalDecl `β .implicit (.sort v) fun β => do
+      withLocalDeclD `f (← mkArrow α β) fun f => do
+        let motive ← forallTelescope matchType.bindingDomain! fun xs _ => mkLambdaFVars xs α
+        let motive' ← forallTelescope matchType.bindingDomain! fun xs _ => mkLambdaFVars xs β
+        let matchType ← instantiateForall matchType #[motive]
+        forallBoundedTelescope matchType info.numDiscrs fun discrs matchType => do
+        forallBoundedTelescope matchType info.altNumParams.size fun alts _ => do
+          let lhs := mkAppN (.const name us) (params ++ #[motive] ++ discrs ++ alts)
+          let lhs := .app f lhs
+          let alts' ← alts.mapM fun alt => do
+            let altType ← inferType alt
+            forallTelescope altType fun ys _ =>
+              if ys.size == 1 && altType.bindingDomain!.isConstOf ``Unit then
+                mkLambdaFVars ys (mkApp f (mkApp alt (.const ``Unit.unit [])))
+              else
+                mkLambdaFVars ys (mkApp f (mkAppN alt ys))
+          let rhs := mkAppN (.const name us') (params ++ #[motive'] ++ discrs ++ alts')
+          let type ← mkEq lhs rhs
+          mkForallFVars (#[α,β,f] ++ params ++ discrs ++ alts) type
+    let value ← mkFreshExprSyntheticOpaqueMVar type
+    TermElabM.run' do withoutErrToSorry do
+      runTactic value.mvarId! (← `(Parser.Term.byTactic| by intros; split <;> rfl)).raw .term
+    let value ← instantiateMVars value
+    let decl := Declaration.thmDecl { name := name ++ `lifter, levelParams, type, value }
+    addDecl decl
+
+def isMatchLiftName (env : Environment) (name : Name) : Bool :=
+  if let .str p "lifter" := name
+  then
+    (getMatcherInfoCore? env p).isSome
+  else
+    false
+
+def mkMatchLifterApp (f : Expr) (matcherApp : MatcherApp) : MetaM (Option Expr) := do
+  let some (α, β) := (← inferType f).arrow? |
+    trace[lift_match] "Cannot lift match: {f} is dependent"
+    return none
+  let lifterName := matcherApp.matcherName ++ `lifter
+  let _ ← realizeGlobalName lifterName
+  let lifterArgs := #[α,β,f] ++ matcherApp.params ++ matcherApp.discrs ++ matcherApp.alts
+  -- using mkAppOptM to instantiate the level params
+  let e ← mkAppOptM lifterName (lifterArgs.map some)
+  return some e
+
+builtin_initialize
+  registerReservedNamePredicate isMatchLiftName
+
+  registerReservedNameAction fun name => do
+    if isMatchLiftName (← getEnv) name then
+      let .str p _ := name | return false
+      MetaM.run' <| deriveMatchLift p
+      return true
+    return false
+
+  Lean.registerTraceClass `lift_match

src/Lean/Meta/Match/MatcherApp/Transform.lean

@@ -5,7 +5,7 @@ Authors: Leonardo de Moura, Joachim Breitner
 -/
 
 prelude
-import Lean.Meta.Match
+import Lean.Meta.Match.MatchEqs
 import Lean.Meta.InferType
 import Lean.Meta.Check
 import Lean.Meta.Tactic.Split

src/Lean/Meta/Tactic.lean

@@ -41,3 +41,4 @@ import Lean.Meta.Tactic.FunInd
 import Lean.Meta.Tactic.Rfl
 import Lean.Meta.Tactic.Rewrites
 import Lean.Meta.Tactic.Grind
+import Lean.Meta.Tactic.LiftMatch

src/Lean/Meta/Tactic/LiftMatch.lean

@@ -0,0 +1,177 @@
+/-
+Copyright (c) 2024 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joachim Breitner
+-/
+
+prelude
+import Init.Simproc
+import Lean.Meta.Match.MatcherApp.Basic
+import Lean.Meta.Match.Lift
+import Lean.Meta.KAbstract
+import Lean.Elab.Tactic.Conv.Basic
+
+/-!
+This module implements the `liftMatch` simpproc and the `lift_match` conv tactic.
+-/
+
+
+open Lean Meta Elab Term
+
+namespace Lean.Meta
+
+section LiftMatch
+
+inductive MatcherOrIteApp where
+  | matcher (matcherApp : MatcherApp)
+  | ite (dite : Bool) (α P inst t e : Expr)
+
+private def _root_.Lean.Expr.constLams? : Expr → Option Expr
+  | .lam _ _ b _ => constLams? b
+  | e => if e.hasLooseBVars then none else some e
+
+def MatcherOrIteApp.motive? : MatcherOrIteApp -> Option Expr
+  | matcher matcherApp => matcherApp.motive.constLams?
+  | ite _ α _ _ _ _ => some α
+
+def mkMatchOrIteLiftApp (f : Expr) (moi : MatcherOrIteApp) : MetaM (Option Expr) := do
+  match moi with
+  | .matcher matcherApp => mkMatchLifterApp f matcherApp
+  | .ite dite α P inst t e =>
+      let some (_, β) := (← inferType f).arrow? |
+        trace[lift_match] "Cannot lift match: {f} is dependent"
+        return none
+      let lifterName := if dite then ``apply_dite else ``apply_ite
+      let e ← mkAppOptM lifterName #[some α, some β, some f, some P, some inst, some t, some e]
+      return some e
+
+
+/--
+Like matchMatcherApp, but also matches ite/dite.
+-/
+private def matchMatcherOrIteApp? (e : Expr) : MetaM (Option MatcherOrIteApp) := do
+  match_expr e with
+  | ite α P inst t e => return some (.ite false α P inst t e)
+  | dite α P inst t e => return some (.ite true α P inst t e)
+  | _ =>
+    if let some matcherApp ← matchMatcherApp? e then
+      if matcherApp.remaining.isEmpty then
+        return .some (.matcher matcherApp)
+  return none
+
+/--
+Finds the first possible liftable match (or (d)ite).
+-/
+private partial def findMatchToLift? (e : Expr) (far : Bool) (depth : Nat := 0) : MetaM (Option (Expr × MatcherOrIteApp)) := do
+  if !far && depth > 1 then
+    return none
+
+  if e.isApp then
+    if depth > 0 then
+      if let some matcherApp ← matchMatcherOrIteApp? e then
+        return some (e, matcherApp)
+
+
+    let args := e.getAppArgs
+    if e.isAppOf ``ite then
+      -- Special-handling for if-then-else:
+      -- * We do not want to lift out of the branches.
+      -- * We want to be able to lift out of
+      --      @ite ([ ] = true) (instDecidableEqBool [ ] true) t e
+      --   but doing it one application at a time does not work due to the dependency.
+      --   So to work around this, we do not bump the depth counter here.
+      if h : args.size > 1 then
+        if let some r ← findMatchToLift? args[1] far depth then
+          return some r
+    else
+      for a in args do
+        if let some r ← findMatchToLift? a far (depth + 1) then
+          return some r
+
+  if e.isLet then
+    if let some r ← findMatchToLift? e.letValue! far (depth + 1) then
+      return some r
+
+  return none
+
+/--
+In `e`, try to find a `match` expression or `ite` application that we can lift
+to the root. Returns the new expression `s` and a proof of `e = s`.
+-/
+def findAndLiftMatch (e : Expr) (far : Bool) : MetaM (Option (Expr × Expr)) := do
+  unless far do
+    -- In the simproc: We could, but for now we do not lift out of props
+    if ← Meta.isProp e then return none
+  let some (me, matcherApp) ← findMatchToLift? e far| return none
+  -- We do not handle dependent motives
+  let some α := matcherApp.motive? |
+    trace[lift_match] "Cannot lift match: motive depends on targets"
+    return none
+  -- Using kabstract, rather than just abstracting over the single occurrence of `me` in `e` with
+  -- helps if later arguments depend on the abstracted argument,
+  -- in particular with ``ite's `Decidable c` parameter
+  let f := (mkLambda `x .default α (← kabstract e me)).eta
+  -- Abstracting over the argument can result in a type incorrect `f` (like in `rw`)
+  unless (← isTypeCorrect f) do
+    trace[lift_match] "Cannot lift match: context is not type correct"
+    return none
+  let some proof ← mkMatchOrIteLiftApp f matcherApp | return none
+  let type ← inferType proof
+  let some (_, _, rhs) := type.eq?
+    | throwError "lift_match: Unexpected non-equality type:{indentExpr type}"
+  return some (rhs, proof)
+
+end LiftMatch
+
+end Lean.Meta
+
+/-!
+The simproc tactic
+-/
+
+section Simproc
+
+/--
+Lifts out `match` expressions, or, equivalently, pushes function applications into the
+branches of a match. For example it can rewrite
+```
+f (match o with | some x => x + 1 | none => 0)
+```
+to
+```
+match o with | some x => f (x + 1) | none => f 0
+```
+
+For the purposes of this simproc, `if-then-else` expressions are treated like `match` expressions.
+
+It can only lift matches with a non-dependent motive, no extra arguments and when the context
+(here `fun x => f x`) is type-correct and is not a proposition.
+
+It is recommended to enable this simproc only selectively, and not by default, as it looks for
+match expression to lift at every step of the simplifier.
+
+Also see the `conv`-tactic `lift_match`.
+-/
+builtin_simproc_decl liftMatch (_) := fun e => do
+  let some (rhs, proof) ← findAndLiftMatch (far := false) e | return .continue
+  return .visit { expr := rhs, proof? := some proof }
+
+end Simproc
+
+
+/-!
+The conv tactic
+-/
+
+namespace Lean.Elab.Tactic.Conv
+
+@[builtin_tactic Lean.Parser.Tactic.Conv.liftMatch]
+def evalLiftMatch : Tactic := fun _ => do
+  let mvarId ← getMainGoal
+  mvarId.withContext do
+    let lhs ← getLhs
+    let some (rhs, proof) ← findAndLiftMatch (far := true) lhs
+      | throwError "cannot find match to lift"
+    updateLhs rhs proof
+
+end Lean.Elab.Tactic.Conv

tests/lean/run/lift_match.lean

@@ -0,0 +1,311 @@
+import Lean.Elab.Command
+import Lean.Meta.Match.MatchEqsExt
+
+def test1 : Nat → Nat
+  | 0 => 0
+  | _+1 => 42
+
+-- set_option pp.match false
+
+/--
+info: test1.match_1.lifter.{u, v} {α : Sort u} {β : Sort v} (f : α → β) (x✝ : Nat) (h_1 : Unit → (fun x => α) 0)
+  (h_2 : (n : Nat) → (fun x => α) n.succ) :
+  f
+      (match x✝ with
+      | 0 => h_1 ()
+      | n.succ => h_2 n) =
+    match x✝ with
+    | 0 => f (h_1 ())
+    | n.succ => f (h_2 n)
+-/
+#guard_msgs in
+#check test1.match_1.lifter
+
+def test2 (α β) : α ∨ β → γ → (β ∨ α) ∧ γ
+  | .inl x, y => ⟨.inr x, y⟩
+  | .inr x, y => ⟨.inl x, y⟩
+
+set_option pp.proofs true in
+/--
+info: test2.match_1.lifter {α β : Prop} (f : α → β) {γ : Prop} (α✝ β✝ : Prop) (x✝ : α✝ ∨ β✝) (x✝¹ : γ)
+  (h_1 : ∀ (x : α✝) (y : γ), (fun x x => α) (Or.inl x) y) (h_2 : ∀ (x : β✝) (y : γ), (fun x x => α) (Or.inr x) y) :
+  f
+      (match x✝, x✝¹ with
+      | Or.inl x, y => h_1 x y
+      | Or.inr x, y => h_2 x y) =
+    match x✝, x✝¹ with
+    | Or.inl x, y => f (h_1 x y)
+    | Or.inr x, y => f (h_2 x y)
+-/
+#guard_msgs in
+#check test2.match_1.lifter
+
+-- This fails if there is no splitter theorem for a match
+
+/--
+error: Failed to realize constant Nat.lt_or_gt_of_ne.match_1.lifter:
+  Cannot construct match lifter:
+    Could not construct splitter for Nat.lt_or_gt_of_ne.match_1
+---
+error: Failed to realize constant Nat.lt_or_gt_of_ne.match_1.lifter:
+  Cannot construct match lifter:
+    Could not construct splitter for Nat.lt_or_gt_of_ne.match_1
+---
+error: unknown identifier 'Nat.lt_or_gt_of_ne.match_1.lifter'
+-/
+#guard_msgs in
+#check Nat.lt_or_gt_of_ne.match_1.lifter
+
+-- A typical example
+
+theorem List.filter_map' (f : β → α) (l : List β) : filter p (map f l) = map f (filter (p ∘ f) l) := by
+  induction l <;> simp [filter, *, liftMatch]
+
+-- Using the lift_match conv tactic
+
+theorem List.filter_map'' (f : β → α) (l : List β) : filter p (map f l) = map f (filter (p ∘ f) l) := by
+  induction l
+  · simp
+  · simp only [filter]
+    conv => rhs; lift_match
+    simp only [Function.comp_apply, map_cons, *]
+
+-- Using the lift_match tactic
+
+-- This works in principle, but isn't very useful, because the simplifier, even with
+-- `(contextual := true)`, does not simplify the duplicated match target in the alternatives.
+-- Looks like that would require congruence theorems for matchers.
+
+/-- warning: declaration uses 'sorry' -/
+#guard_msgs in
+theorem List.filter_map''' (f : β → α) (l : List β) : filter p (map f l) = map f (filter (p ∘ f) l) := by
+  induction l
+  · simp
+  · simp only [filter]
+    lift_match
+    simp (config := {contextual := true}) [*]
+    -- fail
+    sorry
+
+
+-- A simple example
+
+example (o : Option Bool) :
+  (match o with | some b => b | none => false)
+    = !(match o with | some b => !b | none => true) := by
+  simp [liftMatch]
+
+-- Can float out of ite-condition
+/--
+error: tactic 'fail' failed
+o : Option Bool
+P : Nat → Prop
+⊢ P
+    (match o with
+    | some b => if b = true then 1 else 2
+    | none => if True then 1 else 2)
+-/
+#guard_msgs in
+example (o : Option Bool) (P : Nat → Prop):
+  P (if (match o with | some b => b | none => true) then 1 else 2) := by
+  simp only [liftMatch]
+  fail
+
+-- Cannot lift out of ite-branch
+example (b : Bool) (o : Option Bool) (P : Bool → Prop) (abort : ∀ b, P b):
+  P (if b then (match o with | some b => b | none => true) else b) := by
+  fail_if_success simp only [liftMatch]
+  apply abort
+
+-- Can float out of a match target (aka case-of-case)
+/--
+error: tactic 'fail' failed
+o : Option Bool
+P : Nat → Prop
+⊢ P
+    (match o with
+    | some b =>
+      match b with
+      | true => 1
+      | false => 2
+    | none => 1)
+-/
+#guard_msgs in
+example (o : Option Bool) (P : Nat → Prop):
+  P (match (match o with | some b => b | none => true) with | true => 1 | false => 2) := by
+  simp only [liftMatch]
+  fail
+
+-- Dependent motive; must not rewrite
+
+set_option trace.lift_match true in
+/-- info: [lift_match] Cannot lift match: motive depends on targets -/
+#guard_msgs in
+example (o : Option Bool) (motive : Bool → Type) (P : {b : Bool} → motive b → Prop)
+  (f : (x : Bool) → motive x) (g : {x : Bool} → motive x → motive x)
+  (abort : ∀ b (x : motive b), P x) :
+  P (g (match (motive := ∀ b, motive b.isSome) o with | some _ => f true | none => f false)) := by
+  fail_if_success simp [liftMatch]
+  apply abort
+
+-- Dependent context; must not rewrite
+
+set_option trace.lift_match true in
+/-- info: [lift_match] Cannot lift match: f is dependent -/
+#guard_msgs in
+example (o : Option Bool) (motive : Bool → Type) (P : {b : Bool} → motive b → Prop)
+  (f : (x : Bool) → motive x)
+  (abort : ∀ b (x : motive b), P x) :
+  P (f (match (motive := ∀ _, Bool) o with | some b => b | none => false)) := by
+  fail_if_success simp [liftMatch]
+  apply abort
+
+-- Context depends on the concrete value of the match, must not rewrite
+
+set_option trace.lift_match true in
+/-- info: [lift_match] Cannot lift match: context is not type correct -/
+#guard_msgs in
+example (o : Option Bool)
+  (f : (x : Bool) → (h : x = (match o with | some b => b | none => false)) → Bool)
+  (abort : ∀ (P : Prop), P) :
+  f (match (motive := ∀ _, Bool) o with | some b => b | none => false) rfl = true := by
+  fail_if_success simp [liftMatch]
+  apply abort
+
+-- Can float out of a let (Only relevant with zeta := false)
+
+/--
+error: tactic 'fail' failed
+o : Option Bool
+P : Bool → Prop
+⊢ P
+    (match o with
+    | some b =>
+      let b := b;
+      !b
+    | none =>
+      let b := true;
+      !b)
+-/
+#guard_msgs in
+example (o : Option Bool) (P : Bool → Prop):
+  P (let b := match o with | some b => b | none => true; !b) := by
+  simp -zeta only [liftMatch]
+  fail
+
+/-
+This following code tries to create all float theorems for all matches found in the environment.
+-/
+
+-- At the time of writing, the following two quite large matches fail by running out of heartbeat
+-- #check Lean.Expr.name?.match_1.float
+-- #check Lean.Meta.reduceNat?.match_1.float
+
+/-
+open Lean Meta in
+run_meta do
+  for es in (Match.Extension.extension.toEnvExtension.getState (← getEnv)).importedEntries do
+    for e in es do
+      -- Let's not look at matchers that eliminate to Prop only
+      if e.info.uElimPos?.isNone then continue
+      withCurrHeartbeats do
+        tryCatchRuntimeEx do
+          let hasSplitter ← try
+              discard <| Lean.Meta.Match.getEquationsFor e.name
+              pure true
+            catch _ => pure false
+          if hasSplitter then
+            let floatName := e.name ++ `float
+            unless (← hasConst floatName) do
+              executeReservedNameAction floatName
+         fun ex =>
+          logInfo m!"Failed to handle {e.name}:{ex.toMessageData}"
+-/
+
+-- Testing if-then-else
+
+/--
+error: tactic 'fail' failed
+P : Nat → Prop
+f : Nat → Nat
+b : Bool
+⊢ P (if b = true then f 1 else f 2)
+-/
+#guard_msgs in
+example (P : Nat → Prop) (f : Nat → Nat) (b : Bool) :
+  P (f (if b then 1 else 2)) := by
+  simp only [liftMatch]
+  fail
+
+-- Dependent f
+
+/--
+error: simp made no progress
+---
+info: [lift_match] Cannot lift match: f is dependent
+-/
+#guard_msgs in
+set_option trace.lift_match true in
+example (P : {n : Nat} → Fin n → Prop) (f : (n : Nat) → Fin n) (b : Bool) :
+  P (f (if b then 1 else 2)) := by
+  simp only [liftMatch]
+  fail
+
+-- Somewhat dependent f, but abstracting still succeeds
+
+/--
+error: tactic 'fail' failed
+P : Nat → Prop
+f : (n : Nat) → DecidableEq (Fin n) → Nat
+b : Bool
+⊢ P (if b = true then f 1 inferInstance else f 2 inferInstance)
+-/
+#guard_msgs in
+example (P : Nat → Prop) (f : (n : Nat) → DecidableEq (Fin n) → Nat) (b : Bool) :
+  P (f (if b then 1 else 2) inferInstance) := by
+  simp only [liftMatch]
+  fail
+
+-- Testing dependent if-then-else
+
+/--
+error: tactic 'fail' failed
+P : Nat → Prop
+f : Nat → Nat
+b : Bool
+⊢ P (if h : b = true then f 1 else f 2)
+-/
+#guard_msgs in
+example (P : Nat → Prop) (f : Nat → Nat) (b : Bool) :
+  P (f (if h : b then 1 else 2)) := by
+  simp only [liftMatch]
+  fail
+
+-- Dependent f
+
+/--
+error: simp made no progress
+---
+info: [lift_match] Cannot lift match: f is dependent
+-/
+#guard_msgs in
+set_option trace.lift_match true in
+example (P : {n : Nat} → Fin n → Prop) (f : (n : Nat) → Fin n) (b : Bool) :
+  P (f (if h : b then 1 else 2)) := by
+  simp only [liftMatch]
+  fail
+
+-- Somewhat dependent f, but abstracting still succeeds
+
+/--
+error: tactic 'fail' failed
+P : Nat → Prop
+f : (n : Nat) → DecidableEq (Fin n) → Nat
+b : Bool
+⊢ P (if h : b = true then f 1 inferInstance else f 2 inferInstance)
+-/
+#guard_msgs in
+example (P : Nat → Prop) (f : (n : Nat) → DecidableEq (Fin n) → Nat) (b : Bool) :
+  P (f (if h : b then 1 else 2) inferInstance) := by
+  simp only [liftMatch]
+  fail