Total and partial maps in Lean 4 (inspired by Software Foundations)

maps.lean

/-
  Total and partial maps

  This is inspired by [the relevant chapter in Software Foundations]. Unlike
  Software Foundations, these maps are also parameterised over a `Key`
  type, which must supply an implementation of `DecidableEq`.

  [Software Foundations]: https://softwarefoundations.cis.upenn.edu/lf-current/Maps.html
-/
namespace Maps

  /-
    Total maps
  -/
  structure TotalMap (Key : Type) [keyEq : DecidableEq Key] (Entry : Type) where
    get : Key → Entry


  namespace TotalMap

    variables {Key : Type} [keyEq : DecidableEq Key] {Entry : Type}


    /- Operations -/

    /- Update an entry in the map -/
    @[reducible]
    def update (map : TotalMap Key Entry) (key : Key) (entry : Entry) : TotalMap Key Entry := {
      get := fun key' =>
          if key = key' then entry else map.get key'
    }

    /- Construct an empty map -/
    @[reducible]
    def empty (default : Entry) : TotalMap Key Entry := {
      get := fun key => default
    }


    /- Examples -/

    private def exampleTotalMap :=
      empty (default := false)
        |> update (key := 1) (entry := false)
        |> update (key := 3) (entry := true)

    example : exampleTotalMap.get 0 = false := rfl
    example : exampleTotalMap.get 1 = false := rfl
    example : exampleTotalMap.get 2 = false := rfl
    example : exampleTotalMap.get 3 = true := rfl


    /- Properties -/

    /- Eta-conversion -/
    theorem eta : ∀ (map : TotalMap Key Entry), map = { get := map.get }
      | { get := get } => rfl

    /- Congruence of `get` -/
    private theorem congrGet {map₁ map₂ : TotalMap Key Entry} (h : map₁.get = map₂.get) :
      map₁ = map₂
    := by
      rw [eta map₁, eta map₂]
      apply congrArg
      exact h

    /- The empty map returns its default element for all keys -/
    theorem getEmpty (key : Key) (default : Entry) :
      (empty (default := default)).get key = default
    :=
      rfl

    /- If we update a `map` at a `key` with a new `entry` and then look up `key`
       in the map resulting from the update, we get back `entry`:  -/
    theorem getUpdateEq (map : TotalMap Key Entry) key entry :
      (map.update key entry).get key = entry
    :=
      ifPos (c := key = key) rfl

    /- If we update a `map` at `key₁` and then look up a different `key₂` in
       the resulting map, we get the same result that `map` would have given -/
    theorem getUpdateNeq (map : TotalMap Key Entry) key₁ key₂ entry (h : key₁ ≠ key₂) :
      (map.update key₁ entry).get key₂ = map.get key₂
    :=
      ifNeg (c := key₁ = key₂) (Ne.intro h)

    /- If we update a `map` at `key` with `entry₁` and then update again with
       the same `key` and `entry₂`, the resulting map behaves the same (ie.
       gives the same result when applied to any key) as the simpler map
       obtained by performing just the second update on `map`. -/
    theorem updateShadow (map : TotalMap Key Entry) entry₁ entry₂ key :
      (map.update key entry₁).update key entry₂ = map.update key entry₂
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key'
      match decEq key key' with
        | isTrue h => repeat rw [ifPos h]
        | isFalse h => repeat rw [ifNeg h]

    /- If we update a `map` to assign `key` the same entry as it already has in
       `map`, then the result is equal to `map` -/
    theorem updateSame (map : TotalMap Key Entry) key :
      map.update key (map.get key) = map
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key'
      match decEq key key' with
        | isTrue h => rw [ifPos h]; rw [h]
        | isFalse h => rw [ifNeg h]

    /- If we update a `map` at two distinct keys, it doesn't matter in which
       order we do the updates. -/
    theorem updatePermute entry₁ entry₂ key₁ key₂ (map : TotalMap Key Entry) (h : key₂ ≠ key₁) :
      (map.update key₂ entry₂).update key₁ entry₁
        = (map.update key₁ entry₁).update key₂ entry₂
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key
      match decEq key₁ key, decEq key₂ key with
        | isTrue q, isTrue u =>
            rw [ifPos q, ifPos u]
            exact absurd (Eq.trans u q.symm) h
        | isTrue q, isFalse u => repeat rw [ifPos q]; repeat rw [ifNeg u]
        | isFalse q, isTrue u => repeat rw [ifNeg q]; repeat rw [ifPos u]
        | isFalse q, isFalse u => repeat rw [ifNeg q]; repeat rw [ifNeg u]

  end TotalMap


  /-
    Partial Maps
  -/
  structure PartialMap (Key : Type) [keyEq: DecidableEq Key] (Entry : Type) where
    get : Key → Option Entry

  -- TODO: Lean 4's positivity checker can't see through definitions yet :(
  -- abbrev PartialMap (Key : Type) [keyEq: DecidableEq Key] (Entry : Type) :=
  --   TotalMap Key (Option Entry)

  namespace PartialMap

    variables {Key : Type} [keyEq: DecidableEq Key] {Entry : Type}


    /- Operations -/

    /- Update an entry in the map -/
    @[reducible]
    def update (map : PartialMap Key Entry) (key : Key) (entry : Entry) : PartialMap Key Entry := {
      get := fun key' =>
        if key = key' then some entry else map.get key'
    }

    /- Construct an empty map -/
    @[reducible]
    def empty : PartialMap Key Entry := {
      get := fun key => none
    }


    /- Examples -/

    private def exampleMap :=
      empty
        |> update (key := "Church") (entry := false)
        |> update (key := "Turing") (entry := true)

    example : exampleMap.get "Church" = false := rfl
    example : exampleMap.get "Turing" = true := rfl
    example : exampleMap.get "Gödel" = none := rfl


    /- Properties -/

    /- Eta-conversion -/
    theorem eta : ∀ (map : PartialMap Key Entry), map = { get := map.get }
      | { get := get } => rfl

    /- Congruence of `get` -/
    private theorem congrGet {map₁ map₂ : PartialMap Key Entry} (h : map₁.get = map₂.get) :
      map₁ = map₂
    := by
      rw [eta map₁, eta map₂]
      apply congrArg
      exact h

    /- The empty map returns its default element for all keys -/
    theorem getEmpty (key : Key) :
      (empty (Entry := Entry)).get key = none
    :=
      rfl

    /- If we update a `map` at a `key` with a new `entry` and then look up `key`
       in the map resulting from the update, we get back `entry`:  -/
    theorem getUpdateEq (map : PartialMap Key Entry) key entry :
      (map.update key entry).get key = entry
    :=
      ifPos (c := key = key) rfl

    /- If we update a `map` at `key₁` and then look up a different `key₂` in
       the resulting map, we get the same result that `map` would have given -/
    theorem getUpdateNeq (map : PartialMap Key Entry) key₁ key₂ entry (h : key₁ ≠ key₂) :
      (map.update key₁ entry).get key₂ = map.get key₂
    :=
      ifNeg (c := key₁ = key₂) (Ne.intro h)

    /- If we update a `map` at `key` with `entry₁` and then update again with
       the same `key` and `entry₂`, the resulting map behaves the same (ie.
       gives the same result when applied to any key) as the simpler map
       obtained by performing just the second update on `map`. -/
    theorem updateShadow (map : PartialMap Key Entry) entry₁ entry₂ key :
      (map.update key entry₁).update key entry₂ = map.update key entry₂
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key'
      match decEq key key' with
        | isTrue h => repeat rw [ifPos h]
        | isFalse h => repeat rw [ifNeg h]

    /- If we update a `map` to assign `key` the same entry as it already has in
       `map`, then the result is equal to `map` -/
    theorem updateSame (map : PartialMap Key Entry) key entry (h : map.get key = some entry) :
      map.update key entry = map
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key'
      match decEq key key' with
        | isTrue q =>
          rw [ifPos q, q.symm]
          exact h.symm
        | isFalse q => rw [ifNeg q]

    /- If we update a `map` at two distinct keys, it doesn't matter in which
       order we do the updates. -/
    theorem updatePermute entry₁ entry₂ key₁ key₂ (map : PartialMap Key Entry) (h : key₂ ≠ key₁) :
      (map.update key₂ entry₂).update key₁ entry₁
        = (map.update key₁ entry₁).update key₂ entry₂
    := by
      apply congrGet
      simp -- FIXME: avoid `simp`?
      funext key
      match decEq key₁ key, decEq key₂ key with
        | isTrue q, isTrue u =>
            rw [ifPos q, ifPos u]
            exact absurd (Eq.trans u q.symm) h
        | isTrue q, isFalse u => repeat rw [ifPos q]; repeat rw [ifNeg u]
        | isFalse q, isTrue u => repeat rw [ifNeg q]; repeat rw [ifPos u]
        | isFalse q, isFalse u => repeat rw [ifNeg q]; repeat rw [ifNeg u]

    /- A proof that all the entries in `map₁` are also in `map₂` -/
    def Inclusion (map₁ map₂ : PartialMap Key Entry) :=
      ∀ key entry,
          map₁.get key = some entry
            → map₂.get key = some entry

    /- Map updates preserve inclusion -/
    theorem inclusionUpdate (map₁ map₂: PartialMap Key Entry) key entry (h : Inclusion map₁ map₂) :
        Inclusion (map₁.update key entry) (map₂.update key entry)
    := by
      intro key' entry'
      match decEq key key' with
        | isTrue q =>
            simp -- FIXME: avoid `simp`?
            repeat rw [ifPos q]
            exact id
        | isFalse q =>
            repeat rw [getUpdateNeq (h := q)]
            exact h key' entry'

  end PartialMap

end Maps