TokenArray.lean

import ZipCommon.Binary
import Zip.Native.Wide
import Std.Tactic.BVDecide

/-! # Packed token container for the LZ77 stream

The DEFLATE encoder emits a stream of `LZ77Token`s.  In production these are
bit-packed one-per-`UInt32` by `packTok`/`unpackTok` (see `Zip/Native/Deflate.lean`)
and accumulated in an `Array UInt32`.  Lean stores `Array UInt32` with 8-byte
boxed slots, so 45.9M tokens cost ~367 MB — the dominant page-fault / system-time
cost on a whole-tar compress.

This file introduces `TokenArray`, a newtype over a `ByteArray` that stores each
token in exactly 4 little-endian bytes (halving the footprint to 4 B/token).  It
is a *pure addition*: nothing consumes it yet.  Later stages retype the LZ77
producers and consumers to `TokenArray`, one `*Correct.lean` file at a time.

Halving the token bytes only pays off if the per-token pack/unpack stays as
cheap as the boxed `Array UInt32`'s one call per element.  Two runtime measures
keep it there (both measured on silesia.tar, same-worktree sandwich):

* the accumulators seed with `emptyWithCapacity data.size`, so the backing
  `ByteArray` never grows by doubling — the doubling transient (up to ~2× the
  live token bytes) is what regressed *peak* RSS in the un-presized first
  attempt (`#2866`/`#2867`);
* `push`/`get` compile (via `@[csimp]`, not the unverified `@[implemented_by]`)
  to the single-call wide primitives `ByteArray.pushUInt64LE` /
  `ByteArray.ugetUInt32LE` (`Zip/Native/Wide.lean`) rather than four
  `ByteArray.push` / `getElem` calls each.  The 4→1 call-count reduction matches
  the boxed one-call-per-token cost; without it the byte layout costs ~+3–13%
  compress CPU (worst on the greedy tier, where the token stream is longest).
  The reference bodies stay the four-op forms, and `pushUInt32LE_eq_impl` /
  `get_eq_impl` prove the wide forms equal them, so the compiled runtime swaps
  under a kernel-checked proof and every bridge/refinement proof below is
  unchanged.

Net on silesia.tar: peak RSS −25% to −30% (L1–L6) with compress wall neutral to
~3% faster and byte-identical output.

The proof-facing model stays `Array UInt32`: `TokenArray.toArray` exposes the
`Array UInt32` view, and the four bridge lemmas

* `empty_toArray`   `TokenArray.empty.toArray = #[]`
* `push_toArray`    `(ta.push w).toArray = ta.toArray.push w`
* `size_toArray`    `ta.size = ta.toArray.size`
* `get_toArray`     `ta.get i h = ta.toArray[i]`

let the existing packed-layer proofs port by substituting the array-map lemmas
(e.g. `Array.map_push`) with `push_toArray`, while the additive-delta lemmas in
`DeflateFreqsAdditive` keep reasoning on `Array UInt32`.

The `TokenArray` invariant `bytes.size % 4 = 0` is carried as a proof field
(erased at runtime, so the footprint is exactly the `ByteArray`).  The invariant
is precisely 4-byte *alignment* of the underlying `ByteArray` — the constructor
accepts any byte length that is a multiple of four, not only streams built via
`empty`/`push`.  In practice every `TokenArray` value is produced by `empty`,
`push`, and `extract`, each of which discharges the alignment obligation, and it
is that alignment invariant (not the construction provenance) that makes the
bridge lemmas unconditional. -/

namespace ByteArray

/-- Runtime implementation of `pushUInt32LE`: append the four little-endian
    bytes of `v` in a single FFI call via the wide store `pushUInt64LE` (which,
    on the exclusive-with-slack hot path, is one wide store into capacity slack
    plus a size bump — the C writes eight byte-stores the optimizer coalesces).
    The four low bytes of `v.toUInt64` LSB-first are exactly the four bytes of
    `v`, so this agrees with the reference body below byte-for-byte.  The point
    is the call count: 4 `ByteArray.push` calls per token collapse to 1, matching
    the boxed `Array.push`'s one-call-per-token cost while keeping the 4 B/token
    footprint. -/
@[inline] def pushUInt32LEImpl (b : ByteArray) (v : UInt32) : ByteArray :=
  b.pushUInt64LE v.toUInt64 4 (by
    have h : (4 : USize).toNat = 4 :=
      USize.toNat_ofNat_of_lt (Nat.lt_of_lt_of_le (show (4:Nat) < 2 ^ 32 by omega) USize.le_size)
    omega)

/-- Append a `UInt32` as four little-endian bytes (low byte first), matching
    `Binary.writeUInt32LE` / `Binary.readUInt32LE`.  The reference body (four
    `push`es) is the specification; the compiled path (`pushUInt32LEImpl`) does it
    in one wide FFI call, and `pushUInt32LE_eq_impl` (a `@[csimp]` theorem, below)
    *proves* the two agree — so the runtime swap carries no trust of its own; the
    only trusted edge left is the wide-store `@[extern]` on `pushUInt64LE` itself.

    TODO upstream to lean-zip-common (alongside `Binary.writeUInt32LE`). -/
def pushUInt32LE (b : ByteArray) (v : UInt32) : ByteArray :=
  (((b.push v.toUInt8).push (v >>> 8).toUInt8).push (v >>> 16).toUInt8).push (v >>> 24).toUInt8

/-- The wide-store implementation computes exactly the four-`push` reference body,
    so `@[csimp]` justifies compiling `pushUInt32LE` as `pushUInt32LEImpl` with a
    kernel-checked proof instead of an unverified `@[implemented_by]` assertion. -/
@[csimp] theorem pushUInt32LE_eq_impl : @pushUInt32LE = @pushUInt32LEImpl := by
  funext b v
  have e0 : (v.toUInt64).toUInt8 = v.toUInt8 := by bv_decide
  have e1 : (v.toUInt64 >>> 8).toUInt8 = (v >>> 8).toUInt8 := by bv_decide
  have e2 : (v.toUInt64 >>> 8 >>> 8).toUInt8 = (v >>> 16).toUInt8 := by bv_decide
  have e3 : (v.toUInt64 >>> 8 >>> 8 >>> 8).toUInt8 = (v >>> 24).toUInt8 := by bv_decide
  have h4 : (4 : USize).toNat = 4 :=
    USize.toNat_ofNat_of_lt (Nat.lt_of_lt_of_le (show (4:Nat) < 2 ^ 32 by omega) USize.le_size)
  simp only [pushUInt32LEImpl, pushUInt64LE, h4, pushLEBytes, pushUInt32LE, e0, e1, e2, e3]

@[simp] theorem size_pushUInt32LE (b : ByteArray) (v : UInt32) :
    (b.pushUInt32LE v).size = b.size + 4 := by
  simp [pushUInt32LE, ByteArray.size_push]

/-- Underlying byte list of `pushUInt32LE`: the original bytes followed by the
    four little-endian bytes of `v`. -/
theorem pushUInt32LE_data_toList (b : ByteArray) (v : UInt32) :
    (b.pushUInt32LE v).data.toList =
      b.data.toList ++
        [v.toUInt8, (v >>> 8).toUInt8, (v >>> 16).toUInt8, (v >>> 24).toUInt8] := by
  simp [pushUInt32LE, ByteArray.data_push, Array.toList_push]

private theorem length_data_toList (b : ByteArray) : b.data.toList.length = b.size := by
  rw [Array.length_toList]; rfl

/-- `pushUInt32LE` preserves the earlier bytes. -/
theorem getElem_pushUInt32LE_lt (b : ByteArray) (v : UInt32) {j : Nat}
    (hj : j < b.size) (h : j < (b.pushUInt32LE v).size) :
    (b.pushUInt32LE v)[j]'h = b[j]'hj := by
  rw [ByteArray.getElem_eq_getElem_data, ByteArray.getElem_eq_getElem_data,
      ← Array.getElem_toList, ← Array.getElem_toList]
  apply Option.some.inj
  rw [← List.getElem?_eq_getElem, ← List.getElem?_eq_getElem, pushUInt32LE_data_toList,
      List.getElem?_append_left (by rw [length_data_toList]; exact hj)]

/-- The `k`-th appended byte (`k < 4`) of `pushUInt32LE` sits at offset `b.size + k`. -/
theorem getElem_pushUInt32LE_offset (b : ByteArray) (v : UInt32) {k : Nat}
    (hk : k < 4) (h : b.size + k < (b.pushUInt32LE v).size) :
    (b.pushUInt32LE v)[b.size + k]'h =
      [v.toUInt8, (v >>> 8).toUInt8, (v >>> 16).toUInt8, (v >>> 24).toUInt8][k]'(by simpa using hk) := by
  rw [ByteArray.getElem_eq_getElem_data, ← Array.getElem_toList]
  apply Option.some.inj
  rw [← List.getElem?_eq_getElem, ← List.getElem?_eq_getElem, pushUInt32LE_data_toList,
      List.getElem?_append_right (by rw [length_data_toList]; omega)]
  simp

/-- The first appended byte of `pushUInt32LE` sits at offset `b.size` (the
    `k = 0` case of `getElem_pushUInt32LE_offset`, stated without the `+ 0`). -/
theorem getElem_pushUInt32LE_size (b : ByteArray) (v : UInt32)
    (h : b.size < (b.pushUInt32LE v).size) :
    (b.pushUInt32LE v)[b.size]'h = v.toUInt8 := by
  have := getElem_pushUInt32LE_offset b v (k := 0) (by omega) h
  simpa using this

/-- Reassembling the four little-endian bytes of a `UInt32` recovers it. -/
theorem uint32_le_roundtrip (v : UInt32) :
    v.toUInt8.toUInt32 ||| ((v >>> 8).toUInt8.toUInt32 <<< 8)
      ||| ((v >>> 16).toUInt8.toUInt32 <<< 16) ||| ((v >>> 24).toUInt8.toUInt32 <<< 24) = v := by
  bv_decide

end ByteArray

/-- A packed LZ77 token stream: one token per four little-endian bytes.
    The `aligned` field records that the byte length is a multiple of four; it is
    a `Prop` (erased at runtime), so the runtime footprint is just `bytes`. -/
structure TokenArray where
  bytes : ByteArray
  aligned : bytes.size % 4 = 0

namespace TokenArray

/-- The empty token stream. -/
def empty : TokenArray := ⟨ByteArray.empty, by simp [ByteArray.size_empty]⟩

/-- An empty token stream that pre-allocates room for `n` tokens (`4 * n` bytes).
    The matcher seeds its accumulator with this, mirroring the boxed model's
    `Array.emptyWithCapacity data.size`: the per-token `push`es then grow the
    backing `ByteArray` in place instead of reallocating by doubling.  Dropping
    this pre-sizing was the regression in the first unboxing attempt
    (`#2866`/`#2867`) — the doubling growth transient (up to ~2× the live token
    bytes momentarily) dominated *peak* RSS, and the realloc churn cost per-token
    CPU, so the byte-identical refactor regressed both axes.  In the model this is
    the empty stream (`ByteArray.emptyWithCapacity` ignores the capacity in Lean
    and only hints the runtime allocator), so it shares every bridge lemma with
    `empty`. -/
def emptyWithCapacity (n : Nat) : TokenArray :=
  ⟨ByteArray.emptyWithCapacity (4 * n), by
    show (ByteArray.emptyWithCapacity (4 * n)).size % 4 = 0
    rfl⟩

/-- Number of tokens (four bytes each). -/
def size (ta : TokenArray) : Nat := ta.bytes.size / 4

/-- Append one packed token (its little-endian `UInt32` word). -/
def push (ta : TokenArray) (w : UInt32) : TokenArray :=
  ⟨ta.bytes.pushUInt32LE w, by
    have := ta.aligned
    rw [ByteArray.size_pushUInt32LE]; omega⟩

theorem size_push (ta : TokenArray) (w : UInt32) : (ta.push w).size = ta.size + 1 := by
  have := ta.aligned
  simp only [push, size, ByteArray.size_pushUInt32LE]
  omega

private theorem byte_bound (ta : TokenArray) {i : Nat} (h : i < ta.size) :
    4 * i + 3 < ta.bytes.size := by
  have := ta.aligned
  simp only [size] at h
  omega

/-- A `Nat` below `USize.size` round-trips through `USize` unchanged (local copy
    of `Zip.Native.Deflate.toUSize_toNat_of_lt`, which lives downstream). -/
private theorem toUSize_toNat_of_lt {n : Nat} (h : n < USize.size) : n.toUSize.toNat = n := by
  simp only [Nat.toUSize]; exact Nat.mod_eq_of_lt h

/-- Runtime implementation of `get` via a single wide load.  When the backing
    `ByteArray` is `USize`-addressable — the runtime-checked `hsz` guard, always
    true for a real in-memory array — the four token bytes are read as one
    `ugetUInt32LE` FFI call instead of four bounds-checked `getElem`s; otherwise
    it falls back to the byte-recombination formula (the reference body below).
    Both branches compute the same `UInt32`, so this agrees with the model; the
    point, as with `pushUInt32LEImpl`, is the 4→1 call-count reduction on the
    hot token-consumer loops (emit, freq histogram, block-split sizing).  This
    mirrors the input-reader guard used by `lz77Greedy.hash3` (#2706). -/
@[inline] def getImpl (ta : TokenArray) (i : Nat) (h : i < ta.size) : UInt32 :=
  have hb := ta.byte_bound h
  if hsz : ta.bytes.size.toUSize.toNat = ta.bytes.size then
    ta.bytes.ugetUInt32LE (4 * i).toUSize (by
      have hds : ta.bytes.size < USize.size := by
        rw [← hsz]; exact USize.toNat_lt_two_pow_numBits _
      rw [toUSize_toNat_of_lt (show 4 * i < USize.size by omega)]; omega)
  else
    (ta.bytes[4 * i]'(by omega)).toUInt32
      ||| ((ta.bytes[4 * i + 1]'(by omega)).toUInt32 <<< 8)
      ||| ((ta.bytes[4 * i + 2]'(by omega)).toUInt32 <<< 16)
      ||| ((ta.bytes[4 * i + 3]'(by omega)).toUInt32 <<< 24)

/-- Read the `i`-th token as a little-endian `UInt32` (proven-in-bounds). -/
def get (ta : TokenArray) (i : Nat) (h : i < ta.size) : UInt32 :=
  have hb := ta.byte_bound h
  (ta.bytes[4 * i]'(by omega)).toUInt32
    ||| ((ta.bytes[4 * i + 1]'(by omega)).toUInt32 <<< 8)
    ||| ((ta.bytes[4 * i + 2]'(by omega)).toUInt32 <<< 16)
    ||| ((ta.bytes[4 * i + 3]'(by omega)).toUInt32 <<< 24)

/-- The wide-load implementation returns exactly the byte-recombination reference
    body on both branches of its `USize`-addressability guard, so `@[csimp]`
    compiles `get` as `getImpl` under a kernel-checked proof rather than an
    unverified `@[implemented_by]`; the only trusted edge left is the wide-load
    `@[extern]` on `ugetUInt32LE`. -/
@[csimp] theorem get_eq_impl : @get = @getImpl := by
  funext ta i h
  have hb := ta.byte_bound h
  rw [getImpl]
  split
  · rename_i hsz
    have hds : ta.bytes.size < USize.size := by
      rw [← hsz]; exact USize.toNat_lt_two_pow_numBits _
    have hidx : ((4 * i : Nat).toUSize).toNat = 4 * i :=
      toUSize_toNat_of_lt (show 4 * i < USize.size by omega)
    rw [get, ByteArray.ugetUInt32LE]
    simp only [hidx]
  · rw [get]

/-- The `Array UInt32` model: the token at each index. -/
def toArray (ta : TokenArray) : Array UInt32 :=
  Array.ofFn (n := ta.size) (fun i => ta.get i.val i.isLt)

/-! ## Bridge lemmas relating `TokenArray` to its `Array UInt32` model. -/

@[simp] theorem empty_toArray : empty.toArray = #[] := by
  apply Array.eq_empty_of_size_eq_zero
  simp only [toArray, Array.size_ofFn, size, empty, ByteArray.size_empty]

/-- `emptyWithCapacity` views to the empty `Array UInt32` — pre-sizing is a
    runtime allocation hint only, so it shares `empty`'s model.  Drops in for
    `empty_toArray` wherever the matcher seeds with a pre-sized accumulator. -/
@[simp] theorem emptyWithCapacity_toArray (n : Nat) : (emptyWithCapacity n).toArray = #[] := by
  apply Array.eq_empty_of_size_eq_zero
  have hsz : (ByteArray.emptyWithCapacity (4 * n)).size = 0 := rfl
  simp only [toArray, Array.size_ofFn, size, emptyWithCapacity, hsz, Nat.zero_div]

theorem size_toArray (ta : TokenArray) : ta.size = ta.toArray.size := by
  simp only [toArray, Array.size_ofFn]

theorem get_toArray (ta : TokenArray) (i : Nat) (h : i < ta.size) :
    ta.get i h = ta.toArray[i]'(by rw [← size_toArray]; exact h) := by
  simp only [toArray, Array.getElem_ofFn]

/-- `get` ignores a subsequent `push` on earlier tokens. -/
private theorem get_push_lt (ta : TokenArray) (w : UInt32) {i : Nat} (h : i < ta.size)
    (h' : i < (ta.push w).size) :
    (ta.push w).get i h' = ta.get i h := by
  have hb := ta.byte_bound h
  simp only [get, push]
  rw [ByteArray.getElem_pushUInt32LE_lt ta.bytes w (by omega),
      ByteArray.getElem_pushUInt32LE_lt ta.bytes w (by omega),
      ByteArray.getElem_pushUInt32LE_lt ta.bytes w (by omega),
      ByteArray.getElem_pushUInt32LE_lt ta.bytes w (by omega)]

/-- `get` at the fresh index reads back the just-pushed token. -/
private theorem get_push_eq (ta : TokenArray) (w : UInt32)
    (h' : ta.size < (ta.push w).size) :
    (ta.push w).get ta.size h' = w := by
  have haligned := ta.aligned
  have hsz : 4 * ta.size = ta.bytes.size := by simp only [size]; omega
  simp only [get, push, hsz]
  rw [ByteArray.getElem_pushUInt32LE_size ta.bytes w (by simp),
      ByteArray.getElem_pushUInt32LE_offset ta.bytes w (k := 1) (by omega) (by simp),
      ByteArray.getElem_pushUInt32LE_offset ta.bytes w (k := 2) (by omega) (by simp),
      ByteArray.getElem_pushUInt32LE_offset ta.bytes w (k := 3) (by omega) (by simp)]
  simpa using ByteArray.uint32_le_roundtrip w

/-- Byte-level slice on token boundaries: tokens `[i, j)` as a fresh
    `TokenArray`.  Because every token is exactly four bytes, the token range
    `[i, j)` is the byte range `[4·i, 4·j)`, and the alignment invariant is
    preserved (a difference of multiples of four).  This is the packed twin of
    `Array.extract` — the shared-block split family slices the token stream with
    it, and `extract_toArray` bridges it to the `Array UInt32` model. -/
def extract (ta : TokenArray) (i j : Nat) : TokenArray :=
  ⟨ta.bytes.extract (4 * i) (4 * j), by
    have := ta.aligned
    rw [ByteArray.size_extract]; omega⟩

@[simp] theorem size_extract (ta : TokenArray) (i j : Nat) :
    (ta.extract i j).size = min j ta.size - i := by
  have := ta.aligned
  simp only [extract, size, ByteArray.size_extract]
  omega

/-- Reading token `k` of `ta.extract i j` reads token `i + k` of `ta`. -/
theorem get_extract (ta : TokenArray) (i j k : Nat) (h : k < (ta.extract i j).size)
    (h' : i + k < ta.size) :
    (ta.extract i j).get k h = ta.get (i + k) h' := by
  simp only [get, extract]
  have key : ∀ (a b : Nat) (hla : a < (ta.bytes.extract (4 * i) (4 * j)).size)
      (hlb : b < ta.bytes.size), 4 * i + a = b →
      (ta.bytes.extract (4 * i) (4 * j))[a]'hla = ta.bytes[b]'hlb := by
    intro a b hla hlb hab
    rw [ByteArray.getElem_extract]; congr 1
  rw [key (4 * k) (4 * (i + k)) _ _ (by omega),
      key (4 * k + 1) (4 * (i + k) + 1) _ _ (by omega),
      key (4 * k + 2) (4 * (i + k) + 2) _ _ (by omega),
      key (4 * k + 3) (4 * (i + k) + 3) _ _ (by omega)]

/-- The load-bearing slice bridge: extracting a token range from the container
    matches extracting it from the `Array UInt32` model. -/
@[simp] theorem extract_toArray (ta : TokenArray) (i j : Nat) :
    (ta.extract i j).toArray = ta.toArray.extract i j := by
  apply Array.ext
  · rw [← size_toArray, size_extract, Array.size_extract, ← size_toArray]
  · intro k h1 h2
    have hk : k < (ta.extract i j).size := by rw [← size_toArray] at h1; exact h1
    have hik : i + k < ta.size := by
      rw [size_extract] at hk; omega
    rw [← get_toArray (ta.extract i j) k hk, get_extract ta i j k hk hik,
        Array.getElem_extract, get_toArray]

/-- The load-bearing bridge lemma: pushing a token onto the container matches
    pushing its word onto the `Array UInt32` model.  This is the drop-in
    replacement for `Array.map_push` in the ported packed-layer proofs. -/
@[simp] theorem push_toArray (ta : TokenArray) (w : UInt32) :
    (ta.push w).toArray = ta.toArray.push w := by
  apply Array.ext
  · rw [← size_toArray, size_push, size_toArray, Array.size_push]
  · intro i h1 h2
    have hlt : i < (ta.push w).size := by rw [← size_toArray] at h1; exact h1
    have hisize : i < ta.size + 1 := by rw [← size_push]; exact hlt
    rw [← get_toArray (ta.push w) i hlt]
    rcases Nat.lt_succ_iff_lt_or_eq.mp hisize with hi | hi
    · rw [get_push_lt ta w hi hlt, get_toArray ta i hi,
          Array.getElem_push_lt (show i < ta.toArray.size by rw [← size_toArray]; exact hi)]
    · subst hi
      rw [get_push_eq, Array.getElem_push, dif_neg (by rw [← size_toArray]; omega)]

end TokenArray