|
| 1 | +import Std.Data.HashMap |
| 2 | + |
| 3 | +open Std |
| 4 | + |
| 5 | +structure IndexMap |
| 6 | + (α : Type u) (β : Type v) [BEq α] [Hashable α] where |
| 7 | + private indices : HashMap α Nat |
| 8 | + private keys : Array α |
| 9 | + private values : Array β |
| 10 | + private size_keys' : keys.size = values.size := by grind |
| 11 | + private WF : ∀ (i : Nat) (a : α), |
| 12 | + keys[i]? = some a ↔ indices[a]? = some i := by grind |
| 13 | + |
| 14 | +namespace IndexMap |
| 15 | + |
| 16 | +variable {α : Type u} {β : Type v} [BEq α] [Hashable α] |
| 17 | +variable {m : IndexMap α β} {a : α} {b : β} {i : Nat} |
| 18 | + |
| 19 | +@[inline] def size (m : IndexMap α β) : Nat := |
| 20 | + m.values.size |
| 21 | + |
| 22 | +@[local grind =] private theorem size_keys : m.keys.size = m.size := |
| 23 | + m.size_keys' |
| 24 | + |
| 25 | +@[local grind =] private theorem size_values : m.values.size = m.size := rfl |
| 26 | + |
| 27 | +def emptyWithCapacity (capacity := 8) : IndexMap α β where |
| 28 | + indices := HashMap.emptyWithCapacity capacity |
| 29 | + keys := Array.emptyWithCapacity capacity |
| 30 | + values := Array.emptyWithCapacity capacity |
| 31 | + |
| 32 | +@[inline] def contains (m : IndexMap α β) |
| 33 | + (a : α) : Bool := |
| 34 | + m.indices.contains a |
| 35 | + |
| 36 | +instance : Membership α (IndexMap α β) where |
| 37 | + mem m a := a ∈ m.indices |
| 38 | + |
| 39 | +instance {m : IndexMap α β} {a : α} : Decidable (a ∈ m) := |
| 40 | + inferInstanceAs (Decidable (a ∈ m.indices)) |
| 41 | + |
| 42 | +-- theorem getElem_indices_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 43 | +-- m.indices[a] < m.size := by |
| 44 | +-- grind |
| 45 | + |
| 46 | +@[local grind _=_] private theorem mem_indices |
| 47 | + {m : IndexMap α β} {a : α} : |
| 48 | + a ∈ m.indices ↔ a ∈ m := Iff.rfl |
| 49 | + |
| 50 | +variable [LawfulBEq α] [LawfulHashable α] |
| 51 | + |
| 52 | +attribute [local grind _=_] IndexMap.WF |
| 53 | + |
| 54 | +private theorem getElem_indices_lt {h : a ∈ m} : m.indices[a] < m.size := by |
| 55 | + have : m.indices[a]? = some m.indices[a] := by grind |
| 56 | + grind |
| 57 | + |
| 58 | +attribute [local grind] getElem_indices_lt |
| 59 | + |
| 60 | +grind_pattern getElem_indices_lt => m.indices[a] |
| 61 | + |
| 62 | +macro_rules | `(tactic| get_elem_tactic_extensible) => `(tactic| grind) |
| 63 | + |
| 64 | +attribute [local grind] size |
| 65 | + |
| 66 | +instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where |
| 67 | + getElem m a h := |
| 68 | + m.values[m.indices[a]] |
| 69 | + getElem? m a := |
| 70 | + m.indices[a]?.bind (fun i => (m.values[i]?)) |
| 71 | + getElem! m a := |
| 72 | + m.indices[a]?.bind (fun i => (m.values[i]?)) |>.getD default |
| 73 | + |
| 74 | +@[local grind =] private theorem getElem_def |
| 75 | + (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 76 | + m[a] = m.values[m.indices[a]'h] := |
| 77 | + rfl |
| 78 | +@[local grind =] private theorem getElem?_def |
| 79 | + (m : IndexMap α β) (a : α) : |
| 80 | + m[a]? = m.indices[a]?.bind (fun i => (m.values[i]?)) := |
| 81 | + rfl |
| 82 | +@[local grind =] private theorem getElem!_def |
| 83 | + [Inhabited β] (m : IndexMap α β) (a : α) : |
| 84 | + m[a]! = (m.indices[a]?.bind (m.values[·]?)).getD default := |
| 85 | + rfl |
| 86 | + |
| 87 | +instance : LawfulGetElem (IndexMap α β) α β (fun m a => a ∈ m) where |
| 88 | + getElem?_def := by grind |
| 89 | + getElem!_def := by grind |
| 90 | + |
| 91 | +@[inline] def insert (m : IndexMap α β) (a : α) (b : β) : IndexMap α β := |
| 92 | + match h : m.indices[a]? with |
| 93 | + | some i => |
| 94 | + { indices := m.indices |
| 95 | + keys := m.keys.set i a |
| 96 | + values := m.values.set i b } |
| 97 | + | none => |
| 98 | + { indices := m.indices.insert a m.size |
| 99 | + keys := m.keys.push a |
| 100 | + values := m.values.push b } |
| 101 | + |
| 102 | +-- @[inline] def eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := |
| 103 | +-- match h : m.indices[a]? with |
| 104 | +-- | some i => |
| 105 | +-- if w : i = m.size - 1 then |
| 106 | +-- { indices := m.indices.erase a |
| 107 | +-- keys := m.keys.pop |
| 108 | +-- values := m.values.pop } |
| 109 | +-- else |
| 110 | +-- let lastKey := m.keys.back |
| 111 | +-- let lastValue := m.values.back |
| 112 | +-- { indices := (m.indices.erase a).insert lastKey i |
| 113 | +-- keys := m.keys.pop.set i lastKey |
| 114 | +-- values := m.values.pop.set i lastValue } |
| 115 | +-- | none => m |
| 116 | + |
| 117 | +@[local grind .] |
| 118 | +private theorem WF' (i : Nat) (a : α) (h₁ : i < m.keys.size) (h₂ : a ∈ m) : |
| 119 | + m.keys[i] = a ↔ m.indices[a] = i := by |
| 120 | + have := m.WF i a |
| 121 | + grind |
| 122 | + |
| 123 | +example {m : IndexMap α β} {a : α} {h : a ∈ m} : |
| 124 | + m.keys[m.indices[a]'h] = a := by grind |
| 125 | + |
| 126 | +@[inline] def eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := |
| 127 | + match h : m.indices[a]? with |
| 128 | + | some i => |
| 129 | + if w : i = m.size - 1 then |
| 130 | + { indices := m.indices.erase a |
| 131 | + keys := m.keys.pop |
| 132 | + values := m.values.pop } |
| 133 | + else |
| 134 | + let lastKey := m.keys.back |
| 135 | + let lastValue := m.values.back |
| 136 | + { indices := (m.indices.erase a).insert lastKey i |
| 137 | + keys := m.keys.pop.set i lastKey |
| 138 | + values := m.values.pop.set i lastValue } |
| 139 | + | none => m |
| 140 | + |
| 141 | +@[inline] def findIdx? (m : IndexMap α β) (a : α) : Option Nat := |
| 142 | + m.indices[a]? |
| 143 | + |
| 144 | +@[inline] def findIdx (m : IndexMap α β) (a : α) |
| 145 | + (h : a ∈ m := by get_elem_tactic) : Nat := |
| 146 | + m.indices[a] |
| 147 | + |
| 148 | +@[inline] def getIdx? (m : IndexMap α β) (i : Nat) : Option β := |
| 149 | + m.values[i]? |
| 150 | + |
| 151 | +@[inline] def getIdx (m : IndexMap α β) (i : Nat) |
| 152 | + (h : i < m.size := by get_elem_tactic) : β := |
| 153 | + m.values[i] |
| 154 | + |
| 155 | +/-! ### Verification theorems (not exhaustive) -/ |
| 156 | + |
| 157 | +@[grind =] |
| 158 | +theorem mem_insert (m : IndexMap α β) (a a' : α) (b : β) : |
| 159 | + a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m := by |
| 160 | + grind +locals |
| 161 | + |
| 162 | +@[grind =] |
| 163 | +theorem getElem_insert (m : IndexMap α β) (a a' : α) (b : β) (h : a' ∈ m.insert a b) : |
| 164 | + (m.insert a b)[a'] = if h' : a' == a then b else m[a'] := by |
| 165 | + grind +locals |
| 166 | + |
| 167 | +theorem findIdx_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 168 | + m.findIdx a h < m.size := by |
| 169 | + grind +locals |
| 170 | + |
| 171 | +grind_pattern findIdx_lt => m.findIdx a h |
| 172 | + |
| 173 | +@[grind =] |
| 174 | +theorem findIdx_insert_self (m : IndexMap α β) (a : α) (b : β) : |
| 175 | + (m.insert a b).findIdx a = if h : a ∈ m then m.findIdx a else m.size := by |
| 176 | + grind +locals |
| 177 | + |
| 178 | +@[grind =] |
| 179 | +theorem findIdx?_eq (m : IndexMap α β) (a : α) : |
| 180 | + m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none := by |
| 181 | + grind +locals |
| 182 | + |
| 183 | +@[grind =] |
| 184 | +theorem getIdx_findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 185 | + m.getIdx (m.findIdx a) = m[a] := by grind +locals |
| 186 | + |
| 187 | +omit [LawfulBEq α] [LawfulHashable α] in |
| 188 | +@[grind =] |
| 189 | +theorem getIdx?_eq (m : IndexMap α β) (i : Nat) : |
| 190 | + m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none := by |
| 191 | + grind +locals |
| 192 | + |
| 193 | +private theorem getElem_keys_mem {m : IndexMap α β} {i : Nat} (h : i < m.size) : |
| 194 | + m.keys[i] ∈ m := by |
| 195 | + have : m.indices[m.keys[i]]? = some i := by grind |
| 196 | + grind |
| 197 | + |
| 198 | +local grind_pattern getElem_keys_mem => m.keys[i] |
| 199 | + |
| 200 | +theorem getElem?_eraseSwap (m : IndexMap α β) (a a' : α) : |
| 201 | + (m.eraseSwap a)[a']? = if a' == a then none else m[a']? := by |
| 202 | + grind +locals |
| 203 | + |
| 204 | +@[grind =] |
| 205 | +theorem mem_eraseSwap (m : IndexMap α β) (a a' : α) : |
| 206 | + a' ∈ m.eraseSwap a ↔ a' ≠ a ∧ a' ∈ m := by |
| 207 | + grind +locals |
| 208 | + |
| 209 | +theorem getElem_eraseSwap (m : IndexMap α β) (a a' : α) (h : a' ∈ m.eraseSwap a) : |
| 210 | + (m.eraseSwap a)[a'] = m[a'] := by |
| 211 | + grind +locals |
| 212 | + |
| 213 | +local macro_rules | `(tactic| get_elem_tactic_extensible) => `(tactic| grind) |
| 214 | + |
| 215 | +open Std |
| 216 | + |
| 217 | +structure IndexMap |
| 218 | + (α : Type u) (β : Type v) [BEq α] [Hashable α] where |
| 219 | + private indices : HashMap α Nat |
| 220 | + private keys : Array α |
| 221 | + private values : Array β |
| 222 | + private size_keys' : keys.size = values.size := by grind |
| 223 | + private WF : ∀ (i : Nat) (a : α), |
| 224 | + keys[i]? = some a ↔ indices[a]? = some i := by grind |
| 225 | + |
| 226 | +namespace IndexMap |
| 227 | + |
| 228 | +variable {α : Type u} {β : Type v} [BEq α] [Hashable α] |
| 229 | +variable {m : IndexMap α β} {a : α} {b : β} {i : Nat} |
| 230 | + |
| 231 | +@[inline] def size (m : IndexMap α β) : Nat := |
| 232 | + m.values.size |
| 233 | + |
| 234 | +@[local grind =] private theorem size_keys : m.keys.size = m.size := |
| 235 | + m.size_keys' |
| 236 | + |
| 237 | +@[local grind =] private theorem size_values : m.values.size = m.size := rfl |
| 238 | + |
| 239 | +def emptyWithCapacity (capacity := 8) : IndexMap α β where |
| 240 | + indices := HashMap.emptyWithCapacity capacity |
| 241 | + keys := Array.emptyWithCapacity capacity |
| 242 | + values := Array.emptyWithCapacity capacity |
| 243 | + |
| 244 | +instance : EmptyCollection (IndexMap α β) where |
| 245 | + emptyCollection := emptyWithCapacity |
| 246 | + |
| 247 | +instance : Inhabited (IndexMap α β) where |
| 248 | + default := ∅ |
| 249 | + |
| 250 | +@[inline] def contains (m : IndexMap α β) (a : α) : Bool := |
| 251 | + m.indices.contains a |
| 252 | + |
| 253 | +instance : Membership α (IndexMap α β) where |
| 254 | + mem m a := a ∈ m.indices |
| 255 | + |
| 256 | +instance {m : IndexMap α β} {a : α} : Decidable (a ∈ m) := |
| 257 | + inferInstanceAs (Decidable (a ∈ m.indices)) |
| 258 | + |
| 259 | +@[local grind _=_] private theorem mem_indices |
| 260 | + {m : IndexMap α β} {a : α} : |
| 261 | + a ∈ m.indices ↔ a ∈ m := Iff.rfl |
| 262 | + |
| 263 | +@[inline] def findIdx? (m : IndexMap α β) (a : α) : Option Nat := |
| 264 | + m.indices[a]? |
| 265 | + |
| 266 | +@[inline] def findIdx (m : IndexMap α β) (a : α) |
| 267 | + (h : a ∈ m := by get_elem_tactic) : Nat := |
| 268 | + m.indices[a] |
| 269 | + |
| 270 | +@[inline] def getIdx? (m : IndexMap α β) (i : Nat) : Option β := |
| 271 | + m.values[i]? |
| 272 | + |
| 273 | +@[inline] def getIdx (m : IndexMap α β) (i : Nat) |
| 274 | + (h : i < m.size := by get_elem_tactic) : β := |
| 275 | + m.values[i] |
| 276 | + |
| 277 | +variable [LawfulBEq α] [LawfulHashable α] |
| 278 | + |
| 279 | +attribute [local grind _=_] IndexMap.WF |
| 280 | + |
| 281 | +private theorem getElem_indices_lt |
| 282 | + {h : a ∈ m} : m.indices[a] < m.size := by |
| 283 | + have : m.indices[a]? = some m.indices[a] := by grind |
| 284 | + grind |
| 285 | + |
| 286 | +grind_pattern getElem_indices_lt => m.indices[a] |
| 287 | + |
| 288 | +instance : GetElem? (IndexMap α β) α β (fun m a => a ∈ m) where |
| 289 | + getElem m a h := |
| 290 | + m.values[m.indices[a]] |
| 291 | + getElem? m a := |
| 292 | + m.indices[a]?.bind (fun i => (m.values[i]?)) |
| 293 | + getElem! m a := |
| 294 | + m.indices[a]?.bind (fun i => (m.values[i]?)) |>.getD default |
| 295 | + |
| 296 | +@[local grind =] private theorem getElem_def |
| 297 | + (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 298 | + m[a] = m.values[m.indices[a]'h] := |
| 299 | + rfl |
| 300 | +@[local grind =] private theorem getElem?_def |
| 301 | + (m : IndexMap α β) (a : α) : |
| 302 | + m[a]? = m.indices[a]?.bind (fun i => (m.values[i]?)) := |
| 303 | + rfl |
| 304 | +@[local grind =] private theorem getElem!_def |
| 305 | + [Inhabited β] (m : IndexMap α β) (a : α) : |
| 306 | + m[a]! = (m.indices[a]?.bind (m.values[·]?)).getD default := |
| 307 | + rfl |
| 308 | + |
| 309 | +instance : LawfulGetElem (IndexMap α β) α β (fun m a => a ∈ m) where |
| 310 | + getElem?_def := by grind |
| 311 | + getElem!_def := by grind |
| 312 | + |
| 313 | +@[inline] def insert (m : IndexMap α β) (a : α) (b : β) : IndexMap α β := |
| 314 | + match h : m.indices[a]? with |
| 315 | + | some i => |
| 316 | + { indices := m.indices |
| 317 | + keys := m.keys.set i a |
| 318 | + values := m.values.set i b } |
| 319 | + | none => |
| 320 | + { indices := m.indices.insert a m.size |
| 321 | + keys := m.keys.push a |
| 322 | + values := m.values.push b } |
| 323 | + |
| 324 | +instance : Singleton (α × β) (IndexMap α β) := |
| 325 | + ⟨fun ⟨a, b⟩ => (∅ : IndexMap α β).insert a b⟩ |
| 326 | + |
| 327 | +instance : Insert (α × β) (IndexMap α β) := |
| 328 | + ⟨fun ⟨a, b⟩ s => s.insert a b⟩ |
| 329 | + |
| 330 | +instance : LawfulSingleton (α × β) (IndexMap α β) := |
| 331 | + ⟨fun _ => rfl⟩ |
| 332 | + |
| 333 | +@[local grind .] |
| 334 | +private theorem WF' (i : Nat) (a : α) (h₁ : i < m.keys.size) (h₂ : a ∈ m) : |
| 335 | + m.keys[i] = a ↔ m.indices[a] = i := by |
| 336 | + have := m.WF i a |
| 337 | + grind |
| 338 | + |
| 339 | +/-- |
| 340 | +Erase the key-value pair with the given key, |
| 341 | +moving the last pair into its place in the order. |
| 342 | +If the key is not present, the map is unchanged. |
| 343 | +-/ |
| 344 | +@[inline] def eraseSwap (m : IndexMap α β) (a : α) : IndexMap α β := |
| 345 | + match h : m.indices[a]? with |
| 346 | + | some i => |
| 347 | + if w : i = m.size - 1 then |
| 348 | + { indices := m.indices.erase a |
| 349 | + keys := m.keys.pop |
| 350 | + values := m.values.pop } |
| 351 | + else |
| 352 | + let lastKey := m.keys.back |
| 353 | + let lastValue := m.values.back |
| 354 | + { indices := (m.indices.erase a).insert lastKey i |
| 355 | + keys := m.keys.pop.set i lastKey |
| 356 | + values := m.values.pop.set i lastValue } |
| 357 | + | none => m |
| 358 | + |
| 359 | +/-! ### Verification theorems (not exhaustive) -/ |
| 360 | + |
| 361 | +@[grind =] |
| 362 | +theorem mem_insert (m : IndexMap α β) (a a' : α) (b : β) : |
| 363 | + a' ∈ m.insert a b ↔ a' = a ∨ a' ∈ m := by |
| 364 | + grind +locals |
| 365 | + |
| 366 | +@[grind =] |
| 367 | +theorem getElem_insert (m : IndexMap α β) (a a' : α) (b : β) (h : a' ∈ m.insert a b) : |
| 368 | + (m.insert a b)[a'] = if h' : a' == a then b else m[a'] := by |
| 369 | + grind +locals |
| 370 | + |
| 371 | +theorem findIdx_lt (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 372 | + m.findIdx a h < m.size := by |
| 373 | + grind +locals |
| 374 | + |
| 375 | +grind_pattern findIdx_lt => m.findIdx a h |
| 376 | + |
| 377 | +@[grind =] |
| 378 | +theorem findIdx_insert_self (m : IndexMap α β) (a : α) (b : β) : |
| 379 | + (m.insert a b).findIdx a = if h : a ∈ m then m.findIdx a else m.size := by |
| 380 | + grind +locals |
| 381 | + |
| 382 | +@[grind =] |
| 383 | +theorem findIdx?_eq (m : IndexMap α β) (a : α) : |
| 384 | + m.findIdx? a = if h : a ∈ m then some (m.findIdx a h) else none := by |
| 385 | + grind +locals |
| 386 | + |
| 387 | +@[grind =] |
| 388 | +theorem getIdx_findIdx (m : IndexMap α β) (a : α) (h : a ∈ m) : |
| 389 | + m.getIdx (m.findIdx a) = m[a] := by grind +locals |
| 390 | + |
| 391 | +omit [LawfulBEq α] [LawfulHashable α] in |
| 392 | +@[grind =] |
| 393 | +theorem getIdx?_eq (m : IndexMap α β) (i : Nat) : |
| 394 | + m.getIdx? i = if h : i < m.size then some (m.getIdx i h) else none := by |
| 395 | + grind +locals |
| 396 | + |
| 397 | +end IndexMap |
0 commit comments