Add rewrite rule: List.any with equality lambda to List.elem

Blaster/Optimize/Rewriting/OptimizeApp.lean

@@ -7,6 +7,7 @@ import Blaster.Optimize.Rewriting.OptimizeDecideBoolBinary
 import Blaster.Optimize.Rewriting.OptimizeExists
 import Blaster.Optimize.Rewriting.OptimizeInt
 import Blaster.Optimize.Rewriting.OptimizeITE
+import Blaster.Optimize.Rewriting.OptimizeListAny
 import Blaster.Optimize.Rewriting.OptimizeNat
 import Blaster.Optimize.Rewriting.OptimizeString
 import Blaster.Optimize.OptimizeStack
@@ -64,6 +65,7 @@ def optimizeAppAux (f : Expr) (args: Array Expr) : TranslateEnvT Expr := do
   if let some e ← optimizePropBinary? f args then return e
   if let some e ← optimizeBoolNot? f args then return e
   if let some e ← optimizeBoolBinary? f args then return e
+  if let some e ← optimizeListAny? f args then return e
   if let some e ← optimizeEquality? f args then return e
   if let some e ← optimizeNat? f args then return e
   if let some e ← optimizeInt? f args then return e

Blaster/Optimize/Rewriting/OptimizeListAny.lean

@@ -0,0 +1,56 @@
+import Lean
+import Blaster.Optimize.Rewriting.Utils
+
+open Lean Meta
+namespace Blaster.Optimize
+
+/-- Detect `List.any l (λ x => BEq.beq x k)` or `List.any l (λ x => BEq.beq k x)`
+    and rewrite to `List.elem k l`.
+
+    @List.any : {α : Type} → List α → (α → Bool) → Bool
+    args[0] = α, args[1] = l, args[2] = f
+
+    @BEq.beq : {α : Type} → [inst : BEq α] → α → α → Bool
+    beqArgs[0] = α, beqArgs[1] = inst, beqArgs[2] = lhs, beqArgs[3] = rhs
+
+    @List.elem : {α : Type} → [inst : BEq α] → α → List α → Bool
+
+    Uses definitional equality checking to handle custom types with derived BEq
+    instances, where the elaborator may expand (==) into primitive constructs.
+-/
+
+def optimizeListAny? (f : Expr) (args : Array Expr) : TranslateEnvT (Option Expr) := do
+  let Expr.const ``List.any _ := f | return none
+  if args.size != 3 then return none
+  let α := args[0]!
+  let l := args[1]!
+  let lam := args[2]!
+  let Expr.lam binderName binderType body binderInfo := lam | return none
+  withLocalDecl binderName binderInfo binderType λ x => do
+    let body := body.instantiate1 x
+    let beqType := mkApp (mkConst ``BEq [levelZero]) α
+    let inst ← try synthInstance beqType
+      catch _ => do
+      return none
+    let kMVar ← mkFreshExprMVar α
+    let rhs_candidate := mkAppN (mkConst ``BEq.beq [levelZero]) #[α, inst, x, kMVar]
+    let rhs_matches ← isDefEq body rhs_candidate
+    let k ← if rhs_matches then do
+        let k ← instantiateMVars kMVar
+        pure k
+      else do
+        let kMVar2 ← mkFreshExprMVar α
+        let lhs_candidate := mkAppN (mkConst ``BEq.beq [levelZero]) #[α, inst, kMVar2, x]
+        let lhs_matches ← isDefEq body lhs_candidate
+        if lhs_matches then do
+          let k ← instantiateMVars kMVar2
+          pure k
+        else do
+          return none
+    if k.hasLooseBVars then return none
+    setRestart
+    let u <- getLevel α
+    let u' := u.normalize.dec.getD Lean.Level.zero
+    return mkAppN (mkConst ``List.elem [u']) #[α, inst, k, l]
+
+end Blaster.Optimize

Tests/Optimize.lean

@@ -1,15 +1,15 @@
-
 import Tests.Optimize.FunPropagation
 import Tests.Optimize.NormChoiceApp
 import Tests.Optimize.OptimizeBEq
 import Tests.Optimize.OptimizeBool
 import Tests.Optimize.OptimizeConst
+import Tests.Optimize.OptimizeDecide
 import Tests.Optimize.OptimizeEq
 import Tests.Optimize.OptimizeExists
-import Tests.Optimize.OptimizeDecide
 import Tests.Optimize.OptimizeForAll
 import Tests.Optimize.OptimizeITE
 import Tests.Optimize.OptimizeLet
+import Tests.Optimize.OptimizeListAny
 import Tests.Optimize.OptimizeMatch
 import Tests.Optimize.OptimizeNat
 import Tests.Optimize.OptimizeProp

Tests/Optimize/OptimizeConst/NormHOF.lean

@@ -441,12 +441,12 @@ def mapOptionDefault (x : Nat) (y : Option Nat) : Nat :=
 -- ∀ (x : Nat) (xs : List Nat),
 --   List.any xs (λ a => a == x) → ∃ (p : Prop), p ∈ (List.map (Eq x) xs) ∧ p ===>
 -- ∀ (x : Nat) (xs : List Nat),
---   true = List.any xs (λ a => x == a) → ∃ (p : Prop), p ∧ List.Mem p (List.map (λ a => Eq x a) xs)
+--   true = List.elem x xs → ∃ (p : Prop), p ∧ List.Mem p (List.map (λ a => x = a) xs)
 #testOptimize [ "ConstPartialFunArg_4" ]
     ∀ (x : Nat) (xs : List Nat),
       List.any xs (λ a => a == x) → ∃ (p : Prop), p ∈ (List.map (Eq x) xs) ∧ p ===>
     ∀ (x : Nat) (xs : List Nat),
-      true = List.any xs (λ a => x == a) → ∃ (p : Prop), p ∧ List.Mem p (List.map (λ a => Eq x a) xs)
+      true = List.elem x xs → ∃ (p : Prop), p ∧ List.Mem p (List.map (λ a => Eq x a) xs)
 
 /-! Test cases to ensure that opaque functions passed as arguments are properly normalized. -/
 
@@ -473,11 +473,9 @@ def mapOptionDefault (x : Nat) (y : Option Nat) : Nat :=
 
 def boolWapper (f : Nat → Nat → Bool) (x : Nat) (y : Nat) := f x y
 
--- ∀ (x : Nat) (xs : List Nat), List.any xs (beqWapper Nat.beq x) → List.contains xs x ===>
--- ∀ (x : Nat) (xs : List Nat), true = List.any xs (λ a => x == a) → true = List.elem x xs
+-- ∀ (x : Nat) (xs : List Nat), List.any xs (beqWapper Nat.beq x) → List.contains xs x ===> True
 #testOptimize [ "ConstNormOpaqueFunArg_5" ]
-  ∀ (x : Nat) (xs : List Nat), List.any xs (boolWapper Nat.beq x) → List.contains xs x ===>
-  ∀ (x : Nat) (xs : List Nat), true = List.any xs (λ a => x == a) → true = List.elem x xs
+  ∀ (x : Nat) (xs : List Nat), List.any xs (boolWapper Nat.beq x) → List.contains xs x ===> True
 
 -- ∀ (x : Nat) (xs : List Nat), List.any (λ a => x ≤ a) xs = List.any xs (Nat.ble x)  ===> True
 #testOptimize [ "ConstNormOpaqueFunArg_6" ]

Tests/Optimize/OptimizeListAny.lean

@@ -0,0 +1,27 @@
+import Lean
+import Tests.Utils
+import Tests.Smt.Benchmarks.ValidatorsExamples.PlutusLedgerAPI.VDummy
+
+open Lean Elab Command Term
+open Tests.ValidatorsExamples.PlutusLedgerAPI
+
+namespace Test.OptimizeListAny
+
+-- List.any l (λ x => x == k) ===> true = List.elem k l
+#testOptimize [ "ListAnyElem_1" ]
+  ∀ (l : List Nat) (k : Nat), List.any l (λ x => x == k) ===>
+  ∀ (l : List Nat) (k : Nat), true = List.elem k l
+
+-- List.any l (λ x => k == x) ===> true = List.elem k l
+#testOptimize [ "ListAnyElem_2" ]
+  ∀ (l : List Nat) (k : Nat), List.any l (λ x => k == x) ===>
+  ∀ (l : List Nat) (k : Nat), true = List.elem k l
+
+-- Custom type with deriving BEq
+#testOptimize [ "ListAnyElem_3" ]
+  ∀ (l : List VerificationKeyHash) (k : VerificationKeyHash),
+    List.any l (λ x => x == k) ===>
+  ∀ (l : List VerificationKeyHash) (k : VerificationKeyHash),
+    true = List.elem k l
+
+end Test.OptimizeListAny