Harvey–van der Hoeven machine development, part 1: stack machine and multitape-TM compilation

-- Harvey–van der Hoeven multitape-TM development (NLogN), generated by skeleton subtraction.

import Definitions.Def_IntMul_MultitapeModel
import Definitions.Def_IntMul_HvdH_ConvC
import Definitions.Def_IntMul_HvdH_ExpApprox
import Definitions.Def_IntMul_HvdH_MachineRef
import Definitions.Def_IntMul_HvdH_Numeric
import Definitions.Def_IntMul_HvdH_RealExp
import Definitions.Def_IntMul_HvdH_Twiddle
import Mathlib
import Mathlib.Tactic

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.SM
/-!
# A multi-stack machine language with procedures

`Stmt S P` is a structured program over `S` stacks of symbols (`Fin 16`) that may call any of
`P` procedures (recursively).  It reads the input word with `inRd` and writes output bits with
`out`.  `Exec` is its big-step semantics; the cost is the number of steps of the flat machine
the program compiles to (`NLogN.Flat`); each flat step is at most two Turing machine steps.
-/

namespace NLogN

/-- Stack symbols. -/
abbrev Sym := Fin 16

/-- Codes of input symbols. -/
def codeBit (b : Bool) : Sym := if b then 1 else 0
def codeSep : Sym := 2
def codeEnd : Sym := 3

/-- The input word `x # y` as codes. -/
def inputCodes (x y : List Bool) : List Sym := x.map codeBit ++ codeSep :: y.map codeBit

/-- Structured programs. Tests look at the top of one stack (`none` if it is empty). -/
inductive Stmt (S P : ℕ) where
  | skip
  | push (i : Fin S) (a : Sym)
  | pop (i : Fin S)
  | inRd (i : Fin S)
  | out (b : Bool)
  | seq (p q : Stmt S P)
  | ite (i : Fin S) (f : Option Sym → Bool) (p q : Stmt S P)
  | loop (i : Fin S) (f : Option Sym → Bool) (p : Stmt S P)
  | call (j : Fin P)

/-- Machine state: stacks (head of each list is the top), input position, output so far. -/
@[ext] structure St (S : ℕ) where
  stk : Fin S → List Sym
  ip : ℕ
  out : List Bool

namespace St

variable {S : ℕ}

def top (σ : St S) (i : Fin S) : Option Sym := (σ.stk i).head?

def push (σ : St S) (i : Fin S) (a : Sym) : St S :=
  { σ with stk := Function.update σ.stk i (a :: σ.stk i) }

def pop (σ : St S) (i : Fin S) : St S :=
  { σ with stk := Function.update σ.stk i (σ.stk i).tail }

/-- Read the next input code onto stack `i`. -/
def inRd (I : List Sym) (σ : St S) (i : Fin S) : St S :=
  { σ with stk := Function.update σ.stk i (I.getD σ.ip codeEnd :: σ.stk i), ip := σ.ip + 1 }

def outB (σ : St S) (b : Bool) : St S := { σ with out := σ.out ++ [b] }

end St

/-- Big-step semantics with exact cost. `I` is the input word, `prog` the procedures. -/
inductive Exec {S P : ℕ} (I : List Sym) (prog : Fin P → Stmt S P) :
    Stmt S P → St S → St S → ℕ → Prop
  | skip (σ) : Exec I prog .skip σ σ 1
  | push (σ i a) : Exec I prog (.push i a) σ (σ.push i a) 1
  | pop (σ i) : Exec I prog (.pop i) σ (σ.pop i) 1
  | inRd (σ i) : Exec I prog (.inRd i) σ (σ.inRd I i) 1
  | out (σ b) : Exec I prog (.out b) σ (σ.outB b) 1
  | seq {p q σ σ₁ σ₂ c₁ c₂} : Exec I prog p σ σ₁ c₁ → Exec I prog q σ₁ σ₂ c₂ →
      Exec I prog (.seq p q) σ σ₂ (c₁ + c₂)
  | iteT {i f p q σ σ' c} : f (σ.top i) = true → Exec I prog p σ σ' c →
      Exec I prog (.ite i f p q) σ σ' (c + 1)
  | iteF {i f p q σ σ' c} : f (σ.top i) = false → Exec I prog q σ σ' c →
      Exec I prog (.ite i f p q) σ σ' (c + 1)
  | loopT {i f p σ σ₁ σ₂ c₁ c₂} : f (σ.top i) = true → Exec I prog p σ σ₁ c₁ →
      Exec I prog (.loop i f p) σ₁ σ₂ c₂ → Exec I prog (.loop i f p) σ σ₂ (c₁ + c₂ + 1)
  | loopF {i f p σ} : f (σ.top i) = false → Exec I prog (.loop i f p) σ σ 1
  | call {j σ σ' c} : Exec I prog (prog j) σ σ' c → Exec I prog (.call j) σ σ' (c + 2)

namespace Exec

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

end Exec

end NLogN

end

section
-- Source module: NLogN.Logic
/-!
# Program logic for stack programs

`Runs I prog s σ σ' c`: `s` takes `σ` to `σ'` with cost at most `c`.  Besides the rules for each
statement this file provides

* `caseTops`: branch on the tops of several stacks (constant cost);
* `opsStmt`: straight-line lists of primitive operations;
* `machine`: a finite-state transducer whose state lives on a dedicated stack; its correctness
  reduces to iterating a plain Lean function `mstep`.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def Runs (I : List Sym) (prog : Fin P → Stmt S P) (s : Stmt S P) (σ σ' : St S) (c : ℕ) : Prop :=
  ∃ c' ≤ c, Exec I prog s σ σ' c'

namespace Runs

end Runs

/-! ## Branching on stack tops -/

def caseList (i : Fin S) (k : Option Sym → Stmt S P) : List Sym → Stmt S P
  | [] => .skip
  | a :: l => .ite i (fun o => decide (o = some a)) (k (some a)) (caseList i k l)

def caseTop (i : Fin S) (k : Option Sym → Stmt S P) : Stmt S P :=
  .ite i (fun o => o.isNone) (k none) (caseList i k (List.finRange 16))

/-- Branch on the tops of the stacks `w` (in order). -/
def caseTops : List (Fin S) → (List (Option Sym) → Stmt S P) → Stmt S P
  | [], k => k []
  | i :: w, k => caseTop i (fun o => caseTops w (fun os => k (o :: os)))

def St.tops (σ : St S) (w : List (Fin S)) : List (Option Sym) := w.map σ.top

/-! ## Straight-line operations -/

inductive Op (S : ℕ) where
  | push (i : Fin S) (a : Sym)
  | pop (i : Fin S)
  | inRd (i : Fin S)
  | out (b : Bool)

def Op.stmt : Op S → Stmt S P
  | .push i a => .push i a
  | .pop i => .pop i
  | .inRd i => .inRd i
  | .out b => .out b

def Op.run (I : List Sym) : Op S → St S → St S
  | .push i a, σ => σ.push i a
  | .pop i, σ => σ.pop i
  | .inRd i, σ => σ.inRd I i
  | .out b, σ => σ.outB b

def opsStmt : List (Op S) → Stmt S P
  | [] => .skip
  | o :: l => .seq o.stmt (opsStmt l)

def runOps (I : List Sym) : List (Op S) → St S → St S
  | [], σ => σ
  | o :: l, σ => runOps I l (o.run I σ)

/-! ## Finite-state transducers -/

/-- The halting state of a `machine`. -/
def hlt : Sym := 15

/-- A transducer with state stack `q`, watching stacks `w`.  Each step looks at the state and at
the tops of `w` and performs `δ (state :: tops)`, which must replace the state symbol.  It starts
by pushing `q0` and stops when the state is `hlt`, which it then pops. -/
def machine (q : Fin S) (w : List (Fin S)) (q0 : Sym)
    (δ : List (Option Sym) → List (Op S)) : Stmt S P :=
  .seq (.push q q0) (.seq (.loop q (fun o => decide (o ≠ some hlt))
    (caseTops (q :: w) fun os => opsStmt (δ os))) (.pop q))

/-- One step of a transducer, as a function. -/
def mstep (I : List Sym) (q : Fin S) (w : List (Fin S)) (δ : List (Option Sym) → List (Op S))
    (σ : St S) : St S :=
  runOps I (δ (σ.tops (q :: w))) σ

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-- `m` transducer steps lead from `σ` to `σf`, never passing through the halting state. -/
def Reaches (I : List Sym) (q : Fin S) (w : List (Fin S)) (δ : List (Option Sym) → List (Op S))
    (σ σf : St S) (m : ℕ) : Prop :=
  (mstep I q w δ)^[m] σ = σf ∧ ∀ k < m, ((mstep I q w δ)^[k] σ).top q ≠ some hlt

namespace Reaches

variable {q : Fin S} {w : List (Fin S)} {δ : List (Option Sym) → List (Op S)}

end Reaches

end NLogN

end

section
-- Source module: NLogN.Move
/-!
# Basic transducers: moving stack contents
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-- Replace stack `i`. -/
def St.set (σ : St S) (i : Fin S) (l : List Sym) : St S :=
  { σ with stk := Function.update σ.stk i l }

@[simp] theorem St.set_ip (σ : St S) (i : Fin S) (l : List Sym) : (σ.set i l).ip = σ.ip := rfl
@[simp] theorem St.set_out (σ : St S) (i : Fin S) (l : List Sym) : (σ.set i l).out = σ.out := rfl

@[simp] theorem St.push_ip (σ : St S) (i : Fin S) (a : Sym) : (σ.push i a).ip = σ.ip := rfl
@[simp] theorem St.push_out (σ : St S) (i : Fin S) (a : Sym) : (σ.push i a).out = σ.out := rfl
@[simp] theorem St.pop_ip (σ : St S) (i : Fin S) : (σ.pop i).ip = σ.ip := rfl
@[simp] theorem St.pop_out (σ : St S) (i : Fin S) : (σ.pop i).out = σ.out := rfl

@[simp] theorem St.inRd_ip (σ : St S) (i : Fin S) : (σ.inRd I i).ip = σ.ip + 1 := rfl
@[simp] theorem St.inRd_out (σ : St S) (i : Fin S) : (σ.inRd I i).out = σ.out := rfl
@[simp] theorem St.outB_stk (σ : St S) (b : Bool) : (σ.outB b).stk = σ.stk := rfl
@[simp] theorem St.outB_ip (σ : St S) (b : Bool) : (σ.outB b).ip = σ.ip := rfl
@[simp] theorem St.outB_out (σ : St S) (b : Bool) : (σ.outB b).out = σ.out ++ [b] := rfl

/-! ### Normalising chains of `set` -/

/-- Advance the input position. -/
def St.adv (σ : St S) : St S := { σ with ip := σ.ip + 1 }

theorem St.adv_set (σ : St S) (i : Fin S) (l : List Sym) : (σ.set i l).adv = σ.adv.set i l := rfl

theorem St.outB_set (σ : St S) (i : Fin S) (l : List Sym) (b : Bool) :
    (σ.set i l).outB b = (σ.outB b).set i l := rfl

/-! ## Moving symbols from one stack to another -/

/-- Transition of `move`: while `g` is defined on the top of `src`, pop it and push its image
on `dst`. -/
def moveδ (q src dst : Fin S) (g : Sym → Option Sym) : List (Option Sym) → List (Op S)
  | [_, some a] =>
    match g a with
    | some b => [.pop src, .push dst b]
    | none => [.pop q, .push q hlt]
  | _ => [.pop q, .push q hlt]

def move (q src dst : Fin S) (g : Sym → Option Sym) : Stmt S P :=
  machine q [src] 0 (moveδ q src dst g)

/-- Cost of `move` on `m` symbols. -/
def moveCost (m : ℕ) : ℕ := (m + 1) * 38 + 3

end NLogN

end

section
-- Source module: NLogN.Add
/-!
# Binary words on stacks; addition

A word `A : List Bool` (least significant bit first) sits on a stack as `A.map codeBit`, so its
least significant bit is on top.  A word ends where the stack stops showing bits.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def bitOf : Option Sym → Option Bool
  | some 0 => some false
  | some 1 => some true
  | _ => none

@[simp] theorem bitOf_zero : bitOf (some 0) = some false := rfl
@[simp] theorem bitOf_one : bitOf (some 1) = some true := rfl

/-- Value of a little-endian word. -/
def wval : List Bool → ℕ
  | [] => 0
  | b :: l => b.toNat + 2 * wval l

section
variable {q a b d e : Fin S} (hqa : q ≠ a) (hqb : q ≠ b) (hqd : q ≠ d) (hqe : q ≠ e)
  (hab : a ≠ b) (had : a ≠ d) (hae : a ≠ e) (hbd : b ≠ d) (hbe : b ≠ e) (hde : d ≠ e)
  (ra rb Q : List Sym) (hra : bitOf ra.head? = none) (hrb : bitOf rb.head? = none)

end

end NLogN

end

section
-- Source module: NLogN.Zip
/-!
# A generic two-input bit transducer with carry

`zip q a b d e keep f fin` runs least-significant-bit first over the words on top of `a` and `b`
(a missing bit reads as `0`), with a one-bit carry `c`.  At each position it pushes the output
bit of `f x y c` on `d` and continues with the new carry; the bits of `b` are copied onto `e`
when `keep`.  When both words are exhausted it pushes the symbols `fin c` on `d`.
Addition, subtraction, negation and copying are instances.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-- Output bits and final carry. -/
def zipL (f : Bool → Bool → Bool → Bool × Bool) : Bool → List Bool → List Bool → List Bool × Bool
  | c, [], [] => ([], c)
  | c, x :: A, [] => let r := zipL f (f x false c).2 A []; ((f x false c).1 :: r.1, r.2)
  | c, [], y :: B => let r := zipL f (f false y c).2 [] B; ((f false y c).1 :: r.1, r.2)
  | c, x :: A, y :: B => let r := zipL f (f x y c).2 A B; ((f x y c).1 :: r.1, r.2)

def zipδ (q a b d e : Fin S) (keep : Bool) (f : Bool → Bool → Bool → Bool × Bool)
    (fin : Bool → List Sym) : List (Option Sym) → List (Op S)
  | [some c, ta, tb] =>
    if bitOf ta = none ∧ bitOf tb = none then
      (fin (decide (c = 1))).map (Op.push d) ++ [.pop q, .push q hlt]
    else
      let x := (bitOf ta).getD false
      let y := (bitOf tb).getD false
      let r := f x y (decide (c = 1))
      (if (bitOf ta).isSome then [.pop a] else []) ++
        (if (bitOf tb).isSome then .pop b :: (if keep then [.push e (codeBit y)] else []) else [])
        ++ [.push d (codeBit r.1), .pop q, .push q (codeBit r.2)]
  | _ => [.pop q, .push q hlt]

def zip (q a b d e : Fin S) (keep : Bool) (f : Bool → Bool → Bool → Bool × Bool)
    (fin : Bool → List Sym) (c0 : Bool) : Stmt S P :=
  machine q [a, b] (codeBit c0) (zipδ q a b d e keep f fin)

def zipCost (B m : ℕ) : ℕ := (m + 1) * (B + 59) + 3

/-- Pushing a list of symbols in order. -/
def pushAll (l : List Sym) (D : List Sym) : List Sym := l.reverse ++ D

section
variable {q a b d e : Fin S} (hqa : q ≠ a) (hqb : q ≠ b) (hqd : q ≠ d) (hqe : q ≠ e)
  (hab : a ≠ b) (had : a ≠ d) (hae : a ≠ e) (hbd : b ≠ d) (hbe : b ≠ e) (hde : d ≠ e)
  (keep : Bool) (f : Bool → Bool → Bool → Bool × Bool) (fin : Bool → List Sym)
  (ra rb Q : List Sym) (hra : bitOf ra.head? = none) (hrb : bitOf rb.head? = none)

def zipCfg (σ : St S) (c : Sym) (A B : List Bool) (D E : List Sym) : St S :=
  ((((σ.set a (A.map codeBit ++ ra)).set b (B.map codeBit ++ rb)).set d D).set e E).set q (c :: Q)

end

end NLogN

end

section
-- Source module: NLogN.Item
/-!
# Integer items on stacks

An item is an integer (sign–magnitude, minimal binary, least significant bit nearest the header)
or a mark.  On a stack an item is its header symbol followed by its magnitude bits; the next
item's header ends the bits.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Minimal binary -/

/-- Minimal little-endian binary representation. -/
def nbits (m : ℕ) : List Bool :=
  if h : m = 0 then [] else decide (m % 2 = 1) :: nbits (m / 2)
decreasing_by omega

/-- Drop high zero bits. -/
def trimL : List Bool → List Bool
  | [] => []
  | b :: l => if trimL l = [] ∧ b = false then [] else b :: trimL l

/-! ## Items -/

def hdPos : Sym := 4
def hdNeg : Sym := 6
def hdMark : Sym := 8

inductive Item where
  | int (z : ℤ)
  | mark
  deriving DecidableEq

def hdOf (z : ℤ) : Sym := if z < 0 then hdNeg else hdPos

/-- Magnitude bits of an integer. -/
def zbits (z : ℤ) : List Bool := nbits z.natAbs

/-- Bit size of an integer. -/
def zsz (z : ℤ) : ℕ := (zbits z).length

def Item.enc : Item → List Sym
  | .int z => hdOf z :: (zbits z).map codeBit
  | .mark => [hdMark]

def encL (L : List Item) : List Sym := L.flatMap Item.enc

@[simp] theorem encL_nil : NLogN.encL [] = [] := rfl

@[simp] theorem bitOf_hdPos : bitOf (some hdPos) = none := rfl
@[simp] theorem bitOf_hdNeg : bitOf (some hdNeg) = none := rfl
@[simp] theorem bitOf_hdMark : bitOf (some hdMark) = none := rfl

end NLogN

namespace NLogN

/-! ## Bit-level functions for `zip` -/

def fAdd (x y c : Bool) : Bool × Bool := (x ^^ y ^^ c, (x && y) || (c && (x ^^ y)))
def fSub (x y c : Bool) : Bool × Bool := (x ^^ y ^^ c, (!x && (y || c)) || (x && y && c))
def fNeg (x _y c : Bool) : Bool × Bool := (x ^^ c, c || x)
/-- Copy the second word. -/
def fCopy (_x y _c : Bool) : Bool × Bool := (y, false)

end NLogN

end

section
-- Source module: NLogN.IntOps
/-!
# Basic operations on integer items
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Moving bits -/

def bitG (s : Sym) : Option Sym := if s = 0 ∨ s = 1 then some s else none

/-- Move the bits on top of `src` to `dst` (reversing them). -/
def moveBits (q src dst : Fin S) : Stmt S P := move q src dst bitG

/-! ## Trimming high zero bits (most significant bit on top) -/

def trimTop (t : Fin S) : Stmt S P := .loop t (fun o => decide (o = some 0)) (.pop t)

/-! ## Popping and moving items -/

def popBits (t : Fin S) : Stmt S P := .loop t (fun o => (bitOf o).isSome) (.pop t)

def popInt (i : Fin S) : Stmt S P := .seq (.pop i) (popBits i)

/-- Move an integer item with header `h` and bits `l`, via the scratch stack `t`. -/
def mvInt (q i j t : Fin S) (h : Sym) : Stmt S P :=
  .seq (.pop i) (.seq (moveBits q i t) (.seq (moveBits q t j) (.push j h)))

/-- Move the integer item on top of `i` onto `j`, via the scratch stack `t`. -/
def moveInt (q i j t : Fin S) : Stmt S P :=
  caseTop i fun o => if o = some hdNeg then mvInt q i j t hdNeg else mvInt q i j t hdPos

end NLogN

end

section
-- Source module: NLogN.Arith
/-!
# Signed addition of integer items
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def finAdd (c : Bool) : List Sym := if c then [1] else []

/-- Add the magnitudes on top of `i` and `j` (popping them) and push the sum with header `h`
on `k`. -/
def addMag (q i j k t1 t2 : Fin S) (h : Sym) : Stmt S P :=
  .seq (zip q i j t1 t2 false fAdd finAdd false)
    (.seq (trimTop t1) (.seq (moveBits q t1 k) (.push k h)))

def addMagCost (L : ℕ) : ℕ := zipCost 1 L + (2 * (L + 1) + 1) + moveCost (L + 1) + 1

section
variable {q i j k t1 t2 : Fin S} (hqi : q ≠ i) (hqj : q ≠ j) (hqt1 : q ≠ t1) (hqt2 : q ≠ t2)
  (hij : i ≠ j) (hit1 : i ≠ t1) (hit2 : i ≠ t2) (hjt1 : j ≠ t1) (hjt2 : j ≠ t2) (ht12 : t1 ≠ t2)
  (hqk : q ≠ k) (hkt1 : k ≠ t1) (hkt2 : k ≠ t2)

end

def flagB : Sym := 9
def finSub (c : Bool) : List Sym := if c then [flagB] else []
def finNone (_c : Bool) : List Sym := []

def isBitT (o : Option Sym) : Bool := (bitOf o).isSome

/-- Subtract the magnitude on top of `j` from that on top of `i` (popping both); push the
absolute difference on `k`, with header `hpos` if it was positive, `hneg` if negative. -/
def subMag (q i j k t1 t2 : Fin S) (hpos hneg : Sym) : Stmt S P :=
  .seq (zip q i j t1 t2 false fSub finSub false)
    (.ite t1 (fun o => decide (o = some flagB))
      (.seq (.pop t1) (.seq (moveBits q t1 t2) (.seq (zip q t2 i t1 j false fNeg finNone false)
        (.seq (trimTop t1) (.seq (moveBits q t1 k) (.push k hneg))))))
      (.seq (trimTop t1) (.ite t1 isBitT (.seq (moveBits q t1 k) (.push k hpos)) (.push k hdPos))))

def subMagCost (L : ℕ) : ℕ :=
  zipCost 1 L + 1 + (1 + (moveCost L + (zipCost 0 L + (2 * L + 1 + (moveCost L + 1)))))

/-- Signed difference as a header and magnitude. -/
def subHd (a b : ℕ) (hpos hneg : Sym) : Sym := if b < a then hpos else if a < b then hneg else hdPos

section
variable {q i j k t1 t2 : Fin S} (hqi : q ≠ i) (hqj : q ≠ j) (hqt1 : q ≠ t1) (hqt2 : q ≠ t2)
  (hij : i ≠ j) (hit1 : i ≠ t1) (hit2 : i ≠ t2) (hjt1 : j ≠ t1) (hjt2 : j ≠ t2) (ht12 : t1 ≠ t2)
  (hqk : q ≠ k) (hkt1 : k ≠ t1) (hkt2 : k ≠ t2)

end

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-- `k := i + j` on integer items (popping both). Scratch stacks `t1`, `t2`. -/
def addInt (q i j k t1 t2 : Fin S) : Stmt S P :=
  caseTop i fun oi => caseTop j fun oj =>
    .seq (.pop i) (.seq (.pop j)
      (if decide (oi = some hdNeg) = decide (oj = some hdNeg) then
        addMag q i j k t1 t2 (if oi = some hdNeg then hdNeg else hdPos)
      else if oi = some hdNeg then subMag q i j k t1 t2 hdNeg hdPos
      else subMag q i j k t1 t2 hdPos hdNeg))

def addIntCost (L : ℕ) : ℕ := addMagCost L + subMagCost L + 36

section
variable {q i j k t1 t2 : Fin S} (hqi : q ≠ i) (hqj : q ≠ j) (hqt1 : q ≠ t1) (hqt2 : q ≠ t2)
  (hij : i ≠ j) (hit1 : i ≠ t1) (hit2 : i ≠ t2) (hjt1 : j ≠ t1) (hjt2 : j ≠ t2) (ht12 : t1 ≠ t2)
  (hqk : q ≠ k) (hkt1 : k ≠ t1) (hkt2 : k ≠ t2)

end

end NLogN

end

section
-- Source module: NLogN.IntOps2
/-!
# More integer item operations: negate, duplicate, shifts
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Negation -/

def negInt (i : Fin S) : Stmt S P :=
  caseTop i fun o => if o = some hdNeg then .seq (.pop i) (.push i hdPos)
    else .seq (.pop i) (.ite i isBitT (.push i hdNeg) (.push i hdPos))

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Duplication -/

def dupInt (q i j t1 t2 : Fin S) : Stmt S P :=
  caseTop i fun o =>
    let h := if o = some hdNeg then hdNeg else hdPos
    .seq (.pop i) (.seq (zip q j i t1 t2 true fCopy finNone false) (.seq (moveBits q t2 i)
      (.seq (.push i h) (.seq (moveBits q t1 j) (.push j h)))))

def dupCost (L : ℕ) : ℕ := zipCost 0 L + 2 * moveCost L + 21

end NLogN

end

section
-- Source module: NLogN.Take
/-!
# Lockstep transfer: `take`

While the top of `b` is a bit, pop it and push it on `e`; at the same time, if the top of `a` is a
bit, move it to `d`.  So `|Bw|` bits are taken from `a`, where `Bw` is the bit word on `b`.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def takeδ (q a b d e : Fin S) : List (Option Sym) → List (Op S)
  | [_, ta, some sb] =>
    if (bitOf (some sb)).isSome then
      (match ta with
        | some sa => if (bitOf (some sa)).isSome then [.pop a, .push d sa] else []
        | none => []) ++ [.pop b, .push e sb]
    else [.pop q, .push q hlt]
  | _ => [.pop q, .push q hlt]

def take (q a b d e : Fin S) : Stmt S P := machine q [a, b] 0 (takeδ q a b d e)

def takeCost (m : ℕ) : ℕ := (m + 1) * 57 + 3

end NLogN

end

section
-- Source module: NLogN.Shift
/-!
# Shifting and splitting integers by a power of two given as an item

A "ruler" is the item `int (2 ^ m)`; its bits are `m` zeros and a one.
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Bit lemmas -/

/-- Integer with the sign of `z` and magnitude `n`. -/
def sg (z : ℤ) (n : ℕ) : ℤ := if z < 0 then -(n : ℤ) else n

/-! ## Moving the zeros of a ruler -/

def g0 (s : Sym) : Option Sym := if s = 0 then some 0 else none

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Shift left by a ruler -/

def shlR (q i j t1 t2 : Fin S) : Stmt S P :=
  caseTop i fun o =>
    let h := if o = some hdNeg then hdNeg else hdPos
    .seq (.pop i) (.ite i isBitT
      (.seq (.pop j) (.seq (move q j t2 g0) (.seq (zip q t1 t2 i j true fCopy finNone false)
        (.seq (.push j hdPos) (.push i h)))))
      (.push i h))

def shlCost (m : ℕ) : ℕ := moveCost m + zipCost 0 m + 23

/-! ## Split by a ruler -/

def splitR (q i j k t1 t2 : Fin S) : Stmt S P :=
  caseTop i fun o =>
    let h := if o = some hdNeg then hdNeg else hdPos
    .seq (.pop i) (.seq (.pop j) (.seq (move q j t2 g0) (.seq (take q i t2 t1 j)
      (.seq (.push j hdPos) (.seq (trimTop t1) (.seq (moveBits q t1 k)
        (.seq (.ite k isBitT (.push k h) (.push k hdPos))
          (.ite i isBitT (.push i h) (.push i hdPos)))))))))

def splitCost (m : ℕ) : ℕ := moveCost m + takeCost m + 2 * m + moveCost m + 27

end NLogN

end

section
-- Source module: NLogN.Io
/-!
# Reading the input, writing the output, fixed-length copies
-/

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def St.setIp (σ : St S) (n : ℕ) : St S := { σ with ip := n }
def St.setOut (σ : St S) (o : List Bool) : St S := { σ with out := o }

@[simp] theorem St.setIp_stk (σ : St S) (n : ℕ) : (σ.setIp n).stk = σ.stk := rfl
@[simp] theorem St.setIp_ip (σ : St S) (n : ℕ) : (σ.setIp n).ip = n := rfl
@[simp] theorem St.setIp_out (σ : St S) (n : ℕ) : (σ.setIp n).out = σ.out := rfl
@[simp] theorem St.setOut_stk (σ : St S) (o) : (σ.setOut o).stk = σ.stk := rfl
@[simp] theorem St.setOut_ip (σ : St S) (o) : (σ.setOut o).ip = σ.ip := rfl
@[simp] theorem St.setOut_out (σ : St S) (o) : (σ.setOut o).out = o := rfl
@[simp] theorem St.adv_setIp (σ : St S) (n : ℕ) : (σ.setIp n).adv = σ.setIp (n + 1) := rfl
@[simp] theorem St.outB_setOut (σ : St S) (o) (b : Bool) :
    (σ.setOut o).outB b = σ.setOut (o ++ [b]) := rfl
@[simp] theorem St.set_ip' (σ : St S) (i : Fin S) (l) : (σ.set i l).ip = σ.ip := rfl
@[simp] theorem St.setIp_self (σ : St S) : σ.setIp σ.ip = σ := rfl
@[simp] theorem St.setOut_self (σ : St S) : σ.setOut σ.out = σ := rfl
@[simp] theorem St.top_setIp (σ : St S) (n : ℕ) (i : Fin S) : (σ.setIp n).top i = σ.top i := rfl
@[simp] theorem St.top_setOut (σ : St S) (o) (i : Fin S) : (σ.setOut o).top i = σ.top i := rfl

attribute [simp] St.adv_set St.outB_set

/-! ## Reading the input up to a stop symbol -/

def readδ (q x : Fin S) (stop : Sym) : List (Option Sym) → List (Op S)
  | [some s, t] =>
    if s = 0 then [.inRd x, .pop q, .push q 1]
    else if t = some stop then [.pop x, .pop q, .push q hlt] else [.inRd x]
  | _ => [.pop q, .push q hlt]

/-- Read input codes onto `x` up to (and consuming) the first `stop`. -/
def readUntil (q x : Fin S) (stop : Sym) : Stmt S P := machine q [x] 0 (readδ q x stop)

def readCost (m : ℕ) : ℕ := (m + 2) * 39 + 3

section read
variable {q x : Fin S} (hqx : q ≠ x) (stop : Sym) (Q : List Sym)

/-- State `1` of `readUntil`: input position `p`, stack `x` holds `X`. -/
def readCfg (σ : St S) (s : Sym) (X : List Sym) (p : ℕ) : St S :=
  (((σ.setIp p).set x X).set q (s :: Q))

end read

/-! ## Writing a word to the output -/

def outδ (q x : Fin S) : List (Option Sym) → List (Op S)
  | [_, t] =>
    match bitOf t with
    | some b => [.pop x, .out b]
    | none => [.pop q, .push q hlt]
  | _ => [.pop q, .push q hlt]

/-- Output the bits on top of `x` (top first), popping them. -/
def outWord (q x : Fin S) : Stmt S P := machine q [x] 0 (outδ q x)

def outCost (m : ℕ) : ℕ := (m + 1) * 38 + 3

section out
variable {q x : Fin S} (hqx : q ≠ x) (r Q : List Sym) (hr : bitOf r.head? = none)

def outCfg (σ : St S) (s : Sym) (A : List Bool) (O : List Bool) : St S :=
  (((σ.setOut O).set x (A.map codeBit ++ r)).set q (s :: Q))

end out

/-! ## Copying a fixed number of bits -/

/-- The first `n` bits of `A`, padded with zeros. -/
def fitL : ℕ → List Bool → List Bool
  | 0, _ => []
  | n + 1, [] => false :: fitL n []
  | n + 1, b :: A => b :: fitL n A

/-- One bit of `src` (or `0` if it shows no bit) per bit popped from `cnt`, pushed onto `dst`. -/
def padδ (q src cnt dst : Fin S) : List (Option Sym) → List (Op S)
  | [_, ts, tc] =>
    match bitOf tc with
    | some _ =>
      match bitOf ts with
      | some b => [.pop cnt, .pop src, .push dst (codeBit b)]
      | none => [.pop cnt, .push dst 0]
    | none => [.pop q, .push q hlt]
  | _ => [.pop q, .push q hlt]

def padMove (q src cnt dst : Fin S) : Stmt S P := machine q [src, cnt] 0 (padδ q src cnt dst)

def padCost (m : ℕ) : ℕ := (m + 1) * 56 + 3

section pad
variable {q a c d : Fin S} (hqa : q ≠ a) (hqc : q ≠ c) (hqd : q ≠ d) (hac : a ≠ c) (had : a ≠ d)
  (hcd : c ≠ d) (ra rc Q : List Sym) (hra : bitOf ra.head? = none) (hrc : bitOf rc.head? = none)

def padCfg (σ : St S) (s : Sym) (A C : List Bool) (D : List Sym) : St S :=
  (((σ.set a (A.map codeBit ++ ra)).set c (C.map codeBit ++ rc)).set d D).set q (s :: Q)

end pad

end NLogN

end

section
-- Source module: NLogN.Flat
/-!
# Flat code with jumps, and compilation of structured programs into it

Instructions carry absolute addresses. `ret` with an empty control stack halts.
`exec_flat`: a structured run of cost `c` is `c` flat steps.
-/

namespace NLogN

inductive Instr (S : ℕ) where
  | jmp (e : ℕ)
  | push (i : Fin S) (a : Sym) (e : ℕ)
  | pop (i : Fin S) (e : ℕ)
  | inRd (i : Fin S) (e : ℕ)
  | out (b : Bool) (e : ℕ)
  | br (i : Fin S) (f : Option Sym → Bool) (y n : ℕ)
  | call (t e : ℕ)
  | ret

/-- Flat configuration: program counter, state, control stack of return addresses. -/
structure Cfg (S : ℕ) where
  pc : ℕ
  σ : St S
  ctl : List ℕ

variable {S P : ℕ}

/-- One step; `none` when halted (`ret` on an empty control stack). -/
def fstep (I : List Sym) (code : ℕ → Instr S) (c : Cfg S) : Option (Cfg S) :=
  match code c.pc with
  | .jmp e => some ⟨e, c.σ, c.ctl⟩
  | .push i a e => some ⟨e, c.σ.push i a, c.ctl⟩
  | .pop i e => some ⟨e, c.σ.pop i, c.ctl⟩
  | .inRd i e => some ⟨e, c.σ.inRd I i, c.ctl⟩
  | .out b e => some ⟨e, c.σ.outB b, c.ctl⟩
  | .br i f y n => some ⟨if f (c.σ.top i) then y else n, c.σ, c.ctl⟩
  | .call t e => some ⟨t, c.σ, e :: c.ctl⟩
  | .ret => match c.ctl with
    | [] => none
    | e :: l => some ⟨e, c.σ, l⟩

/-- `n` steps from one configuration to another. -/
inductive Steps (I : List Sym) (code : ℕ → Instr S) : Cfg S → ℕ → Cfg S → Prop
  | refl (c) : Steps I code c 0 c
  | step {c c' c'' n} : fstep I code c = some c' → Steps I code c' n c'' → Steps I code c (n + 1) c''

/-! ## Layout -/

/-- Number of instructions of a statement. -/
def Stmt.size : Stmt S P → ℕ
  | .skip => 1
  | .push _ _ => 1
  | .pop _ => 1
  | .inRd _ => 1
  | .out _ => 1
  | .seq p q => p.size + q.size
  | .ite _ _ p q => 1 + p.size + q.size
  | .loop _ _ p => 1 + p.size
  | .call _ => 1

/-- `Placed code s a e ent`: the code of `s` sits at `[a, a + s.size)` and exits to `e`;
procedure `j` starts at `ent j`. -/
def Placed (code : ℕ → Instr S) (ent : Fin P → ℕ) : Stmt S P → ℕ → ℕ → Prop
  | .skip, a, e => code a = .jmp e
  | .push i x, a, e => code a = .push i x e
  | .pop i, a, e => code a = .pop i e
  | .inRd i, a, e => code a = .inRd i e
  | .out b, a, e => code a = .out b e
  | .seq p q, a, e => Placed code ent p a (a + p.size) ∧ Placed code ent q (a + p.size) e
  | .ite i f p q, a, e => code a = .br i f (a + 1) (a + 1 + p.size) ∧
      Placed code ent p (a + 1) e ∧ Placed code ent q (a + 1 + p.size) e
  | .loop i f p, a, e => code a = .br i f (a + 1) e ∧ Placed code ent p (a + 1) a
  | .call j, a, e => code a = .call (ent j) e

/-- The procedures are placed at their entries, each followed by `ret`. -/
def ProgPlaced (code : ℕ → Instr S) (prog : Fin P → Stmt S P) (ent : Fin P → ℕ) : Prop :=
  ∀ j, Placed code ent (prog j) (ent j) (ent j + (prog j).size) ∧
    code (ent j + (prog j).size) = .ret

/-! ## A concrete layout -/

/-- The instruction at address `x` of `s` placed at `a` with exit `e`, if `x` is inside. -/
def Stmt.instrAt (ent : Fin P → ℕ) : Stmt S P → ℕ → ℕ → ℕ → Option (Instr S)
  | .skip, a, e, x => if x = a then some (.jmp e) else none
  | .push i y, a, e, x => if x = a then some (.push i y e) else none
  | .pop i, a, e, x => if x = a then some (.pop i e) else none
  | .inRd i, a, e, x => if x = a then some (.inRd i e) else none
  | .out b, a, e, x => if x = a then some (.out b e) else none
  | .seq p q, a, e, x =>
    if x < a + p.size then p.instrAt ent a (a + p.size) x else q.instrAt ent (a + p.size) e x
  | .ite i f p q, a, e, x =>
    if x = a then some (.br i f (a + 1) (a + 1 + p.size))
    else if x < a + 1 + p.size then p.instrAt ent (a + 1) e x
    else q.instrAt ent (a + 1 + p.size) e x
  | .loop i f p, a, e, x =>
    if x = a then some (.br i f (a + 1) e) else p.instrAt ent (a + 1) a x
  | .call j, a, e, x => if x = a then some (.call (ent j) e) else none

/-- Entry of procedure `j`: procedures are laid out one after another, each followed by `ret`. -/
def entryOf (prog : Fin P → Stmt S P) (j : ℕ) : ℕ :=
  ((List.range j).map fun i => if h : i < P then (prog ⟨i, h⟩).size + 1 else 0).sum

theorem entryOf_succ (prog : Fin P → Stmt S P) (j : Fin P) :
    entryOf prog (j.val + 1) = entryOf prog j + (prog j).size + 1 := by
  simp [entryOf, List.range_succ, j.isLt, Nat.add_assoc]

theorem entryOf_mono (prog : Fin P → Stmt S P) {i j : ℕ} (h : i ≤ j) :
    entryOf prog i ≤ entryOf prog j := by
  induction h with
  | refl => exact le_refl _
  | step _ ih =>
    refine ih.trans ?_
    simp [entryOf, List.range_succ]

/-- Total code length. -/
def codeLen (prog : Fin P → Stmt S P) : ℕ := entryOf prog P

/-- The procedure whose code contains address `x`. -/
def procAt (prog : Fin P → Stmt S P) (x : ℕ) : Option (Fin P) :=
  (List.finRange P).find? (fun j : Fin P => decide (entryOf prog j ≤ x ∧
    x < entryOf prog j + (prog j).size))

/-- The flat code of a program. Addresses outside the code hold `ret`. -/
def codeOf (prog : Fin P → Stmt S P) (x : ℕ) : Instr S :=
  match procAt prog x with
  | some j => ((prog j).instrAt (fun j => entryOf prog j) (entryOf prog j)
      (entryOf prog j + (prog j).size) x).getD .ret
  | none => .ret

/-- Addresses used by the code are below `codeLen`. -/
theorem entryOf_le_codeLen (prog : Fin P → Stmt S P) (j : Fin P) :
    entryOf prog j + (prog j).size + 1 ≤ codeLen prog := by
  rw [← entryOf_succ]; exact entryOf_mono prog j.isLt

end NLogN

end

section
-- Source module: NLogN.Basic
/-!
# A relaxed machine layer over `IntMul.MultitapeTM`

`Prog Γ k` is a deterministic `k`-tape machine over an arbitrary alphabet whose transition
function returns `none` to halt.  Its step semantics silently enforce the platform rules
(`▷` is kept and never written elsewhere, no left move from `▷`, tape `0` is read-only), so
`compile` turns any program into a `MultitapeTM` with one extra step.
-/

namespace IntMul.Mach

open IntMul

variable {Γ : Type} {k : ℕ}

/-- Tape contents and head positions. -/
structure Tapes (Γ : Type) (k : ℕ) where
  cells : Fin k → ℕ → Γ
  head : Fin k → ℕ

/-- Head movement, in the same shape as `MultitapeTM.step`. -/
def moveHead (m : Move) (h : ℕ) : ℕ :=
  match m with
  | Move.left => h - 1
  | Move.stay => h
  | Move.right => h + 1

namespace Tapes

/-- The scanned symbols. -/
def read (T : Tapes Γ k) : Fin k → Γ := fun i => T.cells i (T.head i)

/-- Write and move on every tape. -/
def write (T : Tapes Γ k) (w : Fin k → Γ × Move) : Tapes Γ k where
  cells i := Function.update (T.cells i) (T.head i) (w i).1
  head i := moveHead (w i).2 (T.head i)

end Tapes

/-- The action actually performed: `▷` is kept and never moved left from, `▷` is never written
elsewhere, and tape `0` is read-only. -/
def fix [DecidableEq Γ] (st : Γ) (a : Fin k → Γ) (w : Fin k → Γ × Move) : Fin k → Γ × Move :=
  fun i =>
    if a i = st then (st, if (w i).2 = Move.left then Move.stay else (w i).2)
    else ((if i.val = 0 ∨ (w i).1 = st then a i else (w i).1), (w i).2)

/-- A relaxed program: `δ` returns `none` to halt. -/
structure Prog (Γ : Type) (k : ℕ) where
  Q : Type
  [fintypeQ : Fintype Q]
  q0 : Q
  δ : Q → (Fin k → Γ) → Option (Q × (Fin k → Γ × Move))

attribute [instance] Prog.fintypeQ

section Sem

variable [DecidableEq Γ] (st : Γ)

/-- `Halts st P q T t T'`: started in state `q` on tapes `T`, the program halts after exactly
`t` steps with tapes `T'`. -/
inductive Halts (P : Prog Γ k) : P.Q → Tapes Γ k → ℕ → Tapes Γ k → Prop
  | done {q : P.Q} {T : Tapes Γ k} : P.δ q T.read = none → Halts P q T 0 T
  | step {q q' : P.Q} {T T' : Tapes Γ k} {w : Fin k → Γ × Move} {t : ℕ} :
      P.δ q T.read = some (q', w) → Halts P q' (T.write (fix st T.read w)) t T' →
        Halts P q T (t + 1) T'

/-- Run from the start state. -/
def Runs (P : Prog Γ k) (T : Tapes Γ k) (t : ℕ) (T' : Tapes Γ k) : Prop :=
  Halts st P P.q0 T t T'

variable {st}

end Sem

/-! ## Combinators -/

section CombSem

variable [DecidableEq Γ] {st : Γ}

end CombSem

/-! ## Compilation to the platform model -/

/-- The five distinguished symbols of the platform model. -/
structure Syms (Γ : Type) where
  blank : Γ
  st : Γ
  zero : Γ
  one : Γ
  sep : Γ
  nodup : [blank, st, zero, one, sep].Nodup

variable [DecidableEq Γ] [Fintype Γ]

/-- The platform machine running `P`; one extra step moves to `HALT`. -/
def compile (S : Syms Γ) (hk : 2 ≤ k) (P : Prog Γ k) : MultitapeTM where
  Sym := Γ
  blank := S.blank
  startSym := S.st
  zero := S.zero
  one := S.one
  sep := S.sep
  syms_distinct := S.nodup
  K := Option P.Q
  qStart := some P.q0
  qHalt := none
  start_ne_halt := by simp
  k := k
  two_le_k := hk
  δ q a := match q with
    | none => (none, fun i => (a i, Move.stay))
    | some q => match P.δ q a with
      | none => (none, fun i => (a i, Move.stay))
      | some r => (some r.1, fix S.st a r.2)
  start_preserved := by
    intro q a i hi
    rcases q with _ | q
    · simp [hi]
    · rcases h : P.δ q a with _ | r <;> simp only [h]
      · simp [hi]
      · simp only [fix, hi, if_true]
        split_ifs with hm <;> simp_all
  start_only_at_start := by
    intro q a i hi
    rcases q with _ | q
    · simpa using hi
    · rcases h : P.δ q a with _ | r <;> simp only [h]
      · simpa using hi
      · simp only [fix, hi, if_false]
        split_ifs with h'
        · exact hi
        · simp_all
  halt_fixed := fun _ => rfl
  input_readonly := by
    intro q a
    rcases q with _ | q
    · rfl
    · rcases h : P.δ q a with _ | r <;> simp only [h]
      simp only [fix]
      split_ifs <;> simp_all

/-- A configuration of the compiled machine. -/
def toCfg (S : Syms Γ) (hk : 2 ≤ k) (P : Prog Γ k) (q : Option P.Q) (T : Tapes Γ k) :
    (compile S hk P).Cfg :=
  ⟨q, T.cells, T.head⟩

/-- The platform's initial tapes on input `x#y`. -/
def initTapes (S : Syms Γ) (hk : 2 ≤ k) (x y : List Bool) : Tapes Γ k :=
  ⟨fun i => if i.val = 0 then
      (fun p => match p with
        | 0 => S.st
        | p + 1 => ((x.map (fun b => if b then S.one else S.zero)) ++
            S.sep :: y.map (fun b => if b then S.one else S.zero)).getD p S.blank)
      else (fun p => match p with
        | 0 => S.st
        | _ + 1 => S.blank),
   fun _ => 0⟩

/-- The tape `▷ w □ □ …`. -/
def tapeOf (S : Syms Γ) (w : List Γ) : ℕ → Γ
  | 0 => S.st
  | p + 1 => w.getD p S.blank

/-- Encoding of a bit. -/
def bitSym (S : Syms Γ) (b : Bool) : Γ := if b then S.one else S.zero

end IntMul.Mach

end

section
-- Source module: NLogN.Walk
/-!
# Scans, copies and finite-state walks

Generic single-loop programs and their exact running times.
-/

namespace IntMul.Mach

open IntMul

set_option linter.unusedSectionVars false

variable {Γ : Type} {k : ℕ} [DecidableEq Γ]

/-! ## Tape lemmas -/

/-! ## Writing with `fix` -/

/-- The write vector that changes tape `i` only. -/
def wr1 (a : Fin k → Γ) (i : Fin k) (s : Γ) (m : Move) : Fin k → Γ × Move :=
  fun l => if l = i then (s, m) else (a l, Move.stay)

/-- The write vector that changes tapes `i` and `j` only. -/
def wr2 (a : Fin k → Γ) (i : Fin k) (s : Γ) (m : Move) (j : Fin k) (s' : Γ) (m' : Move) :
    Fin k → Γ × Move :=
  fun l => if l = j then (s', m') else if l = i then (s, m) else (a l, Move.stay)

/-! ## Scans -/

/-! ## Copy -/

/-! ## Finite-state walks to the left -/

/-! ## Single moves -/

end IntMul.Mach

end

section
-- Source module: NLogN.TM
/-!
# Running flat code on the platform's Turing machines

Tapes: `0` input, `1` output, `2` control stack (return addresses), `3 + i` stack `i`.
On every stack tape the head rests on the top symbol (on `▷` when empty), so `push` is two
steps, `pop` and tests are one step, and a call is two steps.  Each flat step is at most two
machine steps (`flat_tm`).
-/

namespace NLogN

open IntMul IntMul.Mach

variable {S L : ℕ}

/-- Tape alphabet: `none` is `▷`, `some none` is the blank. -/
abbrev Γ (L : ℕ) := Option (Option (Sym ⊕ Fin L))

def gst : Γ L := none
def gbl : Γ L := some none
def gsym (a : Sym) : Γ L := some (some (Sum.inl a))
def gret (x : Fin L) : Γ L := some (some (Sum.inr x))

def syms (L : ℕ) : Syms (Γ L) :=
  ⟨gbl, gst, gsym 0, gsym 1, gsym 2, by simp [gbl, gst, gsym]⟩

/-- Return address symbol (blank for an out-of-range address, which never occurs). -/
def retSym (e : ℕ) : Γ L := if h : e < L then gret ⟨e, h⟩ else gbl

def topOf : Γ L → Option Sym
  | some (some (Sum.inl a)) => some a
  | _ => none

def inCode : Γ L → Sym
  | some (some (Sum.inl a)) => a
  | _ => codeEnd

def toPc (e : ℕ) : Option (Fin L) := if h : e < L then some ⟨e, h⟩ else none

def tin : Fin (S + 3) := ⟨0, by omega⟩
def tout : Fin (S + 3) := ⟨1, by omega⟩
def tctl : Fin (S + 3) := ⟨2, by omega⟩
def tst (i : Fin S) : Fin (S + 3) := ⟨i.val + 3, by omega⟩

/-- No change on any tape. -/
def stayAll {k : ℕ} (a : Fin k → Γ L) : Fin k → Γ L × Move := fun l => (a l, Move.stay)

/-- The transition function. State `(pc, false)` starts an instruction, `(pc, true)` is the
second step of a two-step instruction. -/
def δf (code : ℕ → Instr S) (q : Fin L × Bool) (a : Fin (S + 3) → Γ L) :
    Option ((Fin L × Bool) × (Fin (S + 3) → Γ L × Move)) :=
  match code q.1, q.2 with
  | .jmp e, _ => (toPc e).map fun e => ((e, false), stayAll a)
  | .push i _ _, false => some ((q.1, true), wr1 a (tst i) (a (tst i)) Move.right)
  | .push i x e, true => (toPc e).map fun e => ((e, false), wr1 a (tst i) (gsym x) Move.stay)
  | .pop i e, _ => (toPc e).map fun e => ((e, false), wr1 a (tst i) gbl Move.left)
  | .inRd i _, false => some ((q.1, true), wr2 a tin (a tin) Move.right (tst i) (a (tst i)) Move.right)
  | .inRd i e, true => (toPc e).map fun e => ((e, false), wr1 a (tst i) (gsym (inCode (a tin))) Move.stay)
  | .out _ _, false => some ((q.1, true), wr1 a tout (a tout) Move.right)
  | .out b e, true => (toPc e).map fun e =>
      ((e, false), wr1 a tout (gsym (if b then 1 else 0)) Move.stay)
  | .br i f y n, _ => (toPc (if f (topOf (a (tst i))) then y else n)).map fun e => ((e, false), stayAll a)
  | .call _ _, false => some ((q.1, true), wr1 a tctl (a tctl) Move.right)
  | .call t e, true => (toPc t).map fun t => ((t, false), wr1 a tctl (retSym e) Move.stay)
  | .ret, _ => match a tctl with
    | some (some (Sum.inr x)) => some ((x, false), wr1 a tctl gbl Move.left)
    | _ => none

def progOf (code : ℕ → Instr S) (L : ℕ) (pc0 : Fin L) : Prog (Γ L) (S + 3) where
  Q := Fin L × Bool
  q0 := (pc0, false)
  δ := δf code

/-! ## Encoding of configurations -/

/-- Tapes of a flat configuration; `inp` is the input tape. -/
def enc (inp : ℕ → Γ L) (c : Cfg S) : Tapes (Γ L) (S + 3) where
  cells l := if l.val = 0 then inp else if l.val = 1 then tapeOf (syms L) (c.σ.out.map (bitSym (syms L)))
    else if l.val = 2 then tapeOf (syms L) (c.ctl.map retSym).reverse
    else if h : l.val - 3 < S then tapeOf (syms L) ((c.σ.stk ⟨l.val - 3, h⟩).map gsym).reverse
    else tapeOf (syms L) []
  head l := if l.val = 0 then c.σ.ip else if l.val = 1 then c.σ.out.length
    else if l.val = 2 then c.ctl.length
    else if h : l.val - 3 < S then (c.σ.stk ⟨l.val - 3, h⟩).length else 0

end NLogN

namespace NLogN

open IntMul IntMul.Mach

variable {S L : ℕ}

theorem tst_ne_tin (i : Fin S) : tst i ≠ tin := by simp [tst, tin, Fin.ext_iff]
theorem tst_ne_tout (i : Fin S) : tst i ≠ tout := by simp [tst, tout, Fin.ext_iff]
theorem tst_ne_tctl (i : Fin S) : tst i ≠ tctl := by simp [tst, tctl, Fin.ext_iff]

theorem tin_ne_tout : (tin : Fin (S + 3)) ≠ tout := by simp [tin, tout, Fin.ext_iff]
theorem tin_ne_tctl : (tin : Fin (S + 3)) ≠ tctl := by simp [tin, tctl, Fin.ext_iff]
theorem tout_ne_tctl : (tout : Fin (S + 3)) ≠ tctl := by simp [tout, tctl, Fin.ext_iff]

/-! ## Partial runs -/

/-- `t` steps of a program from one state to another (no halting in between). -/
inductive Reach {Γ' : Type} [DecidableEq Γ'] {k : ℕ} (st : Γ') (P : Prog Γ' k) :
    P.Q → Tapes Γ' k → ℕ → P.Q → Tapes Γ' k → Prop
  | refl (q T) : Reach st P q T 0 q T
  | step {q q' q'' : P.Q} {T T'' : Tapes Γ' k} {w : Fin k → Γ' × Move} {n : ℕ} :
      P.δ q T.read = some (q', w) → Reach st P q' (T.write (fix st T.read w)) n q'' T'' →
      Reach st P q T (n + 1) q'' T''

attribute [simp] tst_ne_tin tst_ne_tout tst_ne_tctl tin_ne_tout tin_ne_tctl tout_ne_tctl

@[simp] theorem tst_val (i : Fin S) : (tst i).val = i.val + 3 := rfl
@[simp] theorem tout_val : (tout : Fin (S + 3)).val = 1 := rfl
@[simp] theorem tctl_val : (tctl : Fin (S + 3)).val = 2 := rfl

/-! ## One flat step -/

section Sim

variable (code : ℕ → Instr S) (L) (pc0 : Fin L) (inp : ℕ → Γ L)

/-- All addresses in an instruction are below `L`. -/
def Instr.Valid (L : ℕ) : Instr S → Prop
  | .jmp e => e < L
  | .push _ _ e => e < L
  | .pop _ e => e < L
  | .inRd _ e => e < L
  | .out _ e => e < L
  | .br _ _ y n => y < L ∧ n < L
  | .call t e => t < L ∧ e < L
  | .ret => True

end Sim

end NLogN

end

section
-- Source module: NLogN.Top
/-!
# From stack programs to the platform model

`multipliesAt_of_exec`: if procedure `0` of a stack program, run on input `x # y`, writes
`bin (2n) (val x · val y)` with cost `c`, then the compiled platform machine multiplies at `n`
within `2c + 1` steps.
-/

namespace NLogN

open IntMul IntMul.Mach

variable {S P : ℕ}

/-- The initial state: all stacks empty. -/
def St.init (S : ℕ) : St S := ⟨fun _ => [], 0, []⟩

theorem codeLen_pos (prog : Fin P → Stmt S P) (hP : 0 < P) : 0 < codeLen prog := by
  have := entryOf_le_codeLen prog ⟨0, hP⟩; omega

/-- The platform machine of a stack program (procedure `0` is the main program). -/
def machineOf (prog : Fin P → Stmt S P) (hP : 0 < P) : MultitapeTM :=
  compile (syms (codeLen prog)) (by omega : 2 ≤ S + 3)
    (progOf (codeOf prog) (codeLen prog) ⟨0, codeLen_pos prog hP⟩)

end NLogN

end

section
-- Source module: NLogN.Prims
/-!
# Constants, input and output of integer items
-/

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

/-! ## Constants -/

def constI (i : Fin S) (c : ℤ) : Stmt S P :=
  opsStmt ((((zbits c).map codeBit).reverse ++ [hdOf c]).map (Op.push i))

/-! ## Output -/

def outR (q i j t1 : Fin S) : Stmt S P :=
  .seq (.pop i) (.seq (.pop j) (.seq (padMove q i j t1) (.seq (.pop t1) (outWord q t1))))

def outRCost (m : ℕ) : ℕ := padCost (m + 1) + outCost m + 3

/-! ## Input -/

def fZero (_x _y _c : Bool) : Bool × Bool := (false, false)

/-- Read a bit word `W` (most significant bit first) up to `stop`; push `val W` on `i` and the
ruler `2 ^ |W|` on `j`. -/
def readR (q i j t1 t2 t3 : Fin S) (stop : Sym) : Stmt S P :=
  .seq (readUntil q t1 stop) (.seq (.push j 1) (.seq (zip q t2 t1 j t3 true fZero finNone false)
    (.seq (.push j hdPos) (.seq (trimTop t3) (.seq (moveBits q t3 i) (.push i hdPos))))))

def readRCost (m : ℕ) : ℕ := readCost m + zipCost 0 m + 2 * m + moveCost m + 5

end NLogN

end

section
-- Source module: NLogN.Cat
/-!
# Concatenation: `a * 2^m + b` for `0 ≤ b < 2^m`, in time `O(m)`

`catR q i j k t1 t2`: `i` holds `int a` with `a ≠ 0`, `j` the ruler `int (2^m)`, `k` holds `int b`
with `0 ≤ b < 2^m`.  The result `int (sg a (|a| * 2^m + b))` replaces `a`; `b` is consumed and the
ruler is kept.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

variable {S P : ℕ} {I : List Sym} {prog : Fin P → Stmt S P}

def fFst (x _y _c : Bool) : Bool × Bool := (x, false)

def padL (l : List Bool) (m : ℕ) : List Bool := l ++ List.replicate (m - l.length) false

def catR (q i j k t1 t2 : Fin S) : Stmt S P :=
  caseTop i fun o =>
    let h := if o = some hdNeg then hdNeg else hdPos
    .seq (.pop i) (.seq (.pop j) (.seq (move q j t2 g0) (.seq (.pop k)
      (.seq (zip q k t2 t1 j true fFst finNone false) (.seq (.push j hdPos)
        (.seq (moveBits q t1 i) (.push i h)))))))

def catCost (m : ℕ) : ℕ := 2 * moveCost m + zipCost 0 m + 23

end NLogN

end