MergeSort with Lean

Merge.lean

-- This module serves as the root of the `LeanPlayground` library.
-- Import modules here that should be built as part of the library.
import LeanPlayground.Basic
import Mathlib.Tactic.Tauto
import Mathlib.Order.Basic
import Mathlib
import Mathlib.Tactic.Linarith
import Mathlib.Data.List.Perm.Basic

#eval 2 ≤ 3

-- Note 1: Can't refer to `ℕ` without importing Mathlib
-- Note 2: ≤ creates a Proposition, not a Boolean, but you can still use it in an `if` somehow
-- outputs Prop
#check 2 ≤ 3

-- define the merge function on lists of natural numbers
def mergeNat (xs ys : List Nat) : List Nat :=
  match xs, ys with
  | [], ys => ys
  | xs, [] => xs
  | x :: xs', y :: ys' =>
    if x ≤ y then
      x :: mergeNat xs' (y :: ys')
    else
      y :: mergeNat (x :: xs') ys'

-- let's test the merge function on some simple examples

#eval mergeNat [] []
#eval mergeNat [1, 3, 5] []
#eval mergeNat [] [2, 4, 6]
#eval mergeNat [1, 3, 5] [2, 4, 6]

/-
 Note 3: You need the following annotation for this to typecheck
 [DecidableRel ((· ≤ ·) : α → α → Prop)]
-/

-- #print Ord
-- #print LE


def merge {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys : List α) : List α :=
  match xs, ys with
  | [], ys => ys
  | xs, [] => xs
  | x :: xs', y :: ys' =>
    if x ≤ y then
      x :: merge xs' (y :: ys')
    else
      y :: merge (x :: xs') ys'

def mergeWithOrd {α : Type} [Ord α] (xs ys : List α) : List α :=
  match xs, ys with
  | [], ys => ys
  | xs, [] => xs
  | x :: xs', y :: ys' =>
    -- Note 4: can't use ≤ here
    match compare x y with
    | Ordering.lt => x :: mergeWithOrd xs' (y :: ys')
    | _ => y :: mergeWithOrd (x :: xs') ys'

-- define increasing (by pattern matching and recursion on the list)
def increasing {α : Type} [LE α] (xs : List α) : Prop :=
  match xs with
  | [] => True
  | [_] => True
  | x :: y :: xs' => x ≤ y ∧ increasing (y :: xs')

-- let's prove that [1, 2, 3, 4] is increasing
def increasing1234 : increasing [1, 2, 3, 4] :=
  by
    constructor
    · decide
    · constructor
      · decide
      · constructor
        · decide
        · constructor

def increasing1234' : increasing [1, 2, 3, 4] :=
  by
    dsimp [increasing] -- expands the goal to 1 ≤ 2 ∧ 2 ≤ 3 ∧ 3 ≤ 4 ∧ True
    decide

def increasing1234'' : increasing [1, 2, 3, 4] :=
  by
    -- blasts the goal into many goals
    (repeat constructor)
    <;>
    -- apply decide to each of them
    decide

-- a boolean version of increasing
def isIncreasing {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) : Bool :=
  match xs with
  | [] => true
  | [_] => true
  | x :: y :: xs' => (x ≤ y) && isIncreasing (y :: xs')

-- Let's show that isIncreasing xs iff increasing xs
theorem isIncreasing_iff_increasing {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) :
  isIncreasing xs = true ↔ increasing xs :=
  by
    induction xs with
    | nil =>
      dsimp [isIncreasing, increasing]
      trivial
    | cons x xs ih =>
        cases xs with
        | nil =>
          dsimp [isIncreasing, increasing]
          trivial
        | cons y xs' =>
          dsimp [isIncreasing, increasing]
          -- the goal now has decide (x ≤ y) on one side and x ≤ y on the other
          -- let's case split
          by_cases h : x ≤ y
          · -- case 1: x ≤ y
            -- need to rewrite decide (x ≤ y) using the fact that h : x ≤ y
            -- thankfully, simp understands this idea
            simp [h]
            rw [ih]
          · -- case 2: ¬(x ≤ y)
            -- need to rewrite decide (x ≤ y) using the fact that h : ¬(x ≤ y)
            -- simp does all the heavy lifting here
            simp [h]

-- let's state that increasing is a Decidable Relation
instance {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) : Decidable (increasing xs) :=
  decidable_of_iff (isIncreasing xs = true) (isIncreasing_iff_increasing xs)

-- Note 4: now we can use increasing in an if statement
def increasingOrNot {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) : String :=
  if increasing xs then
    "increasing"
  else
    "not increasing"

-- define increasing2 (as an inductive predicate)
inductive Increasing {α : Type} [LE α] : List α → Prop
| nil : Increasing []
| singleton (x : α) : Increasing [x]
| cons (x y : α) (xs : List α) : x ≤ y → Increasing (y :: xs) → Increasing (x :: y :: xs)

example : Increasing [1, 2, 3, 4] :=
  by
    (repeat constructor) <;> decide

theorem increasing_iff_Increasing {α : Type} [LE α] (xs : List α) : increasing xs ↔ Increasing xs :=
  by
    induction xs with
    | nil =>
      dsimp [increasing]
      -- True ↔ Increasing []
      simp [Increasing.nil]
    | cons x xs ih =>
        cases xs with
        | nil =>
          dsimp [increasing]
          simp [Increasing.singleton]
        | cons y xs' =>
          dsimp [increasing, Increasing]
          rewrite [ih]
          -- x ≤ y ∧ Increasing (y :: xs') ↔ Increasing (x :: y :: xs')
          -- no hypotheses is needed here, we just need to use Increasing.cons
          -- the following short proof was suggested by gemini:
          constructor <;> intro h <;> cases h <;> constructor <;> assumption

instance {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) : Decidable (Increasing xs) :=
  by
    apply decidable_of_iff (isIncreasing xs = true)
    rw [isIncreasing_iff_increasing xs, increasing_iff_Increasing xs]

lemma merge_nil_left {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (ys : List α) :
  merge [] ys = ys :=
  by
    -- rfl doesn't work here for some reason
    simp [merge]

-- Why does rfl not work? That's probably because merge is not what it seems.

#print merge

-- What is this merge._unary
#print merge._unary

/-

the decreasing argument of merge differs in the two branches

if x ≤ y then
      x :: mergeNat xs' (y :: ys')
    else
      y :: mergeNat (x :: xs') ys'

so Lean internally desugars it into a more complex function that uses well-founded recursion

-/

-- Note 5: note that you have to case split here, the proof is not as simple as merge_nil_left
lemma merge_nil_right {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) :
  merge xs [] = xs :=
  by
    -- this won't make progress
    -- simp [merge]
    cases xs with
    | nil =>
      simp [merge]
    | cons x xs =>
      simp [merge]

lemma merge_nil_nil {α : Type} [LE α] [DecidableRel ((· ≤ ·) : α → α → Prop)] :
  forall (xs ys : List α), merge xs ys = [] → xs = [] ∧ ys = [] := by
  intro xs ys h
  match xs, ys with
  | [], [] => trivial
  | (x :: xs), [] =>
    simp [merge] at h
  | [], (y :: ys) =>
    simp [merge] at h
  | (x :: xs), (y :: ys) =>
    simp [merge] at h
    by_cases hxy : x ≤ y
    · simp [hxy] at h
    · simp [hxy] at h

-- Somehow #find doesn't work
-- #find (_ : Nat) + _ = _ + _

-- Note 6: at this point, we need [LinearOrder α] to get properties of ≤ because LE had no axioms attached
-- let's show that merge preserves increasing
theorem merge_increasing {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys : List α) :
  increasing xs → increasing ys → increasing (merge xs ys) :=
  by
    generalize h : xs.length + ys.length = n
    induction n generalizing xs ys with
    | zero =>
      have Lxs : xs.length = 0 := by omega
      have Lys : ys.length = 0 := by omega
      have Lxs' : xs = [] := by
        -- try running exact? or rw? here
        exact List.eq_nil_iff_length_eq_zero.mpr Lxs
      have Lys' : ys = [] := by
        exact List.eq_nil_iff_length_eq_zero.mpr Lys
      rw [Lxs', Lys']
      simp [increasing, merge]
    | succ n ih =>
      match xs, ys with
      | [], ys =>
          simp [merge]
      | xs, [] =>
        rw [merge_nil_right]
        tauto
      | x :: xs', y :: ys' =>
        intro hxs hys
        simp [merge]
        by_cases hxy : x ≤ y
          <;> simp [hxy]
        · -- case 1: x ≤ y
          match xs' with
          | [] =>
            simp [merge]
            constructor
            · assumption
            · assumption
          | x' :: xs'' =>
            set l := merge (x' :: xs'') (y :: ys') with l_eq
            have increasing_l : increasing l := by
              apply ih
              · simp at *
                omega
              · exact hxs.2
              · assumption
            have hxx' : x ≤ x' := by
              exact hxs.1
            by_cases h' : x' ≤ y
              <;> simp [h', merge] at l_eq
              <;> subst l
              <;> rw [l_eq]
              <;> constructor
              <;> (first | assumption | rw [← l_eq]; assumption)
        · -- case 2: ¬(x ≤ y)
          match ys' with
          | [] =>
            simp [merge_nil_right, increasing]
            constructor
            -- Note 7: I spent a lot of time on the following trivial tactic
            -- linarith didn't work, omega isn't available, etc, etc
            · exact Std.le_of_not_ge hxy
            · assumption
          | y' :: ys'' =>
            set l := merge (x :: xs') (y' :: ys'') with l_eq
            have increasing_l : increasing l := by
              apply ih
              · simp at *
                omega
              · assumption
              · exact hys.2
            have hyy' : y ≤ y' := by
              exact hys.1
            by_cases h' : (x ≤ y')
              <;> simp [h', merge] at l_eq
              <;> subst l
              <;> rw [l_eq]
              <;> simp [increasing]
              <;> constructor
            · exact Std.le_of_not_ge hxy
            · rw [l_eq] at increasing_l
              exact increasing_l
            · exact hyy'
            · rw [l_eq] at increasing_l
              exact increasing_l

-- Observation: Gemini will very confidently give you wrong proofs if the proof is longer

/-

Gemini says that you can also organize the proof thusly. It gives some slightly different Induction Hypotheses, so maybe it would work cleanly. I haven't tried it

theorem merge_increasing_2 [LinearOrder α] (xs ys : List α) :
  increasing xs → increasing ys → increasing (merge xs ys) := by
  induction xs, ys using merge.induct with
  | case1 ys => sorry
  | case2 xs => sorry
  | case3 x y xs' ys' hxy ih => sorry
  | case4 x y xs' ys' hxy ih => sorry

-/

def mergeSort {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)]
  (xs : List α) : List α :=
  -- Note 8: note the use of eq_xs. This will generate a useful hypothesis for our proof
  match eq_xs : xs with
  | [] => []
  | [x] => [x]
  | x1 :: x2 :: rest =>
    let n := xs.length / 2
    let sortedFirstHalf := mergeSort (xs.take n)
    let sortedSecondHalf := mergeSort (xs.drop n)
    merge sortedFirstHalf sortedSecondHalf
termination_by xs.length
decreasing_by
  · subst xs
    -- Note 9: Interestingly, simp knows how to simplify List.take (List.length _)
    -- Apparently, you can annotate lemmas with @[simp] to specify that it is a hint
    simp [List.length] at *
    omega
  · subst xs
    simp [List.length] at *
    omega

-- Note 10: I was prepared for this proof to be very difficult, but once I used
-- mergeSort.induct, Copilot practically filled in the rest
-- with very little intervention (I had to name the induction hypotheses in the way that I have)
lemma mergeSort_increasing {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) :
  increasing (mergeSort xs) :=
by
  induction xs using mergeSort.induct with
  | case1 => simp [mergeSort, increasing]
  | case2 => simp [mergeSort, increasing]
  | case3 x1 x2 rest n _ _ ih3 ih4 =>
    simp [mergeSort]
    apply merge_increasing
    · exact ih3
    · exact ih4

#check List.Perm

local infix:50 " ~ " => List.Perm

/- Note 11: There is some variation you could consider here. I believe this is the simplest one to do the induction on.

Some other possibilities include:

lemma merge_perm_app {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys zs: List α) :
  zs = xs ++ ys → (merge xs ys) ~ zs

lemma merge_perm_app {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys: List α) :
 (merge xs ys) ~ (xs ++ ys)

lemma merge_perm_app {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys: List α) :
  forall zs, zs ~ xs ++ ys → (merge xs ys) ~ zs

-/
lemma merge_perm_app {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys: List α) :
  forall zs, zs = xs ++ ys → (merge xs ys) ~ zs :=
  by
    induction xs, ys using merge.induct with
    | case1 ys =>
      simp [merge]
    | case2 xs =>
      simp [merge]
    | case3 x xs y ys hxy ih =>
      intros zs Hzs
      simp [merge, hxy]
      cases zs with
      | nil => nomatch Hzs
      | cons z zs =>
        injection Hzs with Ezx Hzs
        subst z
        apply List.Perm.cons
        -- automatically applies the induction hypothesis ih
        tauto
    | case4 x xs y ys hxy ih =>
      intros zs Hzs
      simp [merge, hxy]
      have Hintermed : (zs ~ y :: x :: xs ++ ys) := by
        subst zs
        calc
          x :: xs ++ y :: ys ~ x :: (xs ++ [y]) ++ ys :=
              by
              apply List.Perm.cons
              simp
          _ ~ x :: ([y] ++ xs) ++ ys :=
              by
              apply List.Perm.cons
              apply List.Perm.append_right
              apply List.perm_append_comm
          _ ~ x :: y :: xs ++ ys := by rfl
          _ ~ y :: x :: xs ++ ys := by tauto
      trans (y :: x :: xs ++ ys)
      · -- ih is used here
        tauto
      · tauto



lemma merge_perm_cons_left {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys : List α) (x : α) :
  merge (x :: xs) ys ~ x :: merge xs ys :=
  by
    calc
      merge (x :: xs) ys ~ (x :: xs) ++ ys := by apply merge_perm_app; simp
      _ ~ x :: (xs ++ ys) := by simp
      _ ~ x :: merge xs ys := by apply List.Perm.cons; symm; simp [merge_perm_app]

lemma merge_perm_cons_right {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys : List α) (y : α) :
  merge xs (y :: ys) ~ y :: merge xs ys :=
  by
    calc
      merge xs (y :: ys) ~ xs ++ (y :: ys)                  := by apply merge_perm_app; simp
                       _ ~ (xs ++ [y]) ++ ys                := by simp
                       _ ~ y :: xs ++ ys                    := by simp
                       _ ~ y :: merge xs ys                 := by apply List.Perm.cons; symm; simp [merge_perm_app]


/-

I thought that the above lemmas would be necessary, but they weren't

I struggled to prove the following for a long time, but later I realized that merge_perm_app actually gives us a "canonical form" for merge xs ys

-/

lemma merge_perm {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs ys: List α) :
  forall xs' ys', xs' ~ xs → ys' ~ ys → (merge xs ys) ~ (merge xs' ys') :=
  by
    intros xs' ys' Hxs Hys
    calc
      (merge xs ys) ~ (xs ++ ys) := by apply merge_perm_app; simp
                  _ ~ (xs' ++ ys) := by apply List.Perm.append_right; tauto
                  _ ~ (xs' ++ ys') := by apply List.Perm.append_left; tauto
                  _ ~ (merge xs' ys') := by symm; apply merge_perm_app; simp

lemma mergeSort_perm {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) :
  (mergeSort xs) ~ xs :=
  by
    induction xs using mergeSort.induct with
    | case1 => simp [mergeSort]
    | case2 => simp [mergeSort]
    | case3 x1 x2 rest n ih1 ih2 ih3 ih4 =>
      simp [mergeSort]
      set l := (x1 :: x2 :: rest) with Leq
      set n' := (rest.length + 1 + 1) / 2 with Eqn
      have Hn : n = n' := by
        subst n
        subst n'
        simp [List.length] at *
      rw [← Hn] at *
      clear Hn Eqn
      set l1 := List.take n l with L1eq
      set l2 := List.drop n l with L2eq
      have H := merge_perm l1 l2 (mergeSort l1) (mergeSort l2) ih3 ih4
      calc
        merge (mergeSort l1) (mergeSort l2) ~ merge l1 l2                     := by symm; simp [H]
                                          _ ~ l1 ++ l2                        := by apply merge_perm_app; simp
                                          _ ~ List.take n l ++ List.drop n l  := by rw [L1eq, L2eq]
                                          _ ~ l                               := by simp

theorem mergeSort_correct {α : Type} [LinearOrder α] [DecidableRel ((· ≤ ·) : α → α → Prop)] (xs : List α) :
  increasing (mergeSort xs) ∧ (mergeSort xs) ~ xs :=
    by exact ⟨mergeSort_increasing xs, mergeSort_perm xs⟩