Implement a simple stack machine and prove things about it in lean4

assignment.lean

section
  abbrev Stack := List Int

  def nAdd (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: [] => ([x], pc)
    | x :: y :: xs => ((x + y) :: xs, pc + 1)
  
  def nSub (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: [] => ([x], pc)
    | x :: y :: xs => ((x - y) :: xs, pc + 1)
  
  def nEq (s : Stack) (pc : Nat): Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: [] => ([x], pc)
    | x :: y :: xs => ((if x = y then 1 else 0) :: xs, pc + 1)
  
  def nDrop (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | _ :: xs => (xs, pc + 1)
  
  def nDup (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: xs => (x :: x :: xs, pc + 1)
  
  def nSwap (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: [] => ([x], pc)
    | x :: y :: xs => (y :: x :: xs, pc + 1)
  
  def nOver (s : Stack) (pc : Nat): Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: [] => ([x], pc)
    | x :: y :: xs => (y :: x :: y :: xs, pc + 1)
  
  def nLit (i : Int) (s : Stack) (pc : Nat) : Stack × Nat := 
    (i :: s, pc + 1)

  def nIf (addr : Nat) (s : Stack) (pc : Nat) : Stack × Nat :=
    match s with
    | [] => ([], pc)
    | x :: xs => if x == 1 then (xs, addr) else (xs, pc + 1)
  
  inductive Instr : Type
  | add : Instr
  | sub : Instr
  | eq  : Instr
  | drop : Instr
  | dup  : Instr
  | swap : Instr
  | over : Instr
  | lit  : Int → Instr
  | if_  : Nat → Instr
  deriving Repr

  abbrev Prog := List Instr

  def execute_instr (i : Instr) (s : Stack) (pc : Nat) : Stack × Nat :=
    match i with
    | Instr.add => nAdd s pc
    | Instr.sub => nSub s pc
    | Instr.eq  => nSub s pc
    | Instr.drop => nDrop s pc
    | Instr.dup => nDup s pc
    | Instr.swap => nSwap s pc
    | Instr.over => nOver s pc
    | Instr.lit i => nLit i s pc
    | Instr.if_ a => nIf a s pc

  def execute_prog (p : Prog) (prog_len : Nat) (s : Stack) (pc : Nat) : Stack × Nat :=
    if prog_len == pc then (s, pc)
    else
      match p with
      | [] => (s, pc)
      | instr :: instrs => 
        let (s', pc') := execute_instr instr s pc
        execute_prog instrs prog_len s' pc'
  
  def stack_ex : Stack := [1, 2, 3, 4]
  def prog_ex : Prog :=
    [
      Instr.dup, -- 0
      Instr.lit 0,  -- 1 
      Instr.eq,  -- 2 
      Instr.if_ 12, -- 3
      Instr.swap, -- 4
      Instr.lit 1, -- 5
      Instr.add, -- 6
      Instr.swap, -- 7
      Instr.lit 1, -- 8
      Instr.sub, -- 9
      Instr.lit 0, -- 10
      Instr.if_ 0 -- 11
    ]

-- Dup
-- Lit 0
-- Eq
-- If 12
-- Swap
-- Lit 1
-- Add
-- Swap
-- Lit 1
-- Sub
-- Lit 0
-- If 0

  -- def final := execute_prog prog_ex (List.length prog_ex) stack_ex 0
  -- #eval final

  -- example (l : List Int)

  theorem involution_nswap 
    (s : Stack) : 
    Prod.fst 
    (execute_instr 
      (Instr.swap) 
      (Prod.fst 
        (execute_instr 
          (Instr.swap) 
            s 
            0)) 0) = s := by
      match s with
      | [] => exact Eq.refl []
      | x :: [] => exact Eq.refl [x]
      | x :: y :: xs => exact Eq.refl (x :: y :: xs)
end