Polynomial-time Computability

import Mathlib.Data.Fintype.Vector
import Mathlib.Data.Nat.Bits
import Mathlib.Data.Nat.BinaryRec
import Mathlib.Analysis.SpecialFunctions.Pow.Real

-- leandown
-- [meta]
-- title = "Polynomial-time computability"
-- group = "Cryptography"
-- [content]

-- [hidden: Definitions]
-- `Bits` (= `BitString`) is the type of bit strings: `List Bool`.
--
-- `bitStringToNat` and `natToBitString` convert between bit strings and natural
-- numbers (LSB-first). `natKthRoot k n` computes $\lfloor n^{1/k} \rfloor$.
/-- The free monoid {0,1}* (Kleene star over the alphabet {0,1}). -/
abbrev Bits := List Bool

/-- Alias for `Bits`; used when emphasising the role of a value as a bit string
    rather than a list. -/
abbrev BitString := Bits

/-- Convert a bit string (LSB-first) to a natural number. -/
def bitStringToNat : Bits → ℕ
  | []      => 0
  | b :: bs => (if b then 1 else 0) + 2 * bitStringToNat bs

/-- Convert a natural number to its LSB-first bit string (using `Nat.bits`). -/
def natToBitString : ℕ → Bits := Nat.bits

/-- The integer k-th root: `⌊n^(1/k)⌋`. Poly-time computable for any fixed `k`. -/
noncomputable def natKthRoot (k n : ℕ) : ℕ :=
  Nat.floor ((n : ℝ) ^ ((1 : ℝ) / (k : ℝ)))
-- [/hidden]

-- # {{definition}}: Polynomial-time arithmetic operations
--
-- `PolyTimeArith` is an inductive predicate on functions $\mathbb{N} \to \mathbb{N}$.
-- A proof of `PolyTimeArith f` is a derivation that `f` is computable in time
-- polynomial in the **bit-length** of its input.
--
-- The constructors are:
-- - `id`, `const k`: identity and constants
-- - `addConst`, `subConst`: $O(n)$
-- - `mulConst`, `divConst`, `modConst`: $O(n^2)$
-- - `comp`, `mul`: closure under composition and pointwise product
-- - `kthRoot`: $\lfloor n^{1/k} \rfloor$ via Newton's method, $O(n^2)$
--
/-- Basic arithmetic operations on `ℕ` that are poly-time in the bit-length of
    their input. -/
inductive PolyTimeArith : (ℕ → ℕ) → Prop where
  | id                             : PolyTimeArith id
  | const (k : ℕ)                  : PolyTimeArith (fun _ => k)
  | addConst (k : ℕ)               : PolyTimeArith (· + k)
  | subConst (k : ℕ)               : PolyTimeArith (· - k)
  | mulConst (k : ℕ)               : PolyTimeArith (· * k)
  | divConst (k : ℕ) (hk : 0 < k) : PolyTimeArith (· / k)
  | modConst (k : ℕ) (hk : 0 < k) : PolyTimeArith (· % k)
  | comp {f g} : PolyTimeArith f → PolyTimeArith g → PolyTimeArith (f ∘ g)
  | mul  {f g} : PolyTimeArith f → PolyTimeArith g →
      PolyTimeArith (fun n => f n * g n)
  | kthRoot (k : ℕ) (hk : 0 < k)  : PolyTimeArith (natKthRoot k)

-- # {{definition}}: Polynomial-time computable functions on bit strings
--
-- `PolyTimeComputable` is an inductive predicate on `BitString → BitString`.
-- **Constructing a term of this type is the proof of poly-time computability** —
-- there is no separate complexity argument to supply.
--
-- The constructors are:
-- - `id`, `const c`: identity and constants
-- - `comp`: closure under composition
-- - `mapBool op`: apply a bitwise operation pointwise, $O(n)$
-- - `arithLift`: lift a `PolyTimeArith` operation via `bitStringToNat`/`natToBitString`
-- - `take k`, `drop k`: prefix/suffix, $O(n)$ for fixed $k$
-- - `arithTake`, `arithDrop`: prefix/suffix of length $\mathrm{op}(n)$
--   for a poly-time arithmetic function $\mathrm{op}$, $O(n)$
-- - `append`: concatenate outputs of two poly-time functions, $O(n)$
-- - `polyIter k`: iterate a poly-time step function $n^k$ times, $O(n^k \cdot T_f(n))$
--
/-- Poly-time computable functions on bit strings, defined inductively over
    natural combinators. Constructing a term of this type *is* the proof of
    poly-time computability. -/
inductive PolyTimeComputable : (Bits → Bits) → Prop where
  | id      : PolyTimeComputable id
  | const   (c : Bits) : PolyTimeComputable (fun _ => c)
  | comp    {f g} : PolyTimeComputable f → PolyTimeComputable g →
      PolyTimeComputable (f ∘ g)
  | mapBool (op : Bool → Bool) : PolyTimeComputable (List.map op)
  | arithLift {op} : PolyTimeArith op →
      PolyTimeComputable (fun x => natToBitString (op (bitStringToNat x)))
  /-- Take the first `k` bits — fixed prefix length. -/
  | take (k : ℕ) : PolyTimeComputable (List.take k)
  /-- Drop the first `k` bits — fixed prefix length. -/
  | drop (k : ℕ) : PolyTimeComputable (List.drop k)
  /-- Take the first `op(|x|)` bits, where `op` is a poly-time arithmetic function
      of the input length. -/
  | arithTake {op} : PolyTimeArith op →
      PolyTimeComputable (fun x => x.take (op x.length))
  /-- Drop the first `op(|x|)` bits, where `op` is a poly-time arithmetic function
      of the input length. -/
  | arithDrop {op} : PolyTimeArith op →
      PolyTimeComputable (fun x => x.drop (op x.length))
  /-- Take the first `⌊√|x|⌋` bits: $O(\sqrt{n})$ to compute the length, $O(n)$
      to scan. Special case of `arithTake` for the square-root function. -/
  | sqrtTake : PolyTimeComputable (fun x => x.take (Nat.sqrt x.length))
  /-- Drop the first `⌊√|x|⌋` bits. Special case of `arithDrop`. -/
  | sqrtDrop : PolyTimeComputable (fun x => x.drop (Nat.sqrt x.length))
  | append {f g} : PolyTimeComputable f → PolyTimeComputable g →
      PolyTimeComputable (fun x => f x ++ g x)
  /-- Iterate a poly-time step function `f` exactly `|x|^k` times.
      Total cost $O(n^k \cdot T_f(n))$, which is polynomial for fixed `k`. -/
  | polyIter (k : ℕ) {f} : PolyTimeComputable f →
      PolyTimeComputable (fun x => Nat.iterate f (x.length ^ k) x)

-- # {{theorem}}: `natToBitString` is a right inverse of `bitStringToNat`
--
-- i.e. `bitStringToNat (natToBitString n) = n` for all `n : ℕ`
--
-- **proof sketch:** induction on the binary representation of `n`
-- using `Nat.binaryRec'`, with the base case $n = 0$ and the inductive step
-- appending a bit via `Nat.bits_append_bit`.
--
/-- Round-trip: converting a natural number to its bit representation and back
    is the identity. -/
lemma bitStringToNat_natToBitString (n : ℕ) :
    bitStringToNat (natToBitString n) = n := by
  -- unfold `natToBitString` to `Nat.bits`
  simp only [natToBitString]
  -- induct on the binary representation of `n`
  induction n using Nat.binaryRec' with
  -- base case: `Nat.bits 0 = []` and `bitStringToNat [] = 0`
  | zero => simp [Nat.bits, bitStringToNat]
  -- inductive step: `n = bit b m`; use `Nat.bits_append_bit` to unfold,
  -- then the inductive hypothesis `ih : bitStringToNat (Nat.bits m) = m`,
  -- and case-split on `b` to close both branches by `omega`
  | bit b m h ih =>
    rw [Nat.bits_append_bit m b h]
    simp only [bitStringToNat, ih]
    cases b <;> simp [Nat.bit_true_apply, Nat.bit_false_apply]; omega

-- # {{theorem}}: `n^k` is poly-time for any fixed `k`
--
-- **proof:** by induction on `k`:
-- - $k = 0$: `fun n => 1` is `PolyTimeArith.const 1`
-- - $k + 1$: `fun n => n^{k+1} = n * n^k` is `PolyTimeArith.mul id (ofPow k)`
--
/-- `n^k` is poly-time for any fixed `k`, derived from repeated pointwise
    multiplication. -/
def PolyTimeArith.ofPow : ∀ k : ℕ, PolyTimeArith (fun n => n ^ k)
  | 0     => by simpa using PolyTimeArith.const 1
  | k + 1 =>
      -- rewrite `n^(k+1)` as `n * n^k` so `.mul .id (.ofPow k)` applies
      have h : (fun n : ℕ => n ^ (k + 1)) = (fun n : ℕ => n * n ^ k) :=
        funext fun n => by rw [pow_succ, mul_comm]
      h ▸ .mul .id (.ofPow k)