Need help with improving a basic Lean example

Stream: new members

[Rick van Oosterhout]

I worked out a basic concept that I want to use to explain some of the concepts of dependent type theory, calculus of (inductive) constructions and Lean, using some analogies with Lego's. I came up with the following definitions and managed to set up a simple theorem and prove it. However, I would like to have some feedback from more experienced Lean users.

Mainly looking for tips to expand the analogy to explain other relevant topics, if you think I missed one. Besides, I would like to get any feedback on the Lean definitions and theorems, on how things could be simplified or clarified further. Any suggestions on my code are welcome, as I want to learn the best practices for Lean.

open Nat

-- Use natural numbers
abbrev ℕ := Nat

-- Basic brick type
inductive Lego : Type
  | brick
  | plate

namespace Lego

-- Function that adds computation to the tower
def brickHeight : Lego → ℕ
  | brick => 3
  | plate => 1

#eval brickHeight plate
#eval brickHeight brick


-- Define the inductive type for a Lego Tower
inductive LegoTower : Type
  | base : LegoTower  -- Base case: base plate
  | stack : Lego → LegoTower → LegoTower  -- Inductive step: add a lego

namespace LegoTower

def t1 : LegoTower := base
def t2 : LegoTower := stack plate base
def t3 : LegoTower := stack plate (stack brick base)


-- Define a dependent type for a tower with exactly 'n' lego's
inductive nLegoTower : ℕ → Type
  | nBase : nLegoTower 0  -- A tower of height 0 (base case)
  | nStack : Lego → (nLegoTower n) → nLegoTower (n + 1)  -- Add a lego

namespace nLegoTower

-- Example of a tower with exactly 3 lego's
def nLegoTower3 : nLegoTower 3 := (nStack brick (nStack plate (nStack brick (nBase))))

def nToLegoTower {n: ℕ} : nLegoTower n → LegoTower
  | nBase => base
  | nStack br rest => stack br (nToLegoTower rest)


-- Tower height
def recTowerHeight (tower : LegoTower) :=
  match tower with
  | base => 0
  | stack br rest => brickHeight br + recTowerHeight rest

def towerHeight (tower : LegoTower) : ℕ := recTowerHeight tower

#eval towerHeight t1
#eval towerHeight t2
#eval towerHeight t3

#eval towerHeight (nToLegoTower nLegoTower3)


-- Basic proof that 'towerHeight (nToLegoTower t) ≤ 3 * n'

-- Lemma that any brick ≤ 3
theorem brick_height_limit (b : Lego) : brickHeight b ≤ 3 := by
  induction b <;> (rw [brickHeight]; simp)

-- Lemma for the inductive step below
theorem n_tower_max_height_ind (n : ℕ) (t : nLegoTower m) :
  towerHeight (nToLegoTower t) ≤ 3 * n → brickHeight b + towerHeight (nToLegoTower t) ≤ 3 * n + 3 := by
    intro h
    calc
      brickHeight b + towerHeight (nToLegoTower t)
        ≤ 3 + towerHeight (nToLegoTower t) := by apply Nat.add_le_add_right (brick_height_limit b)
      _ ≤ 3 + 3 * n := by apply Nat.add_le_add_left h
      _ = 3 * n + 3 := by rw [Nat.add_comm]

-- Complete proof
theorem n_tower_max_height (t : nLegoTower n) : towerHeight (nToLegoTower t) ≤ 3 * n := by
  induction t with
  | nBase =>
    -- Simplify case
    rw [nToLegoTower, towerHeight, recTowerHeight, Nat.mul_zero, Nat.le_zero_eq]
  | nStack b rt ih =>
    -- Move out brickheight b
    rw [nToLegoTower, towerHeight, recTowerHeight, ← towerHeight, Nat.mul_add, Nat.mul_one]
    -- Solve using lemma
    exact n_tower_max_height_ind _ rt ih