Type-Driven Development: insertion sort with proof (O(n²))

InsertionSortList.lean

-- Prompt: "Define a basic linked-list of some type α, and a dependent-type
-- SortedList that provides the proof the list is actually sorted with respect to the ordering on α."

namespace SortList

/-- A basic singly-linked list, generic over a type `α`. -/
inductive List (α : Type) where
  | nil : List α
  | cons : α → List α → List α
deriving Repr

/--
A proof that a `List` is sorted with respect to the `≤` ordering on `α`.

- `nil`: the empty list is sorted.
- `singleton`: a one-element list is sorted.
- `cons x y xs (hxy : x ≤ y) (h : Sorted (.cons y xs))`: prepending `x`
  is sorted if `x ≤ y` and the tail `y :: xs` is itself sorted.
-/
inductive Sorted [LE α] : List α → Prop where
  | nil : Sorted .nil
  | singleton (x : α) : Sorted (.cons x .nil)
  | cons (x y : α) (xs : List α) (hxy : x ≤ y) (h : Sorted (.cons y xs)) :
      Sorted (.cons x (.cons y xs))

/-- A sorted list: a `List` bundled with a proof that it is sorted. -/
structure SortedList (α : Type) [LE α] where
  list : List α
  sorted : Sorted list

/-- Decide at runtime whether a `List` is sorted. -/
def List.isSorted [LE α] [DecidableRel (α := α) (· ≤ ·)] : List α → Bool
  | .nil => true
  | .cons _ .nil => true
  | .cons x (.cons y xs) => x ≤ y && (cons y xs).isSorted

/-- Insert an element into its sorted position in a list. -/
def List.insert [LE α] [DecidableRel (α := α) (· ≤ ·)] (x : α) : List α → List α
  | .nil => .cons x .nil
  | .cons y ys => if x ≤ y then .cons x (.cons y ys) else .cons y (ys.insert x)

/-- Insertion sort. -/
def List.insertionSort [LE α] [DecidableRel (α := α) (· ≤ ·)] : List α → List α
  | .nil => .nil
  | .cons x xs => (xs.insertionSort).insert x

/-- Inserting into a sorted list preserves sortedness. -/
theorem List.insert_sorted [LE α] [DecidableRel (α := α) (· ≤ ·)]
    (total : ∀ (a b : α), a ≤ b ∨ b ≤ a)
    (x : α) : ∀ (l : List α), Sorted l → Sorted (l.insert x) := by
  intro l
  induction l with
  | nil => intro _; exact .singleton x
  | cons y ys ih =>
    intro h
    unfold List.insert
    split
    case isTrue hxy => exact .cons x y ys hxy h
    case isFalse hxy =>
      have hyx := (total x y).resolve_left hxy
      cases ys with
      | nil =>
        simp [List.insert]
        exact .cons y x .nil hyx (.singleton x)
      | cons z zs =>
        have hyz : y ≤ z := by cases h with | cons _ _ _ h _ => exact h
        have htail : Sorted (.cons z zs) := by cases h with | cons _ _ _ _ h => exact h
        have ih' := ih htail
        -- ih' : Sorted ((.cons z zs).insert x)
        -- goal : Sorted (.cons y ((.cons z zs).insert x))
        unfold List.insert at ih' ⊢
        split
        next hxz =>
          rw [if_pos hxz] at ih'
          exact .cons y x (.cons z zs) hyx ih'
        next hxz =>
          rw [if_neg hxz] at ih'
          exact .cons y z (zs.insert x) hyz ih'

/-- Insertion sort produces a sorted list. -/
theorem List.insertionSort_sorted [LE α] [DecidableRel (α := α) (· ≤ ·)]
    (total : ∀ (a b : α), a ≤ b ∨ b ≤ a)
    : ∀ (l : List α), Sorted (l.insertionSort)
  | .nil => .nil
  | .cons x xs => List.insert_sorted total x xs.insertionSort (insertionSort_sorted total xs)

/-- Sort a list, returning a `SortedList` with proof of sortedness. -/
def List.sort [LE α] [DecidableRel (α := α) (· ≤ ·)]
    (total : ∀ (a b : α), a ≤ b ∨ b ≤ a)
    (l : List α) : SortedList α :=
  ⟨l.insertionSort, List.insertionSort_sorted total l⟩

instance : ToString (List Nat) where
  toString l :=
    let rec go : List Nat → String
      | .nil => "nil"
      | .cons x xs => s!"{x} :: {go xs}"
    go l

instance : ToString (SortedList Nat) where
  toString sl := toString sl.list

def main : IO Unit := do
  let sorted := List.cons 1 (.cons 2 (.cons 3 .nil))
  IO.println s!"List:   {sorted}"
  IO.println s!"Sorted: {sorted.isSorted}"
  IO.println ""
  let unsorted := List.cons 3 (.cons 1 (.cons 2 .nil))
  IO.println s!"List:   {unsorted}"
  IO.println s!"Sorted: {unsorted.isSorted}"
  IO.println ""
  let unsortedSorted := unsorted.sort (fun a b => by omega)
  IO.println s!"List:   {unsortedSorted}"
  IO.println s!"Sorted: {unsortedSorted.list.isSorted}"

end SortList