Harvey–van der Hoeven machine development, part 2: item language, expressions and RM programs

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

import Definitions.Def_IntMul_HvdH_M_Defs1

set_option autoImplicit false
set_option relaxedAutoImplicit false
set_option maxHeartbeats 1000000

section
-- Source module: NLogN.ILang
/-!
# The item language

Programs over `S` stacks of items (integers and marks), compiled to stack programs over `S + 4`
stacks (the last four are scratch).  `compile_correct` turns item-level executions into
`Runs` of the compiled program, with the same cost.
-/

namespace NLogN

open IntMul

variable {S P : ℕ}

structure IState (S : ℕ) where
  G : Fin S → List Item
  ip : ℕ
  out : List Bool

def IState.upd (s : IState S) (i : Fin S) (L : List Item) : IState S :=
  { s with G := Function.update s.G i L }

@[simp] theorem IState.upd_G (s : IState S) (i : Fin S) (L : List Item) (j : Fin S) :
    (s.upd i L).G j = if j = i then L else s.G j := by
  simp [IState.upd, Function.update_apply]

@[simp] theorem IState.upd_ip (s : IState S) (i : Fin S) (L : List Item) : (s.upd i L).ip = s.ip :=
  rfl
@[simp] theorem IState.upd_out (s : IState S) (i : Fin S) (L : List Item) :
    (s.upd i L).out = s.out := rfl

/-! ## Stacks of the compiled program -/

def ι (i : Fin S) : Fin (S + 4) := Fin.castAdd 4 i
def sq (S : ℕ) : Fin (S + 4) := ⟨S, by omega⟩
def s1 (S : ℕ) : Fin (S + 4) := ⟨S + 1, by omega⟩
def s2 (S : ℕ) : Fin (S + 4) := ⟨S + 2, by omega⟩
def s3 (S : ℕ) : Fin (S + 4) := ⟨S + 3, by omega⟩

/-- The machine state of an item state: scratch stacks empty. -/
def encI (s : IState S) : St (S + 4) :=
  ⟨fun x => if h : x.val < S then NLogN.encL (s.G ⟨x.val, h⟩) else [], s.ip, s.out⟩

@[simp] theorem enc_ip (s : IState S) : (encI s).ip = s.ip := rfl
@[simp] theorem enc_out (s : IState S) : (encI s).out = s.out := rfl

@[simp] theorem bitOf_none' : bitOf none = none := rfl

/-! ## Operations -/

inductive IOp (S : ℕ) where
  | add (i j k : Fin S)
  | neg (i : Fin S)
  | dup (i j : Fin S)
  | drop (i : Fin S)
  | mv (i j : Fin S)
  | const (i : Fin S) (c : ℤ)
  | mark (i : Fin S)
  | unmark (i : Fin S)
  | shl (i j : Fin S)
  | split (i j k : Fin S)
  | read (i j : Fin S) (stop : Sym)
  | write (i j : Fin S)
  | cat (i j k : Fin S)

def IOp.stmt : IOp S → Stmt (S + 4) P
  | .add i j k => addInt (sq S) (ι i) (ι j) (ι k) (s1 S) (s2 S)
  | .neg i => negInt (ι i)
  | .dup i j => dupInt (sq S) (ι i) (ι j) (s1 S) (s2 S)
  | .drop i => popInt (ι i)
  | .mv i j => moveInt (sq S) (ι i) (ι j) (s1 S)
  | .const i c => constI (ι i) c
  | .mark i => .push (ι i) hdMark
  | .unmark i => .pop (ι i)
  | .shl i j => shlR (sq S) (ι i) (ι j) (s1 S) (s2 S)
  | .split i j k => splitR (sq S) (ι i) (ι j) (ι k) (s1 S) (s2 S)
  | .read i j stop => readR (sq S) (ι i) (ι j) (s1 S) (s2 S) (s3 S) stop
  | .write i j => outR (sq S) (ι i) (ι j) (s1 S)
  | .cat i j k => catR (sq S) (ι i) (ι j) (ι k) (s1 S) (s2 S)

/-- Item-level semantics of the operations, with costs. -/
inductive IOp.Step (I : List Sym) : IOp S → IState S → IState S → ℕ → Prop
  | add {s : IState S} {i j k : Fin S} {a b : ℤ} {Li Lj : List Item} (hij : i ≠ j)
      (hi : s.G i = .int a :: Li) (hj : s.G j = .int b :: Lj) :
      Step I (.add i j k) s
        (((s.upd i Li).upd j Lj).upd k (.int (a + b) :: ((s.upd i Li).upd j Lj).G k))
        (addIntCost (zsz a + zsz b))
  | neg {s : IState S} {i : Fin S} {a : ℤ} {L : List Item} (hi : s.G i = .int a :: L) :
      Step I (.neg i) s (s.upd i (.int (-a) :: L)) 20
  | dup {s : IState S} {i j : Fin S} {a : ℤ} {L : List Item} (hij : i ≠ j)
      (hi : s.G i = .int a :: L) :
      Step I (.dup i j) s (s.upd j (.int a :: s.G j)) (dupCost (zsz a) + 17)
  | drop {s : IState S} {i : Fin S} {a : ℤ} {L : List Item} (hi : s.G i = .int a :: L) :
      Step I (.drop i) s (s.upd i L) (2 * zsz a + 2)
  | mv {s : IState S} {i j : Fin S} {a : ℤ} {L : List Item} (hij : i ≠ j)
      (hi : s.G i = .int a :: L) :
      Step I (.mv i j) s ((s.upd i L).upd j (.int a :: s.G j)) (2 * moveCost (zsz a) + 19)
  | const {s : IState S} {i : Fin S} (c : ℤ) :
      Step I (.const i c) s (s.upd i (.int c :: s.G i)) (zsz c + 2)
  | mark {s : IState S} {i : Fin S} : Step I (.mark i) s (s.upd i (.mark :: s.G i)) 1
  | unmark {s : IState S} {i : Fin S} {L : List Item} (hi : s.G i = .mark :: L) :
      Step I (.unmark i) s (s.upd i L) 1
  | shl {s : IState S} {i j : Fin S} {a : ℤ} {m : ℕ} {Li Lj : List Item} (hij : i ≠ j)
      (hi : s.G i = .int a :: Li) (hj : s.G j = .int (2 ^ m) :: Lj) :
      Step I (.shl i j) s (s.upd i (.int (a * 2 ^ m) :: Li)) (shlCost m)
  | split {s : IState S} {i j k : Fin S} {a : ℤ} {m : ℕ} {Li Lj : List Item} (hij : i ≠ j)
      (hik : i ≠ k) (hjk : j ≠ k) (hi : s.G i = .int a :: Li) (hj : s.G j = .int (2 ^ m) :: Lj) :
      Step I (.split i j k) s
        ((s.upd i (.int (sg a (a.natAbs / 2 ^ m)) :: Li)).upd k
          (.int (sg a (a.natAbs % 2 ^ m)) :: s.G k)) (splitCost m + 17)
  | read {s : IState S} {i j : Fin S} {stop : Sym} (W : List Bool) (hij : i ≠ j) (hs0 : stop ≠ 0)
      (hs1 : stop ≠ 1) (hk : ∀ k (hk : k < W.length), I.getD (s.ip + k) codeEnd = codeBit W[k])
      (hstop : I.getD (s.ip + W.length) codeEnd = stop) :
      Step I (.read i j stop) s
        { (s.upd i (.int (val W : ℤ) :: s.G i)).upd j
            (.int (2 ^ W.length) :: (s.upd i (.int (val W : ℤ) :: s.G i)).G j) with
          ip := s.ip + W.length + 1 }
        (readRCost W.length)
  | write {s : IState S} {i j : Fin S} {n m : ℕ} {Li Lj : List Item} (hij : i ≠ j)
      (hn : n < 2 ^ m) (hi : s.G i = .int (n : ℤ) :: Li) (hj : s.G j = .int (2 ^ m) :: Lj) :
      Step I (.write i j) s { (s.upd i Li).upd j Lj with out := s.out ++ bin m n } (outRCost m)
  | cat {s : IState S} {i j k : Fin S} {a : ℤ} {b m : ℕ} {Li Lj Lk : List Item} (hij : i ≠ j)
      (hik : i ≠ k) (hjk : j ≠ k) (ha : a ≠ 0) (hb : b < 2 ^ m) (hi : s.G i = .int a :: Li)
      (hj : s.G j = .int (2 ^ m) :: Lj) (hk : s.G k = .int (b : ℤ) :: Lk) :
      Step I (.cat i j k) s ((s.upd i (.int (sg a (a.natAbs * 2 ^ m + b)) :: Li)).upd k Lk)
        (catCost m)

variable {I : List Sym} {prog : Fin P → Stmt (S + 4) P}

end NLogN

end

section
-- Source module: NLogN.IComp
/-!
# Control flow of the item language and the compiler correctness theorem
-/

namespace NLogN

open IntMul

variable {S P : ℕ}

/-- Classes of the top of an item stack, for tests. -/
inductive Cls where
  | none | mark | zero | pos | neg
  deriving DecidableEq

def cls : Option Item → Cls
  | none => .none
  | some .mark => .mark
  | some (.int z) => if z = 0 then .zero else if z < 0 then .neg else .pos

inductive IStmt (S P : ℕ) where
  | skip
  | op (o : IOp S)
  | seq (p q : IStmt S P)
  | ite (i : Fin S) (f : Cls → Bool) (p q : IStmt S P)
  | loop (i : Fin S) (f : Cls → Bool) (p : IStmt S P)
  | call (j : Fin P)

def IState.cl (s : IState S) (i : Fin S) : Cls := cls (s.G i).head?

inductive IExec (I : List Sym) (iprog : Fin P → IStmt S P) :
    IStmt S P → IState S → IState S → ℕ → Prop
  | skip (s) : IExec I iprog .skip s s 1
  | op {o s s' c} : o.Step I s s' c → IExec I iprog (.op o) s s' c
  | seq {p q s s₁ s₂ c₁ c₂} : IExec I iprog p s s₁ c₁ → IExec I iprog q s₁ s₂ c₂ →
      IExec I iprog (.seq p q) s s₂ (c₁ + c₂)
  | iteT {i f p q s s' c} : f (s.cl i) = true → IExec I iprog p s s' c →
      IExec I iprog (.ite i f p q) s s' (c + 25)
  | iteF {i f p q s s' c} : f (s.cl i) = false → IExec I iprog q s s' c →
      IExec I iprog (.ite i f p q) s s' (c + 25)
  | loopT {i f p s s₁ s₂ c₁ c₂} : f (s.cl i) = true → IExec I iprog p s s₁ c₁ →
      IExec I iprog (.loop i f p) s₁ s₂ c₂ → IExec I iprog (.loop i f p) s s₂ (c₁ + c₂ + 30)
  | loopF {i f p s} : f (s.cl i) = false → IExec I iprog (.loop i f p) s s 30
  | call {j s s' c} : IExec I iprog (iprog j) s s' c → IExec I iprog (.call j) s s' (c + 2)

/-! ## Compilation -/

def branchOn (i : Fin S) (k : Cls → Stmt (S + 4) P) : Stmt (S + 4) P :=
  caseTop (ι i) fun o =>
    if o = none then k .none
    else if o = some hdMark then k .mark
    else if o = some hdNeg then k .neg
    else .seq (.pop (ι i)) (.ite (ι i) isBitT (.seq (.push (ι i) hdPos) (k .pos))
      (.seq (.push (ι i) hdPos) (k .zero)))

def flagOf (i : Fin S) (f : Cls → Bool) : Stmt (S + 4) P :=
  branchOn i fun c => .push (s3 S) (if f c then 1 else 0)

def loopTest (o : Option Sym) : Bool := decide (o = some 1)

def compile : IStmt S P → Stmt (S + 4) P
  | .skip => .skip
  | .op o => o.stmt
  | .seq p q => .seq (compile p) (compile q)
  | .ite i f p q => branchOn i fun c => if f c then compile p else compile q
  | .loop i f p => .seq (flagOf i f) (.seq (.loop (s3 S) loopTest
      (.seq (.pop (s3 S)) (.seq (compile p) (flagOf i f)))) (.pop (s3 S)))
  | .call j => .call j

def inner (i : Fin S) (f : Cls → Bool) (p : IStmt S P) : Stmt (S + 4) P :=
  .loop (s3 S) loopTest (.seq (.pop (s3 S)) (.seq (compile p) (flagOf i f)))

variable {I : List Sym} {prog : Fin P → Stmt (S + 4) P}

def encF (i : Fin S) (f : Cls → Bool) (s : IState S) : St (S + 4) :=
  (encI s).push (s3 S) (if f (s.cl i) then 1 else 0)

end NLogN

namespace NLogN

open IntMul

variable {S P : ℕ}

end NLogN

end

section
-- Source module: NLogN.IRuns
/-!
# Reasoning about item-language programs
-/

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

/-- `p` runs from `s` to `s'` with cost at most `c`. -/
def IRuns (I : List Sym) (iprog : Fin P → IStmt S P) (p : IStmt S P) (s s' : IState S) (c : ℕ) :
    Prop := ∃ c' ≤ c, IExec I iprog p s s' c'

namespace IRuns

end IRuns

/-! ## Operations as `IRuns` -/

namespace IRuns

end IRuns

/-! ## State lemmas -/

end NLogN

end

section
-- Source module: NLogN.ISchool
/-!
# Schoolbook multiplication in the item language (base case)
-/

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

def smulBody (X Y W Q R T : Fin S) : IStmt S P :=
  .seq (.op (.split Y Q R)) (.seq (.ite R (fun c => c == .zero) .skip
    (.ite R (fun c => c == .pos) (.seq (.op (.dup X T)) (.op (.add T W W)))
      (.seq (.op (.dup X T)) (.seq (.op (.neg T)) (.op (.add T W W))))))
    (.seq (.op (.drop R)) (.op (.shl X Q))))

/-- `A : x :: y :: L  ↦  x*y :: L`. -/
def smul (A X Y W Q R T : Fin S) : IStmt S P :=
  .seq (.op (.mv A X)) (.seq (.op (.mv A Y)) (.seq (.op (.const W 0)) (.seq (.op (.const Q 2))
    (.seq (.loop Y (fun c => c != .zero) (smulBody X Y W Q R T))
      (.seq (.op (.drop Y)) (.seq (.op (.drop X)) (.seq (.op (.drop Q)) (.op (.mv W A)))))))))

/-- The frame of the schoolbook loop. -/
def schoolSt (s : IState S) (A X Y W Q : Fin S) (L : List Item) (a b c : ℤ) : IState S :=
  ((((s.upd A L).upd X (.int a :: s.G X)).upd Y (.int b :: s.G Y)).upd W (.int c :: s.G W)).upd Q
    (.int 2 :: s.G Q)

def schoolIter (L : ℕ) : ℕ := 1500 * (L + 2)

end NLogN

namespace NLogN

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

def smulCost (N : ℕ) : ℕ := 2000 * (N + 2) ^ 2

end NLogN

end

section
-- Source module: NLogN.IEval
/-!
# A functional evaluator for item-language programs

`IOp.run` and `IStmt.run` compute the result state and cost of a program (with fuel for loops and
calls).  A successful evaluation gives an execution (`IStmt.run_exec`), so straight-line code can
be checked by `simp` evaluation instead of step-by-step reasoning.  Input and output operations
are not evaluated (`none`).
-/

namespace NLogN

open IntMul

variable {S P : ℕ}

/-- The exponent of a ruler. -/
def rulerExp (r : ℤ) : Option ℕ :=
  if r = 2 ^ (Nat.log2 r.toNat) then some (Nat.log2 r.toNat) else none

def IOp.run (s : IState S) : IOp S → Option (IState S × ℕ)
  | .add i j k =>
    match s.G i, s.G j with
    | .int a :: Li, .int b :: Lj =>
      if i = j then none else
        some ((((s.upd i Li).upd j Lj).upd k (.int (a + b) :: ((s.upd i Li).upd j Lj).G k)),
          addIntCost (zsz a + zsz b))
    | _, _ => none
  | .neg i =>
    match s.G i with
    | .int a :: L => some (s.upd i (.int (-a) :: L), 20)
    | _ => none
  | .dup i j =>
    match s.G i with
    | .int a :: _ => if i = j then none else some (s.upd j (.int a :: s.G j), dupCost (zsz a) + 17)
    | _ => none
  | .drop i =>
    match s.G i with
    | .int a :: L => some (s.upd i L, 2 * zsz a + 2)
    | _ => none
  | .mv i j =>
    match s.G i with
    | .int a :: L =>
      if i = j then none else some ((s.upd i L).upd j (.int a :: s.G j), 2 * moveCost (zsz a) + 19)
    | _ => none
  | .const i c => some (s.upd i (.int c :: s.G i), zsz c + 2)
  | .mark i => some (s.upd i (.mark :: s.G i), 1)
  | .unmark i =>
    match s.G i with
    | .mark :: L => some (s.upd i L, 1)
    | _ => none
  | .shl i j =>
    match s.G i, s.G j with
    | .int a :: Li, .int r :: _ =>
      match rulerExp r with
      | some m => if i = j then none else some (s.upd i (.int (a * 2 ^ m) :: Li), shlCost m)
      | none => none
    | _, _ => none
  | .split i j k =>
    match s.G i, s.G j with
    | .int a :: Li, .int r :: _ =>
      match rulerExp r with
      | some m =>
        if i = j ∨ i = k ∨ j = k then none else
          some ((s.upd i (.int (sg a (a.natAbs / 2 ^ m)) :: Li)).upd k
            (.int (sg a (a.natAbs % 2 ^ m)) :: s.G k), splitCost m + 17)
      | none => none
    | _, _ => none
  | .cat i j k =>
    match s.G i, s.G j, s.G k with
    | .int a :: Li, .int r :: _, .int b :: Lk =>
      match rulerExp r with
      | some m =>
        if i = j ∨ i = k ∨ j = k ∨ a = 0 ∨ b < 0 ∨ 2 ^ m ≤ b then none else
          some ((s.upd i (.int (sg a (a.natAbs * 2 ^ m + b.toNat)) :: Li)).upd k Lk, catCost m)
      | none => none
    | _, _, _ => none
  | .read _ _ _ => none
  | .write _ _ => none

/-- Evaluation with fuel (each loop iteration and call uses one unit). -/
def IStmt.run (iprog : Fin P → IStmt S P) : ℕ → IStmt S P → IState S → Option (IState S × ℕ)
  | _, .skip, s => some (s, 1)
  | _, .op o, s => o.run s
  | n, .seq p q, s =>
    match IStmt.run iprog n p s with
    | some (s₁, c₁) =>
      match IStmt.run iprog n q s₁ with
      | some (s₂, c₂) => some (s₂, c₁ + c₂)
      | none => none
    | none => none
  | n, .ite i f p q, s =>
    if f (s.cl i) then (IStmt.run iprog n p s).map (fun r => (r.1, r.2 + 25))
    else (IStmt.run iprog n q s).map (fun r => (r.1, r.2 + 25))
  | 0, .loop _ _ _, _ => none
  | n + 1, .loop i f p, s =>
    if f (s.cl i) then
      match IStmt.run iprog (n + 1) p s with
      | some (s₁, c₁) =>
        match IStmt.run iprog n (.loop i f p) s₁ with
        | some (s₂, c₂) => some (s₂, c₁ + c₂ + 30)
        | none => none
      | none => none
    else some (s, 30)
  | 0, .call _, _ => none
  | n + 1, .call j, s => (IStmt.run iprog n (iprog j) s).map (fun r => (r.1, r.2 + 2))

end NLogN

end

section
-- Source module: NLogN.SSBase
/-!
# Schönhage–Strassen in the item language: stack layout and Fermat-ring arithmetic

The program uses `NS` item stacks with fixed roles, so that stack (in)equalities are decided by
`simp`.  The current ring `ℤ/(2^n+1)` is given by the ruler `2^n` on top of `RN` and the modulus
`2^n+1` on top of `RF`.
-/

namespace NLogN

open IntMul

abbrev NS : ℕ := 32

abbrev RN : Fin NS := 2
abbrev RF : Fin NS := 3
abbrev T1 : Fin NS := 4
abbrev T2 : Fin NS := 5
abbrev T3 : Fin NS := 6
abbrev T4 : Fin NS := 7

/-- `o` evaluates to the state `t` with cost at most `c`. -/
def runsTo {S : ℕ} (o : Option (IState S × ℕ)) (t : IState S) (c : ℕ) : Prop :=
  match o with
  | some r => r.1 = t ∧ r.2 ≤ c
  | none => False

@[simp] theorem cls_none' : cls none = .none := rfl
@[simp] theorem cls_mark' : cls (some .mark) = .mark := rfl

@[simp] theorem IState.cl_def {S : ℕ} (s : IState S) (i : Fin S) : s.cl i = cls (s.G i).head? := rfl

/-! ## Fermat reduction -/

variable {P : ℕ}

/-- Reduce the integer on top of `X` modulo `2^n+1` (needs `|z| < 2^(2n+1)`). -/
def fredA (X : Fin NS) : IStmt NS P :=
  .seq (.op (.split X RN T1)) (.seq (.op (.split X RN T2)) (.seq (.op (.neg T2))
    (.seq (.op (.add T1 T2 T1)) (.op (.add T1 X X)))))

def fadd (X : Fin NS) : IStmt NS P :=
  .ite X (fun c => c == .neg) (.seq (.op (.dup RF T1)) (.op (.add T1 X X))) .skip

def fsub (X : Fin NS) : IStmt NS P :=
  .seq (.op (.dup RF T1)) (.seq (.op (.neg T1)) (.seq (.op (.dup X T2)) (.seq (.op (.add T2 T1 T1))
    (.ite T1 (fun c => c != .neg) (.seq (.op (.drop X)) (.op (.mv T1 X))) (.op (.drop T1))))))

def fred (X : Fin NS) : IStmt NS P :=
  .seq (fredA X) (.seq (fadd X) (.seq (fadd X) (.seq (fsub X) (fsub X))))

end NLogN

namespace NLogN

open IntMul

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt NS P}

end NLogN

end

section
-- Source module: NLogN.HExpr
/-!
# Expressions over register stacks

`Ex S` are integer expressions whose leaves read the top item of a register stack.  `Ex.code z e`
evaluates `e` onto the value stack `z.V`, using the scratch stacks of `z` and leaving every other
stack unchanged (`Ex.code_runs`).
-/

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

/-- The integer on top of stack `i` (`0` if there is none). -/
def IState.top (s : IState S) (i : Fin S) : ℤ :=
  match s.G i with
  | .int a :: _ => a
  | _ => 0

/-- The top of stack `i` is an integer. -/
def IState.IsInt (s : IState S) (i : Fin S) : Prop := ∃ L, s.G i = .int (s.top i) :: L

/-- The scratch stacks used by expression code. -/
structure ExSt (S : ℕ) where
  V : Fin S
  A : Fin S
  B : Fin S
  C : Fin S
  D : Fin S
  E : Fin S
  F : Fin S
  G : Fin S

def ExSt.all (z : ExSt S) : List (Fin S) := [z.V, z.A, z.B, z.C, z.D, z.E, z.F, z.G]

inductive Ex (S : ℕ) where
  | reg (i : Fin S)
  | lit (c : ℤ)
  | add (a b : Ex S)
  | neg (a : Ex S)
  | mul (a b : Ex S)
  /-- `a * r` for a ruler `r = 2^m` -/
  | shl (a r : Ex S)
  /-- `a.tdiv r` for a ruler `r = 2^m` -/
  | quo (a r : Ex S)
  /-- `a.tmod r` for a ruler `r = 2^m` -/
  | rem (a r : Ex S)
  /-- `2^a` for `a ≥ 0` -/
  | pow2 (a : Ex S)
  /-- `if c < 0 then x else y` -/
  | sel (c x y : Ex S)
  /-- `a / b` for `a ≥ 0`, `b > 0` -/
  | ndiv (a b : Ex S)
  /-- `a % b` for `a ≥ 0`, `b > 0` -/
  | nmod (a b : Ex S)

namespace Ex

def regs : Ex S → List (Fin S)
  | reg i => [i]
  | lit _ => []
  | add a b | mul a b | shl a b | quo a b | rem a b | ndiv a b | nmod a b => a.regs ++ b.regs
  | neg a | pow2 a => a.regs
  | sel c x y => c.regs ++ x.regs ++ y.regs

def eval (s : IState S) : Ex S → ℤ
  | reg i => s.top i
  | lit c => c
  | add a b => a.eval s + b.eval s
  | neg a => -a.eval s
  | mul a b => a.eval s * b.eval s
  | shl a r => a.eval s * r.eval s
  | quo a r => (a.eval s).tdiv (r.eval s)
  | rem a r => (a.eval s).tmod (r.eval s)
  | pow2 a => 2 ^ (a.eval s).toNat
  | sel c x y => if c.eval s < 0 then x.eval s else y.eval s
  | ndiv a b => a.eval s / b.eval s
  | nmod a b => a.eval s % b.eval s

/-- Side conditions: registers hold integers, rulers are powers of two, divisions are of
naturals by positive naturals. -/
def ok (s : IState S) : Ex S → Prop
  | reg i => s.IsInt i
  | lit _ => True
  | add a b | mul a b => a.ok s ∧ b.ok s
  | neg a => a.ok s
  | shl a r | quo a r | rem a r => a.ok s ∧ r.ok s ∧ ∃ m : ℕ, r.eval s = 2 ^ m
  | pow2 a => a.ok s ∧ 0 ≤ a.eval s
  | sel c x y => c.ok s ∧ (if c.eval s < 0 then x.ok s else y.ok s)
  | ndiv a b | nmod a b => a.ok s ∧ b.ok s ∧ 0 ≤ a.eval s ∧ 0 < b.eval s

end Ex

/-! ## Long division -/

/-- Cost bound of long division of `a ≥ 0` by `b > 0` with `zsz a, zsz b ≤ n`. -/
def ldivCost (n : ℕ) : ℕ := 1000000 * (n + 2) ^ 2

/-- Cost of `2^a` by repeated doubling. -/
def pow2Cost (a : ℕ) : ℕ := (a + 1) * (addIntCost (zsz (a : ℤ) + 1) + 300) + 300

namespace Ex

def cost (s : IState S) : Ex S → ℕ
  | reg i => dupCost (zsz (s.top i)) + 17
  | lit c => zsz c + 2
  | add a b => a.cost s + b.cost s + 2 * moveCost (zsz (b.eval s)) + 2 * moveCost (zsz (a.eval s)) +
      38 + addIntCost (zsz (b.eval s) + zsz (a.eval s))
  | neg a => a.cost s + 20
  | mul a b => a.cost s + b.cost s + smulCost (zsz (b.eval s) + zsz (a.eval s))
  | shl a r => a.cost s + r.cost s + 2 * moveCost (zsz (r.eval s)) + 19 +
      shlCost (zsz (r.eval s)) + 2 * zsz (r.eval s) + 2
  | quo a r | rem a r => a.cost s + r.cost s + 4 * moveCost (zsz (r.eval s)) +
      splitCost (zsz (r.eval s)) + 4 * zsz (a.eval s) + 4 * zsz (r.eval s) + 100 +
      2 * moveCost (zsz (a.eval s))
  | pow2 a => a.cost s + 2 * moveCost (zsz (a.eval s)) + pow2Cost (a.eval s).toNat +
      2 * zsz (a.eval s) + 100
  | sel c x y => c.cost s + 2 * moveCost (zsz (c.eval s)) + 19 + 25 + 2 * zsz (c.eval s) + 2 +
      (if c.eval s < 0 then x.cost s else y.cost s)
  | ndiv a b | nmod a b => a.cost s + b.cost s + 2 * ldivCost (zsz (a.eval s) + zsz (b.eval s))

end Ex

/-! ## Code -/

/-- `V : a :: L ↦ 2^a :: L` (for `a ≥ 0`), by doubling; uses `A, B, C`. -/
def pow2Code (z : ExSt S) : IStmt S P :=
  .seq (.op (.mv z.V z.C)) (.seq (.op (.const z.V 1)) (.seq (.op (.const z.A 2))
    (.seq (.loop z.C (fun c => c != .zero)
      (.seq (.op (.shl z.V z.A)) (.seq (.op (.const z.B (-1))) (.op (.add z.B z.C z.C)))))
      (.seq (.op (.drop z.A)) (.op (.drop z.C))))))

/-- `F := B - A`. -/
def ldFlag (z : ExSt S) : IStmt S P :=
  .seq (.op (.dup z.B z.F)) (.seq (.op (.dup z.A z.G)) (.seq (.op (.neg z.G))
    (.op (.add z.G z.F z.F))))

/-- Phase 1 body: `A := 2A`, `D := D + 1`. -/
def ldUp (z : ExSt S) : IStmt S P :=
  .seq (.op (.drop z.F)) (.seq (.op (.shl z.A z.E)) (.seq (.op (.const z.G 1))
    (.seq (.op (.add z.G z.D z.D)) (ldFlag z))))

/-- Phase 2 body: `A := A/2`, `C := 2C`, and if `B ≥ A` then `B := B - A`, `C := C + 1`;
`D := D - 1`. -/
def ldDown (z : ExSt S) : IStmt S P :=
  .seq (.op (.split z.A z.E z.G)) (.seq (.op (.drop z.G)) (.seq (.op (.shl z.C z.E))
    (.seq (ldFlag z) (.seq (.ite z.F (fun c => c == .neg) (.op (.drop z.F))
      (.seq (.op (.drop z.B)) (.seq (.op (.mv z.F z.B)) (.seq (.op (.const z.G 1))
        (.op (.add z.G z.C z.C))))))
      (.seq (.op (.const z.G (-1))) (.op (.add z.G z.D z.D)))))))

/-- `V : b :: a :: L ↦ (a / b) :: L` and `G : R ↦ (a % b) :: R`, for `a ≥ 0`, `b > 0`. -/
def ldivCode (z : ExSt S) : IStmt S P :=
  .seq (.seq (.op (.mv z.V z.A)) (.seq (.op (.mv z.V z.B)) (.seq (.op (.const z.C 0))
    (.seq (.op (.const z.D 0)) (.seq (.op (.const z.E 2)) (ldFlag z))))))
    (.seq (.loop z.F (fun c => c != .neg) (ldUp z))
      (.seq (.op (.drop z.F))
        (.seq (.loop z.D (fun c => c != .zero) (ldDown z))
          (.seq (.op (.drop z.E)) (.seq (.op (.drop z.D)) (.seq (.op (.drop z.A))
            (.seq (.op (.mv z.C z.V)) (.op (.mv z.B z.G)))))))))

namespace Ex

def code (z : ExSt S) : Ex S → IStmt S P
  | reg i => .op (.dup i z.V)
  | lit c => .op (.const z.V c)
  | add a b => .seq (a.code z) (.seq (b.code z) (.seq (.op (.mv z.V z.A))
      (.seq (.op (.mv z.V z.B)) (.op (.add z.A z.B z.V)))))
  | neg a => .seq (a.code z) (.op (.neg z.V))
  | mul a b => .seq (a.code z) (.seq (b.code z) (smul z.V z.A z.B z.C z.D z.E z.F))
  | shl a r => .seq (a.code z) (.seq (r.code z) (.seq (.op (.mv z.V z.A))
      (.seq (.op (.shl z.V z.A)) (.op (.drop z.A)))))
  | quo a r => .seq (a.code z) (.seq (r.code z) (.seq (.op (.mv z.V z.A))
      (.seq (.op (.split z.V z.A z.B)) (.seq (.op (.drop z.B)) (.op (.drop z.A))))))
  | rem a r => .seq (a.code z) (.seq (r.code z) (.seq (.op (.mv z.V z.A))
      (.seq (.op (.split z.V z.A z.B)) (.seq (.op (.drop z.V)) (.seq (.op (.mv z.B z.V))
        (.op (.drop z.A)))))))
  | pow2 a => .seq (a.code z) (pow2Code z)
  | sel c x y => .seq (c.code z) (.seq (.op (.mv z.V z.A))
      (.ite z.A (fun k => k == .neg) (.seq (.op (.drop z.A)) (x.code z))
        (.seq (.op (.drop z.A)) (y.code z))))
  | ndiv a b => .seq (a.code z) (.seq (b.code z) (.seq (ldivCode z) (.op (.drop z.G))))
  | nmod a b => .seq (a.code z) (.seq (b.code z) (.seq (ldivCode z)
      (.seq (.op (.drop z.V)) (.op (.mv z.G z.V)))))

end Ex

end NLogN

end

section
-- Source module: NLogN.HExprPow
/-!
# `2^a` by doubling
-/

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

/-- Pairwise distinctness of the scratch stacks. -/
structure ExSt.Good (z : ExSt S) : Prop where
  VA : z.V ≠ z.A
  VB : z.V ≠ z.B
  VC : z.V ≠ z.C
  VD : z.V ≠ z.D
  VE : z.V ≠ z.E
  VF : z.V ≠ z.F
  VG : z.V ≠ z.G
  AB : z.A ≠ z.B
  AC : z.A ≠ z.C
  AD : z.A ≠ z.D
  AE : z.A ≠ z.E
  AF : z.A ≠ z.F
  AG : z.A ≠ z.G
  BC : z.B ≠ z.C
  BD : z.B ≠ z.D
  BE : z.B ≠ z.E
  BF : z.B ≠ z.F
  BG : z.B ≠ z.G
  CD : z.C ≠ z.D
  CE : z.C ≠ z.E
  CF : z.C ≠ z.F
  CG : z.C ≠ z.G
  DE : z.D ≠ z.E
  DF : z.D ≠ z.F
  DG : z.D ≠ z.G
  EF : z.E ≠ z.F
  EG : z.E ≠ z.G
  FG : z.F ≠ z.G

end NLogN

end

section
-- Source module: NLogN.HExprDiv
/-!
# Long division
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

variable {S P : ℕ} {I : List Sym} {iprog : Fin P → IStmt S P}

/-- The number of stacks of the HvdH machine. -/
abbrev HS : ℕ := 64

/-- The expression scratch stacks. -/
def zX : ExSt HS := ⟨56, 57, 58, 59, 60, 61, 62, 63⟩

@[simp] theorem zX_V : zX.V = 56 := rfl
@[simp] theorem zX_A : zX.A = 57 := rfl
@[simp] theorem zX_B : zX.B = 58 := rfl
@[simp] theorem zX_C : zX.C = 59 := rfl
@[simp] theorem zX_D : zX.D = 60 := rfl
@[simp] theorem zX_E : zX.E = 61 := rfl
@[simp] theorem zX_F : zX.F = 62 := rfl
@[simp] theorem zX_G : zX.G = 63 := rfl

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The state of the long division: `V = L`, registers `A … E` on top of the original stacks,
optionally a flag on `F`. -/
def ldSt (s : IState HS) (L : List Item) (a b c d : ℤ) (f : Option ℤ) : IState HS :=
  ((((((s.upd 56 L).upd 57 (.int a :: s.G 57)).upd 58 (.int b :: s.G 58)).upd 59
    (.int c :: s.G 59)).upd 60 (.int d :: s.G 60)).upd 61 (.int 2 :: s.G 61)).upd 62
    (match f with | some x => .int x :: s.G 62 | none => s.G 62)

end NLogN

end

section
-- Source module: NLogN.HExprRun
/-!
# Expression code runs (on the HvdH stacks)
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

namespace Ex

/-- Registers not among the scratch stacks. -/
def Clean (e : Ex HS) : Prop := ∀ i ∈ e.regs, i ∉ zX.all

end Ex

end NLogN

end

section
-- Source module: NLogN.HExprCost
/-!
# A uniform cost bound for expressions

If every value an expression computes (and every register it reads) has at most `B` bits, its
code costs at most `size · 10^7 · (B+2)^2`.
-/

namespace NLogN

open IntMul

variable {S : ℕ}

namespace Ex

def size : Ex S → ℕ
  | reg _ | lit _ => 1
  | add a b | mul a b | shl a b | quo a b | rem a b | ndiv a b | nmod a b => a.size + b.size + 1
  | neg a | pow2 a => a.size + 1
  | sel c x y => c.size + x.size + y.size + 1

/-- Every value computed along the evaluation has at most `B` bits. -/
def Bnd (s : IState S) (B : ℕ) : Ex S → Prop
  | reg i => zsz (s.top i) ≤ B
  | lit c => zsz c ≤ B
  | add a b => a.Bnd s B ∧ b.Bnd s B ∧ zsz (a.eval s + b.eval s) ≤ B
  | mul a b => a.Bnd s B ∧ b.Bnd s B ∧ zsz (a.eval s * b.eval s) ≤ B
  | shl a b => a.Bnd s B ∧ b.Bnd s B ∧ zsz (a.eval s * b.eval s) ≤ B
  | quo a b => a.Bnd s B ∧ b.Bnd s B
  | rem a b => a.Bnd s B ∧ b.Bnd s B
  | ndiv a b => a.Bnd s B ∧ b.Bnd s B ∧ zsz (a.eval s / b.eval s) ≤ B
  | nmod a b => a.Bnd s B ∧ b.Bnd s B ∧ zsz (a.eval s % b.eval s) ≤ B
  | neg a => a.Bnd s B
  | pow2 a => a.Bnd s B ∧ (a.eval s).toNat + 1 ≤ B
  | sel c x y => c.Bnd s B ∧ (if c.eval s < 0 then x.Bnd s B else y.Bnd s B)

end Ex

end NLogN

end

section
-- Source module: NLogN.HReg
/-!
# Register programs

`RS` is a small structured language over the register stacks `0 … 55`: assignments of
expressions, pushes and pops, conditionals and counted loops.  `RS.run` is its functional
semantics on `IState`s, and `RS.code_runs` shows that the compiled `IStmt` realises it with cost
`RS.cost`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

inductive RS where
  | skip
  /-- replace the top of `i` by the value of `e` -/
  | asg (i : Fin HS) (e : Ex HS)
  /-- push the value of `e` onto `i` -/
  | push (i : Fin HS) (e : Ex HS)
  /-- drop the (integer) top of `i` -/
  | pop (i : Fin HS)
  | seq (a b : RS)
  /-- `if c < 0 then a else b` -/
  | ite (c : Ex HS) (a b : RS)
  /-- push `n` on `i`, then run `body; i := i - 1` until the top of `i` is `0`; pop `i`.
  The body sees the remaining count on top of `i`. -/
  | rep (i : Fin HS) (n : Ex HS) (body : RS)

namespace RS

/-- Stacks a program may write. -/
def writes : RS → List (Fin HS)
  | skip => []
  | asg i _ | push i _ | pop i => [i]
  | seq a b | ite _ a b => a.writes ++ b.writes
  | rep i _ b => i :: b.writes

/-- One decrement of the counter on `i`. -/
def decSt (i : Fin HS) (t : IState HS) : IState HS :=
  t.upd i (.int (t.top i - 1) :: (t.G i).tail)

def run : RS → IState HS → IState HS
  | skip, s => s
  | asg i e, s => s.upd i (.int (e.eval s) :: (s.G i).tail)
  | push i e, s => s.upd i (.int (e.eval s) :: s.G i)
  | pop i, s => s.upd i (s.G i).tail
  | seq a b, s => b.run (a.run s)
  | ite c a b, s => if c.eval s < 0 then a.run s else b.run s
  | rep i n b, s =>
    let t := (fun t => decSt i (b.run t))^[(n.eval s).toNat] (s.upd i (.int (n.eval s) :: s.G i))
    t.upd i (t.G i).tail

/-- The state at the start of iteration `k` of `rep i n b` from `s`. -/
def repSt (i : Fin HS) (n : Ex HS) (b : RS) (s : IState HS) (k : ℕ) : IState HS :=
  (fun t => decSt i (b.run t))^[k] (s.upd i (.int (n.eval s) :: s.G i))

/-- The decrement expression. -/
def decE (i : Fin HS) : Ex HS := .add (.reg i) (.lit (-1))

def asgCost (i : Fin HS) (e : Ex HS) (s : IState HS) : ℕ :=
  e.cost s + (2 * zsz (s.top i) + 2) + (2 * moveCost (zsz (e.eval s)) + 19)

def ok : RS → IState HS → Prop
  | skip, _ => True
  | asg i e, s => i.val < 56 ∧ e.Clean ∧ e.ok s ∧ s.IsInt i
  | push i e, s => i.val < 56 ∧ e.Clean ∧ e.ok s
  | pop i, s => i.val < 56 ∧ s.IsInt i
  | seq a b, s => a.ok s ∧ b.ok (a.run s)
  | ite c a b, s => c.Clean ∧ c.ok s ∧ (if c.eval s < 0 then a.ok s else b.ok s)
  | rep i n b, s => i.val < 56 ∧ n.Clean ∧ n.ok s ∧ 0 ≤ n.eval s ∧ i ∉ b.writes ∧
      ∀ k < (n.eval s).toNat, b.ok (NLogN.RS.repSt i n b s k)

def cost : RS → IState HS → ℕ
  | skip, _ => 1
  | asg i e, s => asgCost i e s

  | push i e, s => e.cost s + (2 * moveCost (zsz (e.eval s)) + 19)
  | pop i, s => 2 * zsz (s.top i) + 2
  | seq a b, s => a.cost s + b.cost (a.run s)
  | ite c a b, s => c.cost s + (2 * moveCost (zsz (c.eval s)) + 19) + (2 * zsz (c.eval s) + 2) +
      25 + (if c.eval s < 0 then a.cost s else b.cost s)
  | rep i n b, s => n.cost s + (2 * moveCost (zsz (n.eval s)) + 19) +
      ((∑ k ∈ Finset.range (n.eval s).toNat,
        (b.cost (NLogN.RS.repSt i n b s k) + asgCost i (decE i) (b.run (NLogN.RS.repSt i n b s k)) + 30)) + 30) + 2

def code : RS → IStmt HS P
  | skip => .skip
  | asg i e => .seq (e.code zX) (.seq (.op (.drop i)) (.op (.mv zX.V i)))
  | push i e => .seq (e.code zX) (.op (.mv zX.V i))
  | pop i => .op (.drop i)
  | seq a b => .seq a.code b.code
  | ite c a b => .seq (c.code zX) (.seq (.op (.mv zX.V zX.A))
      (.ite zX.A (fun k => k == .neg) (.seq (.op (.drop zX.A)) a.code)
        (.seq (.op (.drop zX.A)) b.code)))
  | rep i n b => .seq (.seq (n.code zX) (.op (.mv zX.V i)))
      (.seq (.loop i (fun c => c != .zero)
        (.seq b.code (.seq ((decE i).code zX) (.seq (.op (.drop i)) (.op (.mv zX.V i))))))
        (.op (.drop i)))

/-! ## Frame -/

/-! ## Compilation -/

/-! ## Sequences and reasoning lemmas -/

instance (e : Ex HS) : Decidable e.Clean := by unfold Ex.Clean; infer_instance

/-! ## Cost bounds -/

def seqL : List RS → RS
  | [] => skip
  | [a] => a
  | a :: l => NLogN.RS.seq a (seqL l)

end RS

end NLogN

end

section
-- Source module: NLogN.HList
/-!
# Counted loops through a family of states

`rep_states`: if the body maps the `k`-th logical state (with the counter on top of `i`) to the
`(k+1)`-st, the loop computes the last one.  This is the interface used for list processing.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul RS

/-- The `k`-th logical state with the remaining count `N - k` on top of the counter stack. -/
def withCtr (i : Fin HS) (s : IState HS) (N k : ℕ) (t : IState HS) : IState HS :=
  t.upd i (.int ((N : ℤ) - k) :: s.G i)

/-- Run hypothesis of a counted loop through the states `S 0 = s, …, S N`. -/
structure RepStates (i : Fin HS) (n : Ex HS) (b : RS) (s : IState HS) (S : ℕ → IState HS) :
    Prop where
  nonneg : 0 ≤ n.eval s
  hw : i ∉ b.writes
  h0 : S 0 = s
  hG : ∀ k, (S k).G i = s.G i
  step : ∀ k < (n.eval s).toNat,
    b.run (withCtr i s (n.eval s).toNat k (S k)) = withCtr i s (n.eval s).toNat k (S (k + 1))

end NLogN

end

section
-- Source module: NLogN.HRM
/-!
# Register programs with multiplication calls

`RM` extends the register programs `RS` by `mul k i j`, which pushes `top j * top i` onto `k`
by calling the multiplication procedure `pM` with the calling convention

  `ARG = 55 : a :: b :: L  ↦  (a * b) :: L`,

all other stacks unchanged.  The cost of a call and the inputs on which it may be used are a
parameter `M : MulSpec`; `RM.code_runs` holds whenever the procedure meets `MulOK M`.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt HS P}

/-- The argument stack of the multiplication procedure. -/
abbrev ARGM : Fin HS := 55

/-- Cost and domain of the multiplication procedure. -/
structure MulSpec where
  Φ : ℤ → ℤ → ℕ
  dom : ℤ → ℤ → Prop

/-- The procedure `pM` multiplies according to `M`. -/
def MulOK (I : List Sym) (iprog : Fin P → IStmt HS P) (pM : Fin P) (M : MulSpec) : Prop :=
  ∀ (s : IState HS) (a b : ℤ) (L : List Item), M.dom a b → s.G ARGM = .int a :: .int b :: L →
    IRuns I iprog (.call pM) s (s.upd ARGM (.int (a * b) :: L)) (M.Φ a b)

inductive RM where
  | rs (p : RS)
  /-- push `top j * top i` onto `k` -/
  | mul (k i j : Fin HS)
  | seq (a b : RM)
  | ite (c : Ex HS) (a b : RM)
  | rep (i : Fin HS) (n : Ex HS) (body : RM)
  /-- a single item-language operation (e.g. `split`/`cat` with costs independent of the
  untouched high part) -/
  | op (o : IOp HS)
  /-- a call of procedure number `q + 1` (number `0` is the multiplication) with functional
  specification `f`, cost `c`, precondition `pre`, writing only the stacks `w` -/
  | prc (q : ℕ) (f : IState HS → IState HS) (c : IState HS → ℕ) (pre : IState HS → Prop)
      (w : List (Fin HS)) (hf : ∀ s j, j ∉ w → (f s).G j = s.G j)
      (hio : ∀ s, (f s).ip = s.ip ∧ (f s).out = s.out)

/-- The stacks an operation mentions. -/
def IOp.stks : IOp HS → List (Fin HS)
  | .add i j k | .split i j k | .cat i j k => [i, j, k]
  | .neg i | .drop i | .const i _ | .mark i | .unmark i => [i]
  | .dup i j | .mv i j | .shl i j | .read i j _ | .write i j => [i, j]

namespace RM

open RS (decSt decE asgCost)

def writes : RM → List (Fin HS)
  | rs p => p.writes
  | mul k _ _ => [k]
  | seq a b | ite _ a b => a.writes ++ b.writes
  | rep i _ b => i :: b.writes
  | op o => o.stks
  | prc _ _ _ _ w _ _ => w

def run : RM → IState HS → IState HS
  | rs p, s => p.run s
  | mul k i j, s => s.upd k (.int (s.top j * s.top i) :: s.G k)
  | seq a b, s => b.run (a.run s)
  | ite c a b, s => if c.eval s < 0 then a.run s else b.run s
  | rep i n b, s =>
    let t := (fun t => decSt i (b.run t))^[(n.eval s).toNat] (s.upd i (.int (n.eval s) :: s.G i))
    t.upd i (t.G i).tail
  | op o, s => match o.run s with
    | some r => r.1
    | none => s
  | prc _ f _ _ _ _ _, s => f s

def repSt (i : Fin HS) (n : Ex HS) (b : RM) (s : IState HS) (k : ℕ) : IState HS :=
  (fun t => decSt i (b.run t))^[k] (s.upd i (.int (n.eval s) :: s.G i))

def mulCost (M : MulSpec) (k i j : Fin HS) (s : IState HS) : ℕ :=
  (dupCost (zsz (s.top i)) + 17) + (dupCost (zsz (s.top j)) + 17) +
    (M.Φ (s.top j) (s.top i) + 2) + (2 * moveCost (zsz (s.top j * s.top i)) + 19)

def ok (M : MulSpec) : RM → IState HS → Prop
  | rs p, s => p.ok s
  | mul k i j, s => k.val < 55 ∧ i.val < 55 ∧ j.val < 55 ∧ s.IsInt i ∧ s.IsInt j ∧
      M.dom (s.top j) (s.top i)
  | seq a b, s => ok M a s ∧ ok M b (a.run s)
  | ite c a b, s => c.Clean ∧ c.ok s ∧ (if c.eval s < 0 then ok M a s else ok M b s)
  | rep i n b, s => i.val < 56 ∧ n.Clean ∧ n.ok s ∧ 0 ≤ n.eval s ∧ i ∉ b.writes ∧
      ∀ k < (n.eval s).toNat, ok M b (repSt i n b s k)
  | op o, s => (o.run s).isSome
  | prc _ _ _ pre _ _ _, s => pre s

def cost (M : MulSpec) : RM → IState HS → ℕ
  | rs p, s => p.cost s
  | mul k i j, s => mulCost M k i j s
  | seq a b, s => cost M a s + cost M b (a.run s)
  | ite c a b, s => c.cost s + (2 * moveCost (zsz (c.eval s)) + 19) + (2 * zsz (c.eval s) + 2) +
      25 + (if c.eval s < 0 then cost M a s else cost M b s)
  | rep i n b, s => n.cost s + (2 * moveCost (zsz (n.eval s)) + 19) +
      ((∑ k ∈ Finset.range (n.eval s).toNat,
        (cost M b (repSt i n b s k) + asgCost i (decE i) (b.run (repSt i n b s k)) + 30)) + 30) + 2
  | op o, s => match o.run s with
    | some r => r.2
    | none => 0
  | prc _ _ c _ _ _ _, s => c s

def code (pM : Fin P) (pc : ℕ → Fin P) : RM → IStmt HS P
  | rs p => p.code
  | mul k i j => .seq (.op (.dup i ARGM)) (.seq (.op (.dup j ARGM)) (.seq (.call pM)
      (.op (.mv ARGM k))))
  | seq a b => .seq (code pM pc a) (code pM pc b)
  | ite c a b => .seq (c.code zX) (.seq (.op (.mv zX.V zX.A))
      (.ite zX.A (fun k => k == .neg) (.seq (.op (.drop zX.A)) (code pM pc a))
        (.seq (.op (.drop zX.A)) (code pM pc b))))
  | rep i n b => .seq (.seq (n.code zX) (.op (.mv zX.V i)))
      (.seq (.loop i (fun c => c != .zero)
        (.seq (code pM pc b) (.seq ((decE i).code zX) (.seq (.op (.drop i)) (.op (.mv zX.V i))))))
        (.op (.drop i)))
  | op o => .op o
  | prc q _ _ _ _ _ _ => .call (pc (q + 1))

/-! ## Frame -/

variable (I iprog) in
/-- Every called procedure meets its specification. -/
def PrcOK (pc : ℕ → Fin P) : RM → Prop
  | prc q f c pre _ _ _ => ∀ s, pre s → IRuns I iprog (.call (pc (q + 1))) s (f s) (c s)
  | seq a b | ite _ a b => PrcOK pc a ∧ PrcOK pc b
  | rep _ _ b => PrcOK pc b
  | rs _ | mul _ _ _ | op _ => True

/-! ## Reasoning lemmas -/

def seqL : List RM → RM
  | [] => rs .skip
  | [a] => a
  | a :: l => NLogN.RM.seq a (seqL l)

/-! ## Counted loops through a family of states -/

structure RepStates (i : Fin HS) (n : Ex HS) (b : RM) (s : IState HS) (S : ℕ → IState HS) :
    Prop where
  nonneg : 0 ≤ n.eval s
  hw : i ∉ b.writes
  h0 : S 0 = s
  hG : ∀ k, (S k).G i = s.G i
  step : ∀ k < (n.eval s).toNat,
    b.run (withCtr i s (n.eval s).toNat k (S k)) = withCtr i s (n.eval s).toNat k (S (k + 1))

end RM

end NLogN

end

section
-- Source module: NLogN.HExprLin
/-!
# Linear cost bound for multiplication-free expressions

Expressions built from registers, literals, `add`, `neg`, `shl`, `quo`, `rem` and `sel` cost at
most `size · 2000 · (B + 2)` when every value has at most `B` bits.  This is what per-element
list operations use; products of large numbers go through the multiplication procedure instead.
-/

namespace NLogN

open IntMul

variable {S : ℕ}

namespace Ex

/-- No `mul`, `pow2`, `ndiv`, `nmod`. -/
def Lin : Ex S → Prop
  | reg _ | lit _ => True
  | add a b | shl a b | quo a b | rem a b => a.Lin ∧ b.Lin
  | neg a => a.Lin
  | sel c x y => c.Lin ∧ x.Lin ∧ y.Lin
  | mul _ _ | pow2 _ | ndiv _ _ | nmod _ _ => False

end Ex

end NLogN

end

section
-- Source module: NLogN.HRen
/-!
# Renaming stacks and procedures of item-language programs

Given a map of stacks `f : Fin S → Fin S'` and a map of procedure indices `g : Fin P → Fin P'`,
`IStmt.ren f g` renames every stack `i` to `f i` and every call `j` to `g j`.  For injective `f`,
executions are simulated: a run from `s` to `s'` becomes a run from `embed f s t` to `embed f s' t`,
where `IState.embed f s t` is `t` with the stacks in the image of `f` overwritten by those of `s`
(and `ip`, `out` taken from `s`).  Stacks outside the image of `f` are untouched.
-/

namespace NLogN

open IntMul

variable {S S' P P' : ℕ}

/-- Rename the stacks of an operation. -/
def IOp.ren (f : Fin S → Fin S') : IOp S → IOp S'
  | .add i j k => .add (f i) (f j) (f k)
  | .neg i => .neg (f i)
  | .dup i j => .dup (f i) (f j)
  | .drop i => .drop (f i)
  | .mv i j => .mv (f i) (f j)
  | .const i c => .const (f i) c
  | .mark i => .mark (f i)
  | .unmark i => .unmark (f i)
  | .shl i j => .shl (f i) (f j)
  | .split i j k => .split (f i) (f j) (f k)
  | .read i j stop => .read (f i) (f j) stop
  | .write i j => .write (f i) (f j)
  | .cat i j k => .cat (f i) (f j) (f k)

/-- Rename the stacks (by `f`) and the called procedures (by `g`) of a program. -/
def IStmt.ren (f : Fin S → Fin S') (g : Fin P → Fin P') : IStmt S P → IStmt S' P'
  | .skip => .skip
  | .op o => .op (o.ren f)
  | .seq p q => .seq (p.ren f g) (q.ren f g)
  | .ite i c p q => .ite (f i) c (p.ren f g) (q.ren f g)
  | .loop i c p => .loop (f i) c (p.ren f g)
  | .call j => .call (g j)

/-- `t` with the stacks `f x` replaced by `s.G x`; `ip` and `out` are those of `s`. -/
noncomputable def IState.embed (f : Fin S → Fin S') (s : IState S) (t : IState S') : IState S' :=
  ⟨Function.extend f s.G t.G, s.ip, s.out⟩

/-- The state seen through `f`. -/
def IState.pull (f : Fin S → Fin S') (t : IState S') : IState S :=
  ⟨fun x => t.G (f x), t.ip, t.out⟩

@[simp] theorem IState.pull_G (f : Fin S → Fin S') (t : IState S') (x : Fin S) :
    (IState.pull f t).G x = t.G (f x) := rfl
@[simp] theorem IState.pull_ip (f : Fin S → Fin S') (t : IState S') :
    (IState.pull f t).ip = t.ip := rfl
@[simp] theorem IState.pull_out (f : Fin S → Fin S') (t : IState S') :
    (IState.pull f t).out = t.out := rfl

@[simp] theorem IState.embed_ip (f : Fin S → Fin S') (s : IState S) (t : IState S') :
    (IState.embed f s t).ip = s.ip := rfl
@[simp] theorem IState.embed_out (f : Fin S → Fin S') (s : IState S) (t : IState S') :
    (IState.embed f s t).out = s.out := rfl

section Inj

variable {f : Fin S → Fin S'} (hf : Function.Injective f)
include hf

variable {I : List Sym}

end Inj

end NLogN

end

section
-- Source module: NLogN.HRenRM
/-!
# Relocating register programs

For a permutation `σ` of the `HS` stacks fixing every stack `≥ 55` (the multiplication argument
`55` and the expression scratch `56..63`), `Ex.ren σ`, `RS.ren σ`, `RM.ren σ` rename every register
`i` to `σ i`.  On the relocated state `IState.rn σ s` (stack `σ i` holds the old stack `i`) the
renamed program computes the relocated result, with the same side conditions and the same cost.
-/

set_option linter.unusedSimpArgs false

namespace NLogN

open IntMul

/-! ## Relocated states -/

/-- Stack `σ i` of the new state holds stack `i` of `s`; `ip`, `out` unchanged. -/
def IState.rn (σ : Equiv.Perm (Fin HS)) (s : IState HS) : IState HS :=
  ⟨fun y => s.G (σ.symm y), s.ip, s.out⟩

variable (σ : Equiv.Perm (Fin HS))

namespace IState

theorem rn_G' (s : IState HS) (y : Fin HS) : (s.rn σ).G y = s.G (σ.symm y) := rfl
@[simp] theorem rn_ip (s : IState HS) : (s.rn σ).ip = s.ip := rfl
@[simp] theorem rn_out (s : IState HS) : (s.rn σ).out = s.out := rfl

end IState

/-! ## Permutations fixing the reserved stacks -/

/-- `σ` fixes every stack `≥ 55`. -/
def FixHi (σ : Equiv.Perm (Fin HS)) : Prop := ∀ i : Fin HS, 55 ≤ i.val → σ i = i

instance (σ : Equiv.Perm (Fin HS)) : Decidable (FixHi σ) := by unfold FixHi; infer_instance

variable {σ}

/-! ## Expressions -/

namespace Ex

/-- Rename the registers of an expression. -/
def ren (σ : Equiv.Perm (Fin HS)) : Ex HS → Ex HS
  | reg i => reg (σ i)
  | lit c => lit c
  | add a b => add (a.ren σ) (b.ren σ)
  | neg a => neg (a.ren σ)
  | mul a b => mul (a.ren σ) (b.ren σ)
  | shl a r => shl (a.ren σ) (r.ren σ)
  | quo a r => quo (a.ren σ) (r.ren σ)
  | rem a r => rem (a.ren σ) (r.ren σ)
  | pow2 a => pow2 (a.ren σ)
  | sel c x y => sel (c.ren σ) (x.ren σ) (y.ren σ)
  | ndiv a b => ndiv (a.ren σ) (b.ren σ)
  | nmod a b => nmod (a.ren σ) (b.ren σ)

variable (σ)

variable {σ}

end Ex

/-! ## Register programs -/

namespace RS

def ren (σ : Equiv.Perm (Fin HS)) : RS → RS
  | skip => skip
  | asg i e => asg (σ i) (e.ren σ)
  | push i e => push (σ i) (e.ren σ)
  | pop i => pop (σ i)
  | seq a b => NLogN.RS.seq (a.ren σ) (b.ren σ)
  | ite c a b => ite (c.ren σ) (a.ren σ) (b.ren σ)
  | rep i n b => rep (σ i) (n.ren σ) (b.ren σ)

variable (σ)

variable {σ}

end RS

/-! ## Operations -/

/-! ## Register programs with multiplication -/

namespace RM

open RS (decSt decE asgCost)

def ren (σ : Equiv.Perm (Fin HS)) : RM → RM
  | rs p => rs (p.ren σ)
  | mul k i j => mul (σ k) (σ i) (σ j)
  | seq a b => NLogN.RM.seq (a.ren σ) (b.ren σ)
  | ite c a b => ite (c.ren σ) (a.ren σ) (b.ren σ)
  | rep i n b => rep (σ i) (n.ren σ) (b.ren σ)
  | op o => op (o.ren σ)
  | prc q f c pre w hf hio => NLogN.RM.prc q (fun t => (f (t.rn σ⁻¹)).rn σ) (fun t => c (t.rn σ⁻¹))
      (fun t => pre (t.rn σ⁻¹)) (w.map σ)
      (fun t j hj => by
        have hj' : σ.symm j ∉ w := fun h => hj (List.mem_map.2 ⟨_, h, by simp⟩)
        rw [IState.rn_G', hf _ _ hj', IState.rn_G']
        simp [Equiv.Perm.inv_def])
      (fun t => by simpa using hio (t.rn σ⁻¹))

variable (σ)

variable {σ}

/-! ## Reading off specifications -/

end RM

namespace RS

end RS

/-! ## Building permutations from swaps -/

/-- The product `swap a₁ b₁ * swap a₂ b₂ * ⋯` (the last swap is applied first). -/
def swapsPerm : List (Fin HS × Fin HS) → Equiv.Perm (Fin HS)
  | [] => 1
  | (a, b) :: l => Equiv.swap a b * swapsPerm l

/-- All swapped stacks are below `55` (a `Bool` check, closed by `rfl`/`decide` on literals). -/
def swapsLow (l : List (Fin HS × Fin HS)) : Bool := l.all fun p => p.1.val < 55 && p.2.val < 55

end NLogN

end

section
-- Source module: NLogN.HIface
/-!
# Interface for multiplication calls inside one recursive step

`MulBnd M m Φm`: the multiplication procedure accepts every pair of operands of at most `m` bits
and costs at most `Φm` on them.  Subroutines state their costs as
`(linear part) + (number of calls) · Φm`.
-/

namespace NLogN

open IntMul

def MulBnd (M : MulSpec) (m Φm : ℕ) : Prop :=
  ∀ a b : ℤ, zsz a ≤ m → zsz b ≤ m → M.dom a b ∧ M.Φ a b ≤ Φm

namespace RM

end RM

end NLogN

end

section
-- Source module: NLogN.SSLoops
/-!
# Generic counted loops over item arrays
-/

namespace NLogN

open IntMul

variable {P : ℕ} {I : List Sym} {iprog : Fin P → IStmt NS P}

/-- Decrement the counter on top of `C` (uses `T3`). -/
def decC (C : Fin NS) : IStmt NS P := .seq (.op (.const T3 (-1))) (.op (.add C T3 C))

/-- Move `k` integers from `X` to `Y` (reversing them); `k` is the counter on top of `C`. -/
def mvN (X Y C : Fin NS) : IStmt NS P :=
  .loop C (fun c => c != .zero) (.seq (.op (.mv X Y)) (decC C))

def ints (A : List ℤ) : List Item := A.map .int

@[simp] theorem ints_nil : ints [] = [] := rfl
@[simp] theorem ints_cons (a : ℤ) (A : List ℤ) : ints (a :: A) = .int a :: ints A := rfl

/-- All entries have at most `m` bits. -/
def Bnd (m : ℕ) (A : List ℤ) : Prop := ∀ a ∈ A, zsz a ≤ m

end NLogN

end