From bba4ef3728e9c6fd99b3fa3d9dd5dca1411379cc Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Sun, 3 Mar 2024 18:51:04 -0800 Subject: [PATCH] feat: `simproc`s for folding numeric literals (#3586) This PR folds exposed `BitVec` (`Fin`, `UInt??`, and `Int`) ground literals. cc @shigoel --- src/Lean/Meta/Offset.lean | 41 ++++++-------- .../Meta/Tactic/LinearArith/Nat/Basic.lean | 53 +++++++----------- .../Tactic/Simp/BuiltinSimprocs/BitVec.lean | 11 ++++ .../Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean | 9 +++ .../Meta/Tactic/Simp/BuiltinSimprocs/Int.lean | 10 ++++ .../Tactic/Simp/BuiltinSimprocs/UInt.lean | 6 ++ tests/lean/run/foldLits.lean | 56 +++++++++++++++++++ 7 files changed, 129 insertions(+), 57 deletions(-) create mode 100644 tests/lean/run/foldLits.lean diff --git a/src/Lean/Meta/Offset.lean b/src/Lean/Meta/Offset.lean index ba46b435b2..402afec72b 100644 --- a/src/Lean/Meta/Offset.lean +++ b/src/Lean/Meta/Offset.lean @@ -18,9 +18,10 @@ private abbrev withInstantiatedMVars (e : Expr) (k : Expr → OptionT MetaM α) k eNew def isNatProjInst (declName : Name) (numArgs : Nat) : Bool := - (numArgs == 4 && (declName == ``Add.add || declName == ``Sub.sub || declName == ``Mul.mul)) - || (numArgs == 6 && (declName == ``HAdd.hAdd || declName == ``HSub.hSub || declName == ``HMul.hMul)) - || (numArgs == 3 && declName == ``OfNat.ofNat) + (numArgs == 4 && (declName == ``Add.add || declName == ``Sub.sub || declName == ``Mul.mul || declName == ``Div.div || declName == ``Mod.mod || declName == ``NatPow.pow)) + || (numArgs == 5 && (declName == ``Pow.pow)) + || (numArgs == 6 && (declName == ``HAdd.hAdd || declName == ``HSub.hSub || declName == ``HMul.hMul || declName == ``HDiv.hDiv || declName == ``HMod.hMod || declName == ``HPow.hPow)) + || (numArgs == 3 && declName == ``OfNat.ofNat) /-- Evaluate simple `Nat` expressions. @@ -35,31 +36,21 @@ partial def evalNat (e : Expr) : OptionT MetaM Nat := do | _ => failure where visit e := do - let f := e.getAppFn - match f with - | .mvar .. => withInstantiatedMVars e evalNat - | .const c _ => - let nargs := e.getAppNumArgs - if c == ``Nat.succ && nargs == 1 then - let v ← evalNat (e.getArg! 0) - return v+1 - else if c == ``Nat.add && nargs == 2 then - let v₁ ← evalNat (e.getArg! 0) - let v₂ ← evalNat (e.getArg! 1) - return v₁ + v₂ - else if c == ``Nat.sub && nargs == 2 then - let v₁ ← evalNat (e.getArg! 0) - let v₂ ← evalNat (e.getArg! 1) - return v₁ - v₂ - else if c == ``Nat.mul && nargs == 2 then - let v₁ ← evalNat (e.getArg! 0) - let v₂ ← evalNat (e.getArg! 1) - return v₁ * v₂ - else if isNatProjInst c nargs then + match_expr e with + | Nat.succ a => return (← evalNat a) + 1 + | Nat.add a b => return (← evalNat a) + (← evalNat b) + | Nat.sub a b => return (← evalNat a) - (← evalNat b) + | Nat.mul a b => return (← evalNat a) * (← evalNat b) + | Nat.div a b => return (← evalNat a) / (← evalNat b) + | Nat.mod a b => return (← evalNat a) % (← evalNat b) + | Nat.pow a b => return (← evalNat a) ^ (← evalNat b) + | _ => + let e ← instantiateMVarsIfMVarApp e + let f := e.getAppFn + if f.isConst && isNatProjInst f.constName! e.getAppNumArgs then evalNat (← unfoldProjInst? e) else failure - | _ => failure mutual diff --git a/src/Lean/Meta/Tactic/LinearArith/Nat/Basic.lean b/src/Lean/Meta/Tactic/LinearArith/Nat/Basic.lean index 3be8fd9ece..6290ec44fe 100644 --- a/src/Lean/Meta/Tactic/LinearArith/Nat/Basic.lean +++ b/src/Lean/Meta/Tactic/LinearArith/Nat/Basic.lean @@ -74,42 +74,31 @@ def addAsVar (e : Expr) : M LinearExpr := do partial def toLinearExpr (e : Expr) : M LinearExpr := do match e with - | Expr.lit (Literal.natVal n) => return num n - | Expr.mdata _ e => toLinearExpr e - | Expr.const ``Nat.zero .. => return num 0 - | Expr.app .. => visit e - | Expr.mvar .. => visit e - | _ => addAsVar e + | .lit (.natVal n) => return num n + | .mdata _ e => toLinearExpr e + | .const ``Nat.zero .. => return num 0 + | .app .. => visit e + | .mvar .. => visit e + | _ => addAsVar e where visit (e : Expr) : M LinearExpr := do - let f := e.getAppFn - match f with - | Expr.mvar .. => - let eNew ← instantiateMVars e - if eNew != e then - toLinearExpr eNew + match_expr e with + | Nat.succ a => return inc (← toLinearExpr a) + | Nat.add a b => return add (← toLinearExpr a) (← toLinearExpr b) + | Nat.mul a b => + match (← evalNat a |>.run) with + | some k => return mulL k (← toLinearExpr b) + | none => match (← evalNat b |>.run) with + | some k => return mulR (← toLinearExpr a) k + | none => addAsVar e + | _ => + let e ← instantiateMVarsIfMVarApp e + let f := e.getAppFn + if f.isConst && isNatProjInst f.constName! e.getAppNumArgs then + let some e ← unfoldProjInst? e | addAsVar e + toLinearExpr e else addAsVar e - | Expr.const declName .. => - let numArgs := e.getAppNumArgs - if declName == ``Nat.succ && numArgs == 1 then - return inc (← toLinearExpr e.appArg!) - else if declName == ``Nat.add && numArgs == 2 then - return add (← toLinearExpr (e.getArg! 0)) (← toLinearExpr (e.getArg! 1)) - else if declName == ``Nat.mul && numArgs == 2 then - match (← evalNat (e.getArg! 0) |>.run) with - | some k => return mulL k (← toLinearExpr (e.getArg! 1)) - | none => match (← evalNat (e.getArg! 1) |>.run) with - | some k => return mulR (← toLinearExpr (e.getArg! 0)) k - | none => addAsVar e - else if isNatProjInst declName numArgs then - if let some e ← unfoldProjInst? e then - toLinearExpr e - else - addAsVar e - else - addAsVar e - | _ => addAsVar e partial def toLinearCnstr? (e : Expr) : M (Option LinearCnstr) := do let f := e.getAppFn diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean index 0ac8c1ced9..37eaa364bf 100644 --- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean +++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/BitVec.lean @@ -268,4 +268,15 @@ builtin_simproc [simp, seval] reduceAllOnes (allOnes _) := fun e => do let some n ← Nat.fromExpr? n | return .continue return .done { expr := toExpr (allOnes n) } +builtin_simproc [simp, seval] reduceBitVecOfFin (BitVec.ofFin _) := fun e => do + let_expr BitVec.ofFin w v ← e | return .continue + let some w ← evalNat w |>.run | return .continue + let some ⟨_, v⟩ ← getFinValue? v | return .continue + return .done { expr := toExpr (BitVec.ofNat w v.val) } + +builtin_simproc [simp, seval] reduceBitVecToFin (BitVec.toFin _) := fun e => do + let_expr BitVec.toFin _ v ← e | return .continue + let some ⟨_, v⟩ ← getBitVecValue? v | return .continue + return .done { expr := toExpr v.toFin } + end BitVec diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean index 55aa9a2d04..4efce2b547 100644 --- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean +++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Fin.lean @@ -71,4 +71,13 @@ builtin_simproc [simp, seval] isValue ((OfNat.ofNat _ : Fin _)) := fun e => do return .done { expr := e } return .done { expr := toExpr v } +builtin_simproc [simp, seval] reduceFinMk (Fin.mk _ _) := fun e => do + let_expr Fin.mk n v _ ← e | return .continue + let some n ← evalNat n |>.run | return .continue + let some v ← getNatValue? v | return .continue + if h : n > 0 then + return .done { expr := toExpr (Fin.ofNat' v h) } + else + return .continue + end Fin diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Int.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Int.lean index c031524bcc..38b0b1ff71 100644 --- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Int.lean +++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/Int.lean @@ -89,4 +89,14 @@ builtin_simproc [simp, seval] reduceBNe (( _ : Int) != _) := reduceBoolPred `` builtin_simproc [simp, seval] reduceAbs (natAbs _) := reduceNatCore ``natAbs natAbs builtin_simproc [simp, seval] reduceToNat (Int.toNat _) := reduceNatCore ``Int.toNat Int.toNat +builtin_simproc [simp, seval] reduceNegSucc (Int.negSucc _) := fun e => do + let_expr Int.negSucc a ← e | return .continue + let some a ← getNatValue? a | return .continue + return .done { expr := toExpr (-(Int.ofNat a + 1)) } + +builtin_simproc [simp, seval] reduceOfNat (Int.ofNat _) := fun e => do + let_expr Int.ofNat a ← e | return .continue + let some a ← getNatValue? a | return .continue + return .done { expr := toExpr (Int.ofNat a) } + end Int diff --git a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/UInt.lean b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/UInt.lean index 5e9868bb61..955d48232e 100644 --- a/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/UInt.lean +++ b/src/Lean/Meta/Tactic/Simp/BuiltinSimprocs/UInt.lean @@ -60,6 +60,12 @@ builtin_simproc [simp, seval] $(mkIdent `reduceOfNatCore):ident ($ofNatCore _ _) let value := $(mkIdent ofNat) value return .done { expr := toExpr value } +builtin_simproc [simp, seval] $(mkIdent `reduceOfNat):ident ($(mkIdent ofNat) _) := fun e => do + unless e.isAppOfArity $(quote ofNat) 1 do return .continue + let some value ← Nat.fromExpr? e.appArg! | return .continue + let value := $(mkIdent ofNat) value + return .done { expr := toExpr value } + builtin_simproc [simp, seval] $(mkIdent `reduceToNat):ident ($toNat _) := fun e => do unless e.isAppOfArity $(quote toNat.getId) 1 do return .continue let some v ← ($fromExpr e.appArg!) | return .continue diff --git a/tests/lean/run/foldLits.lean b/tests/lean/run/foldLits.lean new file mode 100644 index 0000000000..d3eba2bf83 --- /dev/null +++ b/tests/lean/run/foldLits.lean @@ -0,0 +1,56 @@ +open BitVec + +example : (Fin.mk 5 (by decide) : Fin 10) + 2 = x := by + simp + guard_target =ₛ 7 = x + sorry + +example : (Fin.mk 5 (by decide) : Fin 10) + 2 = x := by + simp (config := { ground := true }) only + guard_target =ₛ 7 = x + sorry + +example : (BitVec.ofFin (Fin.mk 2 (by decide)) : BitVec 32) + 2 = x := by + simp + guard_target =ₛ 4#32 = x + sorry + +example : (BitVec.ofFin (Fin.mk 2 (by decide)) : BitVec 32) + 2 = x := by + simp (config := { ground := true }) only + guard_target =ₛ 4#32 = x + sorry + +example : (BitVec.ofFin 2 : BitVec 32) + 2 = x := by + simp + guard_target =ₛ 4#32 = x + sorry + +example (h : -2 = x) : Int.negSucc 3 + 2 = x := by + simp + guard_target =ₛ -2 = x + assumption + +example (h : -2 = x) : Int.negSucc 3 + 2 = x := by + simp (config := { ground := true }) only + guard_target =ₛ -2 = x + assumption + +example : Int.ofNat 3 + 2 = x := by + simp + guard_target =ₛ 5 = x + sorry + +example : Int.ofNat 3 + 2 = x := by + simp (config := { ground := true }) only + guard_target =ₛ 5 = x + sorry + +example (h : 5 = x) : UInt32.ofNat 2 + 3 = x := by + simp + guard_target =ₛ 5 = x + assumption + +example (h : 5 = x) : UInt32.ofNat 2 + 3 = x := by + simp (config := { ground := true }) only + guard_target =ₛ 5 = x + assumption