From b5bf0b1d05e3da179b0d7fbc4306e9e4e8856f16 Mon Sep 17 00:00:00 2001
From: Kim Morrison <scott.morrison@gmail.com>
Date: Wed, 14 May 2025 07:01:34 +0800
Subject: [PATCH] chore: cleanup of grind if-then-else example (#8326)

---
 tests/lean/run/grind_ite.lean | 42 ++++++++++++++++++++++++++++-------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/tests/lean/run/grind_ite.lean b/tests/lean/run/grind_ite.lean
index 48299c8960..ce9e544fb6 100644
--- a/tests/lean/run/grind_ite.lean
+++ b/tests/lean/run/grind_ite.lean
@@ -1,4 +1,5 @@
 import Std.Data.HashMap.Lemmas
+import Std.Data.TreeMap.Lemmas
 
 /-!
 # If normalization
@@ -115,11 +116,8 @@ set_option grind.warning false
 
 namespace IfExpr
 
--- We tell `grind` it can unfold our definitions above.
-attribute [grind] normalized hasNestedIf hasConstantIf hasRedundantIf disjoint vars eval
-
--- I'd prefer to use `TreeMap` here, but its `getElem?_insert` lemma is not useful.
-attribute [grind] Std.HashMap.getElem?_insert
+-- We tell `grind` to unfold our definitions above.
+attribute [local grind] normalized hasNestedIf hasConstantIf hasRedundantIf disjoint vars eval
 
 /--
 Custom size function for if-expressions, used for proving termination.
@@ -132,6 +130,10 @@ we are allowed to increase the size of the branches by one, and still be smaller
   | var _ => 1
   | .ite i t e => 2 * normSize i + max (normSize t) (normSize e) + 1
 
+-- TODO: ensure that we can use `TreeMap` here instead.
+-- Currently we rely on `HashMap.getElem?_insert`, but the corresponding `TreeMap` lemma is harder to use
+-- and will require some amount of `LawfulOrd` machinery.
+
 def normalize (assign : Std.HashMap Nat Bool) : IfExpr → IfExpr
   | lit b => lit b
   | var v =>
@@ -150,13 +152,37 @@ def normalize (assign : Std.HashMap Nat Bool) : IfExpr → IfExpr
     | some b => normalize assign (ite (lit b) t e)
   termination_by e => e.normSize
 
+-- TODO: add these to the library.
+attribute [grind] getElem?_eq_none_iff
+attribute [grind] Std.HashMap.contains_iff_mem
+
+example (m : Std.HashMap Nat Bool) (v w : Nat) (h : m.contains v) (_ : m[w]? = none) (_ : v = w) : False := by
+  grind
+
+-- TODO: ungrind `List.mem_append_left`/`right`?
+
 theorem normalize_spec (assign : Std.HashMap Nat Bool) (e : IfExpr) :
     (normalize assign e).normalized
     ∧ (∀ f, (normalize assign e).eval f = e.eval fun w => assign[w]?.getD (f w))
-    ∧ ∀ (v : Nat), v ∈ vars (normalize assign e) → assign[v]? = none := by
-  fun_induction normalize <;> grind (gen := 8) (splits := 8)
+    ∧ ∀ (v : Nat), v ∈ vars (normalize assign e) → ¬ v ∈ assign := by
+  fun_induction normalize <;> grind (gen := 7)
 
-/-
+-- We can also prove other variations, where we spell "`v` is not in `assign`"
+-- different ways, and `grind` doesn't mind.
+
+example (assign : Std.HashMap Nat Bool) (e : IfExpr) :
+    (normalize assign e).normalized
+    ∧ (∀ f, (normalize assign e).eval f = e.eval fun w => assign[w]?.getD (f w))
+    ∧ ∀ (v : Nat), v ∈ vars (normalize assign e) → assign.contains v = false := by
+  fun_induction normalize <;> grind (gen := 7)
+
+example (assign : Std.HashMap Nat Bool) (e : IfExpr) :
+    (normalize assign e).normalized
+    ∧ (∀ f, (normalize assign e).eval f = e.eval fun w => assign[w]?.getD (f w))
+    ∧ ∀ (v : Nat), v ∈ vars (normalize assign e) → assign[v]? = none := by
+  fun_induction normalize <;> grind (gen := 8)
+
+/--
 We recall the statement of the if-normalization problem.
 
 We want a function from if-expressions to if-expressions,