Using grind for Ordered Maps

In this section we'll build an example of a new data structure and basic API for it, illustrating the use of grind. The example will be derived from Rust's indexmap data structure.

IndexMap is intended as a replacement for HashMap (in particular, it has fast hash-based lookup), but allowing the user to maintain control of the order of the elements. We won't give a complete API, just set up some basic functions and theorems about them.

The two main functions we'll implement for now are insert and eraseSwap:

  • insert k v checks if k is already in the map. If so, it replaces the value with v, keeping k in the same position in the ordering. If it is not already in the map, insert adds (k, v) to the end of the map.

  • eraseSwap k removes the element with key k from the map, and swaps it with the last element of the map (or does nothing if k is not in the map). (This behavior may be surprising: this function exists because it is an efficient way to an erase element when you don't care about the order of the remaining elements. Another function, not implemented here, would preserve the order of the remaining elements, but at the cost of running in time proportional to the number of elements after the erased element.)

Our goals will be:

Complete encapsulation

The implementation of IndexMap is hidden from the users, and the theorems about the implementation details are private.

Use grind as much as possible

We'll prefer adding a private theorem and annotating it with @[grind] over writing a longer proof whenever practical.

Use auto-parameters as much as possible

Ideally, we don't even need to see the proofs; they should mostly be handled invisibly by grind.

The first step is to import the necessary data structures:

import Std.Data.HashMap

Skeleton

To begin with, we'll write out a skeleton of what we want to achieve, liberally using sorry as a placeholder for all proofs. In particular, this version makes no use of grind.

import Std.Data.HashMap open Std structure IndexMap (α : Type u) (β : Type v) [BEq α] [Hashable α] where indices : HashMap α Nat keys : Array α values : Array β size_keys : keys.size = values.size WF : ∀ (i : Nat) (a : α), keys[i]? = some a ↔ indices[a]? = some i namespace IndexMap variable {α : Type u} {β : Type v} [BEq α] [LawfulBEq α] [Hashable α] [LawfulHashable α] variable {m : IndexMap α β} {a : α} {b : β} {i : Nat} @[inline] def size (m : IndexMap α β) : Nat := m.values.size def declaration uses `sorry`declaration uses `sorry`emptyWithCapacity (capacity := 8) : IndexMap α β where indices := HashMap.emptyWithCapacity capacity keys := Array.emptyWithCapacity capacity values := Array.emptyWithCapacity capacity size_keys := sorry WF := sorry instance : EmptyCollection (IndexMap α β) where emptyCollection := emptyWithCapacity instance : Inhabited (IndexMap α β) where default := ∅ @[inline] def contains (m : IndexMap α β) (a : α) : Bool := m.indices.contains a instance : Membership α (IndexMap α β) where mem m a := a ∈ m.indices instance {m : IndexMap α β} {a : α} : Decidable (a ∈ m) := inferInstanceAs (Decidable (a ∈ m.indices)) @[inline] def findIdx? (m : IndexMap α β) (a : α) : Option Nat := m.indices[a]? @[inline] def findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : Nat := m.indices[a] @[inline] def getIdx? (m : IndexMap α β) (i : Nat) : Option β := m.values[i]? @[inline] def getIdx (m : IndexMap α β) (i : Nat) (h : i < m.size := by get_elem_tactic) : β := m.values[i] declaration uses `sorry`instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where getElem m a h := m.values[m.indices[a]]'(α:Type uβ:Type vinst✝³:BEq αinst✝²:LawfulBEq αinst✝¹:Hashable αinst✝:LawfulHashable αm✝:IndexMap α βa✝:αb:βi:Natm:IndexMap α βa:αh:a ∈ m⊢ m.indices[a] < m.values.size All goals completed! 🐙) getElem? m a := m.indices[a]?.bind (m.values[·]?) getElem! m a := m.indices[a]?.bind (m.values[·]?) |>.getD default declaration uses `sorry`instance : LawfulGetElem (IndexMap α β) α β (fun m a => a ∈ m) where getElem?_def := sorry getElem!_def := sorry @[inline] def declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`insert (m : IndexMap α β) (a : α) (b : β) : IndexMap α β := match h : m.indices[a]? with | some i => { indices := m.indices keys := m.keys.set i a sorry values := m.values.set i b sorry size_keys := sorry WF := sorry } | none => { indices := m.indices.insert a m.size keys := m.keys.push a values := m.values.push b size_keys := sorry WF := sorry } instance : Singleton (α × β) (IndexMap α β) := ⟨fun ⟨a, b⟩ => (∅ : IndexMap α β).insert a b⟩ instance : Insert (α × β) (IndexMap α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩ instance : LawfulSingleton (α × β) (IndexMap α β) := ⟨fun _ => rfl⟩ /-- Erase the key-value pair with the given key, moving the last pair into its place in the order. If the key is not present, the map is unchanged. -/ @[inline] def declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`declaration uses `sorry`eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := match h : m.indices[a]? with | some i => if w : i = m.size - 1 then { indices := m.indices.erase a keys := m.keys.pop values := m.values.pop size_keys := sorry WF := sorry } else let lastKey := m.keys.back sorry let lastValue := m.values.back sorry { indices := (m.indices.erase a).insert lastKey i keys := m.keys.pop.set i lastKey sorry values := m.values.pop.set i lastValue sorry size_keys := sorry WF := sorry } | none => m /-! ### Verification theorems -/ theorem declaration uses `sorry`getIdx_findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : m.getIdx (m.findIdx a h) sorry = m[a] := sorry theorem declaration uses `sorry`mem_insert (m : IndexMap α β) (a a' : α) (b : β) : a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m := α:Type uβ:Type vinst✝³:BEq αinst✝²:LawfulBEq αinst✝¹:Hashable αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:β⊢ a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m All goals completed! 🐙 theorem declaration uses `sorry`getElem_insert (m : IndexMap α β) (a a' : α) (b : β) (h : a' ∈ m.insert a b) : (m.insert a b)[a']'h = if h' : a' == a then b else m[a']'sorry := α:Type uβ:Type vinst✝³:BEq αinst✝²:LawfulBEq αinst✝¹:Hashable αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:βh:a' ∈ m.insert a b⊢ (m.insert a b)[a'] = if h' : (a' == a) = true then b else m[a'] All goals completed! 🐙 theorem declaration uses `sorry`findIdx_insert_self (m : IndexMap α β) (a : α) (b : β) : (m.insert a b).findIdx a sorry = if h : a ∈ m then m.findIdx a h else m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:LawfulBEq αinst✝¹:Hashable αinst✝:LawfulHashable αm:IndexMap α βa:αb:β⊢ (m.insert a b).findIdx a ⋯ = if h : a ∈ m then m.findIdx a h else m.size All goals completed! 🐙 end IndexMap

Header 2

Let's get started. We'll aspire to never writing a proof by hand, and the first step of that is to install auto-parameters for the size_keys and WF field, so we can omit these fields whenever grind can prove them. While we're modifying the definition of IndexMap itself, let's make all the fields private, since we're planning on having complete encapsulation.

open Std structure IndexMap (α : Type u) (β : Type v) [BEq α] [Hashable α] where private indices : HashMap α Nat private keys : Array α private values : Array β private size_keys' : keys.size = values.size := by grind private WF : ∀ (i : Nat) (a : α), keys[i]? = some a ↔ indices[a]? = some i := by grind

For the rest of this tutorial, the following namespace and variable declarations are in effect:

namespace IndexMap variable {α : Type u} {β : Type v} [BEq α] [Hashable α] variable {m : IndexMap α β} {a : α} {b : β} {i : Nat}

Let's give grind access to the definition of size, and size_keys private field:

@[inline] def size (m : IndexMap α β) : Nat := m.values.size @[local grind =] private theorem size_keys : m.keys.size = m.size := m.size_keys' @[local grind =] private theorem size_values : m.values.size = m.size := rfl

Our first sorrys in the draft version are the size_keys and WF fields in our construction of def emptyWithCapacity. Surely these are trivial, and solvable by grind, so we simply delete those fields:

def emptyWithCapacity (capacity := 8) : IndexMap α β where indices := HashMap.emptyWithCapacity capacity keys := Array.emptyWithCapacity capacity values := Array.emptyWithCapacity capacity

Our next task is to deal with the sorry in our construction of the original GetElem? instance:

declaration uses `sorry`instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where getElem m a h := m.values[m.indices[a]]'(α:Type uβ:Type vinst✝³:BEq αinst✝²:LawfulBEq αinst✝¹:Hashable αinst✝:LawfulHashable αm✝:IndexMap α βa✝:αb:βi:Natm:IndexMap α βa:αh:a ∈ m⊢ m.indices[a] < m.values.size All goals completed! 🐙) getElem? m a := m.indices[a]?.bind (m.values[·]?) getElem! m a := m.indices[a]?.bind (m.values[·]?) |>.getD default

The goal at this sorry is

m : IndexMap α β
a : α
h : a ∈ m
⊢ m.indices[a] < m.values.size

Let's try proving this as a stand-alone theorem, via grind, and see where grind gets stuck. Because we've added grind annotations for size and size_keys already, we can safely reformulate the goal as:

theorem getElem_indices_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : m.indices[a] < m.size := α:Type uβ:Type vinst✝¹:BEq αinst✝:Hashable αm:IndexMap α βa:αh:a ∈ m⊢ m.indices[a] < m.size `grind` failed α:Type uβ:Type vinst:BEq αinst_1:Hashable αm:IndexMap α βa:αh:a ∈ mh_1:m.size ≤ m.indices[a]⊢ False
[grind] Goal diagnostics
  • [facts] Asserted facts
    • [prop] a ∈ m
    • [prop] m.size ≤ m.indices[a]
  • [eqc] True propositions
    • [prop] m.size ≤ m.indices[a]
    • [prop] a ∈ m
  • [eqc] Equivalence classes
    • [eqc] {Membership.mem, fun m a => a ∈ m}
  • [ematch] E-matching patterns
    • [thm] getElem?_pos: [@getElem #8 #7 #6 #5 _ #2 #1 _]
    • [thm] HashMap.contains_iff_mem: [@Membership.mem #5 (HashMap _ #4 #3 #2) _ #1 #0]
  • [cutsat] Assignment satisfying linear constraints
    • [assign] m.size := 0
    • [assign] m.indices[a] := 0
All goals completed! 🐙

This fails, and looking at the Goal diagnostics section of the message from grind we see that it hasn't done much:

`grind` failed
α:Type uβ:Type vinst:BEq αinst_1:Hashable αm:IndexMap α βa:αh:a ∈ mh_1:m.size ≤ m.indices[a]⊢ False

[grind] Goal diagnostics
  • [facts] Asserted facts
    • [prop] a ∈ m
    • [prop] m.size ≤ m.indices[a]
  • [eqc] True propositions
    • [prop] m.size ≤ m.indices[a]
    • [prop] a ∈ m
  • [eqc] Equivalence classes
    • [eqc] {Membership.mem, fun m a => a ∈ m}
  • [ematch] E-matching patterns
    • [thm] getElem?_pos: [@getElem #8 #7 #6 #5 _ #2 #1 _]
    • [thm] HashMap.contains_iff_mem: [@Membership.mem #5 (HashMap _ #4 #3 #2) _ #1 #0]
  • [cutsat] Assignment satisfying linear constraints
    • [assign] m.size := 0
    • [assign] m.indices[a] := 0

An immediate problem we can see here is that grind does not yet know that a ∈ m is the same as a ∈ m.indices. Let's add this fact:

@[local grind _=_] private theorem mem_indices {m : IndexMap α β} {a : α} : a ∈ m.indices ↔ a ∈ m := Iff.rfl

However this proof is going to work, we know the following:

  • It must use the well-formedness condition of the map.

  • It can't do so without relating m.indices[a] and m.indices[a]? (because the later is what appears in the well-formedness condition).

  • The expected relationship there doesn't even hold unless the map m.indices satisfies LawfulGetElem, for which we need instances of LawfulBEq α and LawfulHashable α.

Let's configure things so that those are available:

variable [LawfulBEq α] [LawfulHashable α] attribute [local grind _=_] IndexMap.WF

and then give grind one manual hint, to relate m.indices[a] and m.indices[a]?:

private theorem getElem_indices_lt {h : a ∈ m} : m.indices[a] < m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αinst✝¹:LawfulBEq αinst✝:LawfulHashable αh:a ∈ m⊢ m.indices[a] < m.size have : m.indices[a]? = some m.indices[a] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αinst✝¹:LawfulBEq αinst✝:LawfulHashable αh:a ∈ m⊢ m.indices[a] < m.size All goals completed! 🐙 All goals completed! 🐙

With that theorem proved, we want to make it accessible to grind. We could either add @[grind] before the theorem statement, or write attribute [local grind] getElem_indices_lt after the theorem statement. These will use grind's built-in heuristics for deciding a pattern to match the theorem on.

In this case, let's see which patterns the grind attribute generates:

attribute [local Try these: [apply] [grind .] for pattern: [@LE.le `[Nat] `[instLENat] ((@getElem (HashMap #8 `[Nat] #6 #5) _ `[Nat] _ _ (@indices _ #7 _ _ #4) #3 #0) + 1) (@size _ _ _ _ #4)] [apply] [grind →] for pattern: [LawfulBEq #8 #6, LawfulHashable _ _ #5, @Membership.mem _ (IndexMap _ #7 _ _) _ #4 #3]grind] getElem_indices_lt
Try these:
  [apply] [grind
    .] for pattern: [@LE.le `[Nat] `[instLENat] ((@getElem (HashMap #8 `[Nat] #6 #5) _ `[Nat] _ _ (@indices _ #7 _ _ #4) #3 #0) + 1) (@size _ _ _ _ #4)]
  [apply] [grind →] for pattern: [LawfulBEq #8 #6, LawfulHashable _ _ #5, @Membership.mem _ (IndexMap _ #7 _ _) _ #4 #3]

These patterns are not useful. The first is matching on the entire conclusion of the theorem (in fact, a normalized version of it, in which x < y has been replaced by x + 1 ≤ y). The second is too general: it will match any term that includes the theorem's assumptions, ignoring the conclusion.

We want something more general than the entire conclusion, the conclusion should not be ignored. We'd like this theorem to fire whenever grind sees m.indices[a], and so instead of using the attribute we write a custom pattern:

grind_pattern getElem_indices_lt => m.indices[a]

The Lean standard library uses the get_elem_tactic tactic as an auto-parameter for the xs[i] notation (which desugars to GetElem.getElem xs i h, with the proof h generated by get_elem_tactic). We'd like to not only have grind fill in these proofs, but even to be able to omit these proofs. To achieve this, we add the line

macro_rules | `(tactic| get_elem_tactic_extensible) => `(tactic| grind)

(In later versions of Lean this may be part of the built-in behavior.)

We can now return to constructing our GetElem? instance. In order to use the well-formedness condition, grind must be able to unfold size:

attribute [local grind] size

The local modifier restricts this unfolding to the current file. With this in place, we can simply write:

instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where getElem m a h := m.values[m.indices[a]] getElem? m a := m.indices[a]?.bind (fun i => (m.values[i]?)) getElem! m a := m.indices[a]?.bind (fun i => (m.values[i]?)) |>.getD default

with neither any sorrys, nor any explicitly written proofs.

Next, we want to expose the content of these definitions, but only locally in this file:

@[local grind =] private theorem getElem_def (m : IndexMap α β) (a : α) (h : a ∈ m) : m[a] = m.values[m.indices[a]'h] := rfl @[local grind =] private theorem getElem?_def (m : IndexMap α β) (a : α) : m[a]? = m.indices[a]?.bind (fun i => (m.values[i]?)) := rfl @[local grind =] private theorem getElem!_def [Inhabited β] (m : IndexMap α β) (a : α) : m[a]! = (m.indices[a]?.bind (m.values[·]?)).getD default := rfl

Again we're using the @[local grind =] private theorem pattern to hide these implementation details, but allow grind to see these facts locally.

Next, we want to prove the LawfulGetElem instance, and hope that grind can fill in the proofs:

instance : LawfulGetElem (IndexMap α β) α β (fun m a => a ∈ m) where getElem?_def := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αb:βi:Natinst✝¹:LawfulBEq αinst✝:LawfulHashable α⊢ ∀ (c : IndexMap α β) (i : α) [inst : Decidable (i ∈ c)], c[i]? = if h : i ∈ c then some c[i] else none All goals completed! 🐙 getElem!_def := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αb:βi:Natinst✝¹:LawfulBEq αinst✝:LawfulHashable α⊢ ∀ [inst : Inhabited β] (c : IndexMap α β) (i : α), c[i]! = match c[i]? with | some e => e | none => default All goals completed! 🐙

Success!

Let's press onward, and see if we can define insert without having to write any proofs:

@[inline] def insert (m : IndexMap α β) (a : α) (b : β) : IndexMap α β := match h : m.indices[a]? with | some i => { indices := m.indices keys := m.keys.set i a values := m.values.set i b } | none => { indices := m.indices.insert a m.size keys := m.keys.push a values := m.values.push b }

In both branches, grind is automatically proving both the size_keys' and WF fields! Note also in the first branch the set calls m.keys.set i a and m.values.set i b are having their “in-bounds” obligations automatically filled in by grind via the get_elem_tactic auto-parameter.

In eraseSwap, everything goes through cleanly now, with no manual proofs:

@[inline] def eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := match h : m.indices[a]? with | some i => if w : i = m.size - 1 then { indices := m.indices.erase a keys := m.keys.pop values := m.values.pop } else let lastKey := m.keys.back let lastValue := m.values.back { indices := (m.indices.erase a).insert lastKey i keys := m.keys.pop.set i lastKey values := m.values.pop.set i lastValue } | none => m

Finally we turn to the verification theorems about the basic operations, relating getIdx, findIdx, and insert. The proofs all go through effortlessly using grind with the +locals modifier (which tells grind to unfold local definitions):

/-! ### Verification theorems (not exhaustive) -/ @[grind =] theorem mem_insert (m : IndexMap α β) (a a' : α) (b : β) : a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:β⊢ a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m All goals completed! 🐙 @[grind =] theorem getElem_insert (m : IndexMap α β) (a a' : α) (b : β) (h : a' ∈ m.insert a b) : (m.insert a b)[a'] = if h' : a' == a then b else m[a'] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:βh:a' ∈ m.insert a b⊢ (m.insert a b)[a'] = if h' : (a' == a) = true then b else m[a'] All goals completed! 🐙 theorem findIdx_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : m.findIdx a h < m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αh:a ∈ m⊢ m.findIdx a h < m.size All goals completed! 🐙 grind_pattern findIdx_lt => m.findIdx a h @[grind =] theorem findIdx_insert_self (m : IndexMap α β) (a : α) (b : β) : (m.insert a b).findIdx a = if h : a ∈ m then m.findIdx a else m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αb:β⊢ (m.insert a b).findIdx a ⋯ = if h : a ∈ m then m.findIdx a h else m.size All goals completed! 🐙 @[grind =] theorem findIdx?_eq (m : IndexMap α β) (a : α) : m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:α⊢ m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none All goals completed! 🐙 @[grind =] theorem getIdx_findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : m.getIdx (m.findIdx a) = m[a] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αh:a ∈ m⊢ m.getIdx (m.findIdx a h) ⋯ = m[a] All goals completed! 🐙 omit [LawfulBEq α] [LawfulHashable α] in @[grind =] theorem getIdx?_eq (m : IndexMap α β) (i : Nat) : m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none := α:Type uβ:Type vinst✝¹:BEq αinst✝:Hashable αm:IndexMap α βi:Nat⊢ m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none All goals completed! 🐙 private theorem getElem_keys_mem {m : IndexMap α β} {i : Nat} (h : i < m.size) : m.keys[i] ∈ m := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βi:Nath:i < m.size⊢ m.keys[i] ∈ m have : m.indices[m.keys[i]]? = some i := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βi:Nath:i < m.size⊢ m.keys[i] ∈ m All goals completed! 🐙 All goals completed! 🐙 local grind_pattern getElem_keys_mem => m.keys[i] theorem getElem?_eraseSwap (m : IndexMap α β) (a a' : α) : (m.eraseSwap a)[a']? = if a' == a then none else m[a']? := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':α⊢ (m.eraseSwap a)[a']? = if (a' == a) = true then none else m[a']? All goals completed! 🐙 @[grind =] theorem mem_eraseSwap (m : IndexMap α β) (a a' : α) : a' ∈ m.eraseSwap a ↔ a' ≠ a ∧ a' ∈ m := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':α⊢ a' ∈ m.eraseSwap a ↔ a' ≠ a ∧ a' ∈ m All goals completed! 🐙 theorem getElem_eraseSwap (m : IndexMap α β) (a a' : α) (h : a' ∈ m.eraseSwap a) : (m.eraseSwap a)[a'] = m[a'] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':αh:a' ∈ m.eraseSwap a⊢ (m.eraseSwap a)[a'] = m[a'] All goals completed! 🐙

Note that these are part of the public API of IndexMap, so we need to mark them as @[grind], so that users without our internal local grind annotations can still use them in grind proofs.

Putting this all together, our prototype API has reached the following state:

local macro_rules | `(tactic| get_elem_tactic_extensible) => `(tactic| grind) open Std structure IndexMap (α : Type u) (β : Type v) [BEq α] [Hashable α] where private indices : HashMap α Nat private keys : Array α private values : Array β private size_keys' : keys.size = values.size := by grind private WF : ∀ (i : Nat) (a : α), keys[i]? = some a ↔ indices[a]? = some i := by grind namespace IndexMap variable {α : Type u} {β : Type v} [BEq α] [Hashable α] variable {m : IndexMap α β} {a : α} {b : β} {i : Nat} @[inline] def size (m : IndexMap α β) : Nat := m.values.size @[local grind =] private theorem size_keys : m.keys.size = m.size := m.size_keys' @[local grind =] private theorem size_values : m.values.size = m.size := rfl def emptyWithCapacity (capacity := 8) : IndexMap α β where indices := HashMap.emptyWithCapacity capacity keys := Array.emptyWithCapacity capacity values := Array.emptyWithCapacity capacity instance : EmptyCollection (IndexMap α β) where emptyCollection := emptyWithCapacity instance : Inhabited (IndexMap α β) where default := ∅ @[inline] def contains (m : IndexMap α β) (a : α) : Bool := m.indices.contains a instance : Membership α (IndexMap α β) where mem m a := a ∈ m.indices instance {m : IndexMap α β} {a : α} : Decidable (a ∈ m) := inferInstanceAs (Decidable (a ∈ m.indices)) @[local grind _=_] private theorem mem_indices {m : IndexMap α β} {a : α} : a ∈ m.indices ↔ a ∈ m := Iff.rfl @[inline] def findIdx? (m : IndexMap α β) (a : α) : Option Nat := m.indices[a]? @[inline] def findIdx (m : IndexMap α β) (a : α) (h : a ∈ m := by get_elem_tactic) : Nat := m.indices[a] @[inline] def getIdx? (m : IndexMap α β) (i : Nat) : Option β := m.values[i]? @[inline] def getIdx (m : IndexMap α β) (i : Nat) (h : i < m.size := by get_elem_tactic) : β := m.values[i] variable [LawfulBEq α] [LawfulHashable α] attribute [local grind _=_] IndexMap.WF private theorem getElem_indices_lt {h : a ∈ m} : m.indices[a] < m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αinst✝¹:LawfulBEq αinst✝:LawfulHashable αh:a ∈ m⊢ m.indices[a] < m.size have : m.indices[a]? = some m.indices[a] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αinst✝¹:LawfulBEq αinst✝:LawfulHashable αh:a ∈ m⊢ m.indices[a] < m.size All goals completed! 🐙 All goals completed! 🐙 grind_pattern getElem_indices_lt => m.indices[a] instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where getElem m a h := m.values[m.indices[a]] getElem? m a := m.indices[a]?.bind (fun i => (m.values[i]?)) getElem! m a := m.indices[a]?.bind (fun i => (m.values[i]?)) |>.getD default @[local grind =] private theorem getElem_def (m : IndexMap α β) (a : α) (h : a ∈ m) : m[a] = m.values[m.indices[a]'h] := rfl @[local grind =] private theorem getElem?_def (m : IndexMap α β) (a : α) : m[a]? = m.indices[a]?.bind (fun i => (m.values[i]?)) := rfl @[local grind =] private theorem getElem!_def [Inhabited β] (m : IndexMap α β) (a : α) : m[a]! = (m.indices[a]?.bind (m.values[·]?)).getD default := rfl instance : LawfulGetElem (IndexMap α β) α β (fun m a => a ∈ m) where getElem?_def := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αb:βi:Natinst✝¹:LawfulBEq αinst✝:LawfulHashable α⊢ ∀ (c : IndexMap α β) (i : α) [inst : Decidable (i ∈ c)], c[i]? = if h : i ∈ c then some c[i] else none All goals completed! 🐙 getElem!_def := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βa:αb:βi:Natinst✝¹:LawfulBEq αinst✝:LawfulHashable α⊢ ∀ [inst : Inhabited β] (c : IndexMap α β) (i : α), c[i]! = match c[i]? with | some e => e | none => default All goals completed! 🐙 @[inline] def insert (m : IndexMap α β) (a : α) (b : β) : IndexMap α β := match h : m.indices[a]? with | some i => { indices := m.indices keys := m.keys.set i a values := m.values.set i b } | none => { indices := m.indices.insert a m.size keys := m.keys.push a values := m.values.push b } instance : Singleton (α × β) (IndexMap α β) := ⟨fun ⟨a, b⟩ => (∅ : IndexMap α β).insert a b⟩ instance : Insert (α × β) (IndexMap α β) := ⟨fun ⟨a, b⟩ s => s.insert a b⟩ instance : LawfulSingleton (α × β) (IndexMap α β) := ⟨fun _ => rfl⟩ @[local grind .] private theorem WF' (i : Nat) (a : α) (h₁ : i < m.keys.size) (h₂ : a ∈ m) : m.keys[i] = a ↔ m.indices[a] = i := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βinst✝¹:LawfulBEq αinst✝:LawfulHashable αi:Nata:αh₁:i < m.keys.sizeh₂:a ∈ m⊢ m.keys[i] = a ↔ m.indices[a] = i α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αm:IndexMap α βinst✝¹:LawfulBEq αinst✝:LawfulHashable αi:Nata:αh₁:i < m.keys.sizeh₂:a ∈ mthis:m.keys[i]? = some a ↔ m.indices[a]? = some i := WF m i a⊢ m.keys[i] = a ↔ m.indices[a] = i All goals completed! 🐙 /-- Erase the key-value pair with the given key, moving the last pair into its place in the order. If the key is not present, the map is unchanged. -/ @[inline] def eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := match h : m.indices[a]? with | some i => if w : i = m.size - 1 then { indices := m.indices.erase a keys := m.keys.pop values := m.values.pop } else let lastKey := m.keys.back let lastValue := m.values.back { indices := (m.indices.erase a).insert lastKey i keys := m.keys.pop.set i lastKey values := m.values.pop.set i lastValue } | none => m /-! ### Verification theorems (not exhaustive) -/ @[grind =] theorem mem_insert (m : IndexMap α β) (a a' : α) (b : β) : a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:β⊢ a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m All goals completed! 🐙 @[grind =] theorem getElem_insert (m : IndexMap α β) (a a' : α) (b : β) (h : a' ∈ m.insert a b) : (m.insert a b)[a'] = if h' : a' == a then b else m[a'] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αa':αb:βh:a' ∈ m.insert a b⊢ (m.insert a b)[a'] = if h' : (a' == a) = true then b else m[a'] All goals completed! 🐙 theorem findIdx_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : m.findIdx a h < m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αh:a ∈ m⊢ m.findIdx a h < m.size All goals completed! 🐙 grind_pattern findIdx_lt => m.findIdx a h @[grind =] theorem findIdx_insert_self (m : IndexMap α β) (a : α) (b : β) : (m.insert a b).findIdx a = if h : a ∈ m then m.findIdx a else m.size := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αb:β⊢ (m.insert a b).findIdx a ⋯ = if h : a ∈ m then m.findIdx a h else m.size All goals completed! 🐙 @[grind =] theorem findIdx?_eq (m : IndexMap α β) (a : α) : m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:α⊢ m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none All goals completed! 🐙 @[grind =] theorem getIdx_findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : m.getIdx (m.findIdx a) = m[a] := α:Type uβ:Type vinst✝³:BEq αinst✝²:Hashable αinst✝¹:LawfulBEq αinst✝:LawfulHashable αm:IndexMap α βa:αh:a ∈ m⊢ m.getIdx (m.findIdx a h) ⋯ = m[a] All goals completed! 🐙 omit [LawfulBEq α] [LawfulHashable α] in @[grind =] theorem getIdx?_eq (m : IndexMap α β) (i : Nat) : m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none := α:Type uβ:Type vinst✝¹:BEq αinst✝:Hashable αm:IndexMap α βi:Nat⊢ m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none All goals completed! 🐙 end IndexMap

We've now also added verification theorems for eraseSwap operations; the interested reader is encouraged to explore further, and perhaps even releasing a complete IndexMap library!

Summarizing the design principles discussed above about encapsulation:

  • the fields of IndexMap are all private, as these are implementation details.

  • the theorems about these fields are all private, and marked as @[grind], rather than @[grind], as they won't be needed after we've set up the API.

  • the verification theorems are both marked as @[grind], and proved by grind: the annotation is necessary because we want grind to be able to prove these facts even once we're outside the current module, and the @[grind] theorems are no longer available.