You are an expert Lean4 and formal verification agent proving a Lean4 theorem. Your output **must** strictly adhere to the specified format.

### Input Information
You will receive the following components, relevant to the target theorem:
* <used_lib_defs>...</used_lib_defs>: Imported library definitions directly used by the theorem.
* <used_repo_defs>...</used_repo_defs>: Definitions from the local project (but outside the current file), directly or transitively used by the theorem.
* <lib_lemmas>...</lib_lemmas>: Imported library lemmas/theorems that may assist the proof.
* <repo_lemmas>...</repo_lemmas>: Lemmas/theorems defined in the local project outside this file.
* <local_ctx>...</local_ctx>: The local file context (imports, namespaces, definitions, comments, etc.) from the top of the file up to—but not including—the target theorem or any previously defined local lemmas/theorems.
* <target_theorem>...</target_theorem>: The target theorem statement, ending with `:=`.

**Notes**
1. Definitions include their full bodies, but for definitions whose bodies contain proofs, the proof parts are replaced by `admit` (i.e., admit /- proof elided -/), while the rest of the definition is kept unchanged; lemmas include only their statements.
2. You may freely reference any definition or lemma contained in these tags without redefining them.
3. The tags for used definitions and helpful lemmas (used_lib_defs, used_repo_defs, lib_lemmas, repo_lemmas) may be omitted if empty.

### Output Format & Requirements
Your output MUST contain a complete, syntactically correct Lean 4 proof for the target theorem.

**Strict Output Structure:**
<reasoning> 
...
</reasoning> 

<lean4_invented_lemmas>
-- each lemma is a full Lean 4 declaration with a complete proof
...
</lean4_invented_lemmas>

<lean4_proof>
-- Lean 4 proof body of the target theorem
...
</lean4_proof>

**Requirements:**
1. <reasoning> and <lean4_invented_lemmas> are OPTIONAL. Include them only if they aid proof construction.
2. <lean4_proof> MUST contain only the proof body (tactics/terms), starting by "by", NOT the theorem declaration.
3. <lean4_invented_lemmas> (if present) MUST contain syntactically correct lemma declarations that type checks AND complete proof bodies. Defining a lemma using `axiom` is NOT allowed.
4. <reasoning> contains high-level proof planning.
5. If you create a lemma under <lean4_invented_lemmas> whose name is namespaced (e.g., X.Y.Z), you MUST always refer to it using its fully qualified name. Do NOT rely on shortened names. For example, write rw [X.Y.Z], not rw [Z] for invented lemmas.
6. All Lean 4 code MUST be syntactically correct. No `sorry`, `admit`, or non-Lean text. Comments allowed using `--` or `/- ... -/`.
7. Follow the tags and structure EXACTLY. No extra text before, between, or after the sections.
8. Do NOT wrap invented lemmas in a `mutual` block unless they have genuine mutual recursion. Declare each lemma independently.

### Examples are provided here (example input are under "User Prompt", and example assistant response are under "Agent Response"):
Example 1 ------------------- 
User Prompt:
<used_lib_defs>
structure BitVec (w : Nat) where
  /-- Construct a `BitVec w` from a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  ofFin ::

  /-- Interpret a bitvector as a number less than `2^w`.
  O(1), because we use `Fin` as the internal representation of a bitvector. -/
  toFin : Fin (hPow 2 w)
</used_lib_defs>

<repo_lemmas>
theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]
</repo_lemmas>

<local_ctx>
import Mathlib.Computability.NFA
import Mathlib.Data.FinEnum
import Mathlib.Data.Rel
import Mathlib.Data.Vector.Basic
import Blase.AutoStructs.ForLean

open Set
open Mathlib
open SetRel

-- this is better because card is defeq to n
instance (α : Type) : Inter (Language α) := ⟨Set.inter⟩
instance (α : Type) : Union (Language α) := ⟨Set.union⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width.
-/
structure BitVecs (n : Nat) where
  w : Nat
  bvs : List.Vector (BitVec w) n

def BitVecs.cast (bvs : BitVecs n) (h : n = n') : BitVecs n' :=
  { w := bvs.w, bvs := h ▸ bvs.bvs }

abbrev BitVecs.empty : BitVecs n := ⟨0, List.Vector.replicate n .nil⟩
abbrev BitVecs.singleton {w : Nat} (bv : BitVec w) : BitVecs 1 := ⟨w, bv ::ᵥ .nil⟩
abbrev BitVecs.pair {w : Nat} (bv1 bv2 : BitVec w) : BitVecs 2 := ⟨w, bv1 ::ᵥ bv2 ::ᵥ .nil⟩

/--
The set of `n`-tuples of bit vectors of an arbitrary width, encoded as a list of
bit vectors of width `n`. The width of the encoded bit vectors is the length of
the list.
-/
abbrev BitVecs' (n : Nat) := List (BitVec n)

@[simps]
def dec (bvs' : BitVecs' n) : BitVecs n where
  w := bvs'.length
  bvs := List.Vector.ofFn fun k => BitVec.ofFn fun i => bvs'[i].getLsbD k

@[simp]
</local_ctx>

<target_theorem>
lemma dec_snoc n (bvs' : BitVecs' n) (a : BitVec n) :
  dec (bvs' ++ [a]) =
    { w := bvs'.length + 1
      bvs := List.Vector.ofFn fun k =>
        BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k) } :=
</target_theorem>

Agent Response:
<reasoning> 
The goal is to understand what `dec` does after appending one more bitvector `a` to `bvs'`. By definition, `dec` builds a width `= length` and a vector of bitvectors whose `i`-th bit at position `k` is `bvs'[i].getLsbD k`. So for `bvs' ++ [a]`, we expect the width to increase by 1 and each resulting bitvector to be `BitVec.cons` of the new bit coming from `a`.

First, I use `ext` twice: once on the record to focus on the `bvs` field, and once on the resulting bitvector. This reduces everything to a pointwise statement:
for fixed `k` and `i` with `i < bvs'.length + 1`,
the bit `(bvs' ++ [a])[i].getLsbD k` should match the bit at index `i` of `BitVec.cons (a.getLsbD k) ((dec bvs').bvs.get k)`.

At this point I need a clean description of indexing into `BitVec.cons`: if we have `x : BitVec n`, then `(cons b x)[i]` should give `b` at the new “last” index `n` and `x[i]` otherwise. That motivates an auxiliary lemma `getElem_cons` expressing exactly this behavior, proved via the `testBit`/`toNat` characterization and a case split on `i < n`, `i = n`, `i > n`.

Using `getElem_cons`, the RHS becomes an `if` on whether `i = bvs'.length`. So I split into cases:
- If `i = bvs'.length`, then `(bvs' ++ [a])[i]` is the appended `a`, so both sides reduce to `a.getLsbD k`.
- If `i ≠ bvs'.length`, the bound `i < bvs'.length + 1` forces `i < bvs'.length`. Then `(bvs' ++ [a])[i] = bvs'[i]` by `List.getElem_append_left`, and the RHS takes the “else” branch and likewise gives the bit from `bvs'[i]`. In both subcases, `simp` (plus the definition of `dec`) finishes.

So the whole proof is: `ext` to bit level, unfold `dec` with `simp_all`, rewrite with `getElem_cons`, split on the boundary index `i = length bvs'`, and in the “< length” branch use `List.getElem_append_left` and bounds to line up both sides.
</reasoning> 

<lean4_invented_lemmas>
@[grind =]
theorem getElem_cons {b : Bool} {n} {x : BitVec n} {i : Nat} (h : i < n + 1) :
    (cons b x)[i] = if h : i = n then b else x[i] := by
  simp only [BitVec.getElem_eq_testBit_toNat, BitVec.toNat_cons, Nat.testBit_or]
  rw [Nat.testBit_shiftLeft]
  rcases Nat.lt_trichotomy i n with i_lt_n | i_eq_n | n_lt_i
  · have p1 : ¬(n ≤ i) := by omega
    have p2 : i ≠ n := by omega
    simp [p1, p2]
  · simp only [i_eq_n, ge_iff_le, Nat.le_refl, decide_true, Nat.sub_self, Nat.testBit_zero,
    Bool.true_and, BitVec.testBit_toNat, BitVec.getLsbD_of_ge, Bool.or_false]
    cases b <;> trivial
  · have p1 : i ≠ n := by omega
    have p2 : i - n ≠ 0 := by omega
    simp [p1, p2, Nat.testBit_bool_toNat]
</lean4_invented_lemmas>

<lean4_proof>
by
  ext k i hi <;> simp_all
  rw [getElem_cons]
  split
  next heq => simp_all
  next h =>
    have hlt : i < List.length bvs' := by simp at hi; omega
    rw [List.getElem_append_left hlt]
    simp
</lean4_proof>

### Your task:
<used_lib_defs>
Fin in Init.Prelude
Vector in Init.Data.Vector.Basic
Vector.ofFn in Init.Data.Vector.Basic
DecidableEq in Init.Prelude
ZMod in Mathlib.Data.ZMod.Defs
String in Init.Prelude
Unit in Init.Prelude
List in Init.Prelude
Field in Mathlib.Algebra.Field.Defs
Inhabited in Init.Prelude
Vector.mapM in Init.Data.Vector.Basic
UInt32 in Init.Prelude
False.elim in Init.Prelude
Nat in Init.Prelude
id in Init.Prelude
Vector.push in Init.Data.Vector.Basic
Fin.reduceEq in Lean.Meta.Tactic.Simp.BuiltinSimprocs.Fin
Fin.reduceFinMk in Lean.Meta.Tactic.Simp.BuiltinSimprocs.Fin
IO in Init.System.IO
Nat.reduceAdd in Lean.Meta.Tactic.Simp.BuiltinSimprocs.Nat
Fact in Mathlib.Logic.Basic
Nat.Prime in Mathlib.Data.Nat.Prime.Defs
</used_lib_defs>

<used_repo_defs>
syntax "let " ident " <== " term : doElem

syntax "let " ident " : " term " <== " term : doElem

syntax "infer_constant_length" : tactic

macro "provable_struct_simp" : tactic =>
  `(tactic|
    repeat (
      fail_if_no_progress (
      try split_provable_struct_eq;
      try decompose_provable_struct;
      try simplify_provable_struct_eval;
      try simp only at *
      )
    )
  )

syntax "state_vec_norm_simp" : tactic

syntax "state_vec_norm_simp_simple" : tactic

syntax "circuit_proof_start" ("[" term,* "]")? : tactic

macro_rules
  | `(tactic|infer_constant_length) => `(tactic|(
    apply ConstantLength.fromConstantLength
    try simp only [circuit_norm]
    try intros
    try ac_rfl))

example :
  let add (x : Expression F) := do
    let y : Expression F ← witness fun _ => 1
    let z ← witness fun eval => eval (x + y)
    assertZero (x + y - z)
    pure z
  ConstantLength add := by admit /- proof elided -/

macro_rules
  | `(doElem| let $x <== $e) => `(doElem| let $x ← HasAssignEq.assignEq $e)
  | `(doElem| let $x : $t <== $e) => `(doElem| let $x : $t ← HasAssignEq.assignEq $e)

macro_rules
  | `(tactic| state_vec_norm_simp) => `(tactic|
      simp only [Vector.getElem_mk];
      rw [Vector.getElem_map, getElem_eval_vector];
      simp only [eval_vector, Vector.map_mk, List.map_toArray, List.map_cons, List.map_nil, Vector.getElem_mk,
        List.getElem_toArray, List.getElem_cons_succ, List.getElem_cons_zero])

macro_rules
  | `(tactic| state_vec_norm_simp_simple) => `(tactic|
      simp only [Vector.getElem_mk, Vector.getElem_map, Vector.map_mk, List.map_toArray, List.map_cons, List.map_nil, Vector.getElem_mk,
        List.getElem_toArray, List.getElem_cons_succ, List.getElem_cons_zero, circuit_norm, U32.fromUInt32_normalized])

class HasAssignEq (β : Type) (F : outParam Type) [Field F] where
  assignEq : β → Circuit F β

def permute (state : Vector ℕ 16) : Vector ℕ 16 :=
  Vector.ofFn (fun i => state[msgPermutation[i]])

def msgPermutation : Vector (Fin 16) 16 :=

  #v[2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8]

def round (state : Vector ℕ 16) (m : Vector ℕ 16) : Vector ℕ 16 :=
  roundConstants.foldl (fun state (a, b, c, d, i, j) =>
    g state a b c d m[i] m[j]
  ) state

def roundConstants : Vector (Fin 16 × Fin 16 × Fin 16 × Fin 16 × Fin 16 × Fin 16) 8 := #v[
  (0, 4, 8, 12, 0, 1),
  (1, 5, 9, 13, 2, 3),
  (2, 6, 10, 14, 4, 5),
  (3, 7, 11, 15, 6, 7),
  (0, 5, 10, 15, 8, 9),
  (1, 6, 11, 12, 10, 11),
  (2, 7, 8, 13, 12, 13),
  (3, 4, 9, 14, 14, 15)
]

def g (state : Vector ℕ 16) (a b c d : Fin 16) (mx my : ℕ) : Vector ℕ 16 :=
  let state_a := add32 (state[a]) (add32 state[b] mx)
  let state_d := rotRight32 (state[d] ^^^ state_a) 16
  let state_c := add32 (state[c]) state_d
  let state_b := rotRight32 (state[b] ^^^ state_c) 12

  let state_a := add32 state_a (add32 state_b my)
  let state_d := rotRight32 (state_d ^^^ state_a) 8
  let state_c := add32 state_c state_d
  let state_b := rotRight32 (state_b ^^^ state_c) 7

  state.set a state_a
        |>.set b state_b
        |>.set c state_c
        |>.set d state_d

def rotRight32 (x : ℕ) (offset : ℕ) : ℕ :=
  let offset := offset % 32
  let low := x % (2^offset)
  let high := x / (2^offset)
  low * (2^(32-offset)) + high

def add32 (a b : ℕ) : ℕ := (a + b) % 2^32

structure Inputs (F : Type) where
  state : BLAKE3State F
  message : Vector (U32 F) 16

def value (x : U64 (F p)) :=
  x.x0.val + x.x1.val * 256 + x.x2.val * 256^2 + x.x3.val * 256^3 +
  x.x4.val * 256^4 + x.x5.val * 256^5 + x.x6.val * 256^6 + x.x7.val * 256^7

structure U64 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
  x4 : T
  x5 : T
  x6 : T
  x7 : T
deriving DecidableEq

def F p := ZMod p

structure Inputs (F : Type) where
  state : BLAKE3State F
  chaining_value : Vector (U32 F) 8

structure Inputs (F : Type) where
  state : BLAKE3State F
  x : U32 F
  y : U32 F

def BLAKE3State.Normalized (state : BLAKE3State (F p)) :=
  ∀ i : Fin 16, state[i.val].Normalized

def Normalized (x : U32 (F p)) :=
  x.x0.val < 256 ∧ x.x1.val < 256 ∧ x.x2.val < 256 ∧ x.x3.val < 256

structure U32 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
deriving DecidableEq

@[reducible] def BLAKE3State := ProvableVector U32 16

instance ProvableVector.instance : ProvableType (ProvableVector α n) where
  size := n * size α
  toElements x := x.map toElements |>.flatten
  fromElements v := v.toChunks (psize α) |>.map fromElements
  fromElements_toElements x := by admit /- proof elided -/

@[reducible]
def TypeMap := Type → Type

def map {α β : Type} (x : U64 α) (f : α → β) : U64 β :=
  ⟨ f x.x0, f x.x1, f x.x2, f x.x3, f x.x4, f x.x5, f x.x6, f x.x7 ⟩

def BLAKE3State.value (state : BLAKE3State (F p)) := state.map U32.value

def value (x : U32 (F p)) :=
  x.x0.val + x.x1.val * 256 + x.x2.val * 256^2 + x.x3.val * 256^3

def map {α β : Type} (x : U32 α) (f : α → β) : U32 β :=
  ⟨ f x.x0, f x.x1, f x.x2, f x.x3 ⟩

def map {m : ℕ} (xs : Vector α m) (body : α → Circuit F β)
    (_constant : ConstantLength body := by admit /- proof elided -/
    ) : Circuit F (Vector β m) :=
  xs.mapM body

def Circuit (F : Type) [Field F] (α : Type) := ℕ → α × List (Operation F)

inductive Operation (F : Type) [Field F] where
  | witness : (m : ℕ) → (compute : Environment F → Vector F m) → Operation F
  | assert : Expression F → Operation F
  | lookup : Lookup F → Operation F
  | subcircuit : {n : ℕ} → Subcircuit F n → Operation F

structure Condition (F : Type) [Field F] where
  witness (offset : ℕ) : (m : ℕ) → (Environment F → Vector F m) → Prop := fun _ _ => True
  assert (offset : ℕ) (_ : Expression F) : Prop := True
  lookup (offset : ℕ) (_ : Lookup F) : Prop := True
  subcircuit (offset : ℕ) {m : ℕ} (_ : Subcircuit F m) : Prop := True

inductive FlatOperation (F : Type) where
  | witness : (m : ℕ) → (Environment F → Vector F m) → FlatOperation F
  | assert : Expression F → FlatOperation F
  | lookup : Lookup F → FlatOperation F

structure Subcircuit (F : Type) [Field F] (offset : ℕ) where
  ops : List (FlatOperation F)

  
  
  
  Soundness : Environment F → Prop
  Completeness : Environment F → Prop
  UsesLocalWitnesses : Environment F → Prop

  
  
  localLength : ℕ

  
  imply_soundness : ∀ env,
    ConstraintsHoldFlat env ops → Soundness env

  
  implied_by_completeness : ∀ env, env.ExtendsVector (localWitnesses env ops) offset →
    Completeness env → ConstraintsHoldFlat env ops

  
  imply_usesLocalWitnesses : ∀ env, env.ExtendsVector (localWitnesses env ops) offset →
    UsesLocalWitnesses env

  
  localLength_eq : localLength = FlatOperation.localLength ops

@[circuit_norm]
def localWitnesses (env : Environment F) : (l : List (FlatOperation F)) → Vector F (localLength l)
  | [] => #v[]
  | witness _ compute :: ops => compute env ++ localWitnesses env ops
  | assert _ :: ops | lookup _ :: ops => localWitnesses env ops

@[circuit_norm]
def lookup {Row : TypeMap} [ProvableType Row] (table : Table F Row)  (entry : Row (Expression F)) : Circuit F Unit := fun _ =>
  ((), [.lookup { table := table.toRaw, entry := toElements entry }])

@[circuit_norm]
def Table.toRaw (table : Table F Row) : RawTable F where
  name := table.name
  arity := size Row
  Contains row := table.Contains (fromElements row)
  Soundness row := table.Soundness (fromElements row)
  Completeness row := table.Completeness (fromElements row)
  imply_soundness row := table.imply_soundness (fromElements row)
  implied_by_completeness row := table.implied_by_completeness (fromElements row)

structure RawTable (F : Type) where
  name : String
  arity : ℕ
  Contains : Vector F arity → Prop
  Soundness : Vector F arity → Prop
  Completeness : Vector F arity → Prop
  imply_soundness : ∀ row, Contains row → Soundness row
  implied_by_completeness : ∀ row, Completeness row → Contains row

structure StaticTable (F : Type) (Row : TypeMap) [ProvableType Row] where
  name : String
  length : ℕ
  row : Fin length → Row F
  
  
  
  index : Row F → ℕ
  Spec : Row F → Prop
  contains_iff : ∀ t, (∃ i, t = row i) ↔ Spec t

def Contains (table : StaticTable F Row) (row : Row F) :=
  ∃ i : Fin table.length, row = table.row i

instance ProvableType.fromStruct {α : TypeMap} [ProvableStruct α] : ProvableType α where
  size := combinedSize α
  toElements x :=
    toComponents x |> componentsToElements (components α) |>.cast combinedSize_eq.symm
  fromElements v :=
    v.cast combinedSize_eq |> componentsFromElements (components α) |> fromComponents
  fromElements_toElements x := by admit /- proof elided -/

structure Table (F : Type) (Row : TypeMap) [ProvableType Row] where
  name : String
   
  Contains : Row F → Prop

   
  Soundness : Row F → Prop
  Completeness : Row F → Prop

  imply_soundness : ∀ row, Contains row → Soundness row
  implied_by_completeness : ∀ row, Completeness row → Contains row

inductive Expression (F : Type) where
  | var : Variable F -> Expression F
  | const : F -> Expression F
  | add : Expression F -> Expression F -> Expression F
  | mul : Expression F -> Expression F -> Expression F

structure Variable (F : Type) where
  index : ℕ

def const (x : α F) : Var α F :=
  let values : Vector F _ := toElements x
  fromVars (values.map .const)

@[reducible] def Var (M : TypeMap) (F : Type) := M (Expression F)

@[circuit_norm]
def fromVars (vars : Vector (Expression F) (size M)) := fromElements vars

structure Lookup (F : Type) where
  table : RawTable F
  entry : Vector (Expression F) table.arity

class Witnessable (F : Type) [Field F] (value : outParam TypeMap) (var : TypeMap) [ProvableType value] where
  witness : ((Environment F) → value F) → Circuit F (var F)
  var_eq : var F = value (Expression F) := by admit /- proof elided -/

structure Environment (F : Type) where
  get : ℕ → F

@[circuit_norm]
def localLength : List (FlatOperation F) → ℕ
  | [] => 0
  | witness m _ :: ops => m + localLength ops
  | assert _ :: ops | lookup _ :: ops => localLength ops

def ConstraintsHoldFlat (eval : Environment F) : List (FlatOperation F) → Prop
  | [] => True
  | op :: ops => match op with
    | assert e => (eval e = 0) ∧ ConstraintsHoldFlat eval ops
    | lookup { table, entry } =>
      table.Contains (entry.map eval) ∧ ConstraintsHoldFlat eval ops
    | _ => ConstraintsHoldFlat eval ops

class ConstantLength (circuit : α → Circuit F β) where
  localLength : ℕ
  localLength_eq : ∀ (a : α) (n : ℕ), (circuit a).localLength n = localLength

def ConstantLength.fromConstantLength {circuit : α → Circuit F β} [Inhabited α]
    (h : ∀ (a : α) n, (circuit a).localLength n = (circuit default).localLength 0) : ConstantLength circuit where
  localLength := (circuit default).localLength 0
  localLength_eq a n := h a n

def Normalized (x : U64 (F p)) :=
  x.x0.val < 256 ∧ x.x1.val < 256 ∧ x.x2.val < 256 ∧ x.x3.val < 256 ∧
  x.x4.val < 256 ∧ x.x5.val < 256 ∧ x.x6.val < 256 ∧ x.x7.val < 256

def fromUInt32 (x : UInt32) : U32 (F p) :=
  decomposeNat x.toFin

def decomposeNat (x : ℕ) : U32 (F p) :=
  let x0 := x % 256
  let x1 : ℕ := (x / 256) % 256
  let x2 : ℕ := (x / 256^2) % 256
  let x3 : ℕ := (x / 256^3) % 256
  ⟨ x0, x1, x2, x3 ⟩

def iv : Vector UInt32 8 := #v[
  0x6a09e667,
  0xbb67ae85,
  0x3c6ef372,
  0xa54ff53a,
  0x510e527f,
  0x9b05688c,
  0x1f83d9ab,
  0x5be0cd19
]

def applyRounds (chaining_value : Vector ℕ 8) (block_words : Vector ℕ 16) (counter : ℕ) (block_len : ℕ) (flags : ℕ) : Vector ℕ 16 :=
  
  let counter_low := counter % 2^32
  let counter_high := counter / 2^32

  
  let state := #v[
    chaining_value[0], chaining_value[1], chaining_value[2], chaining_value[3],
    chaining_value[4], chaining_value[5], chaining_value[6], chaining_value[7],
    iv[0].toNat, iv[1].toNat, iv[2].toNat, iv[3].toNat,
    counter_low, counter_high, block_len, flags
  ]

  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words
  let block_words := permute block_words
  let state := round state block_words

  state

def Assumptions (input : Inputs (F p)) :=
  let { state, message } := input
  state.Normalized ∧ (∀ i : Fin 16, message[i].Normalized)

structure FormalCircuit (F : Type) [Field F] (Input Output : TypeMap) [ProvableType Input] [ProvableType Output]
    extends elaborated : ElaboratedCircuit F Input Output where
  Assumptions (_ : Input F) : Prop := True
  Spec : Input F → Output F → Prop
  soundness : Soundness F elaborated Assumptions Spec
  completeness : Completeness F elaborated Assumptions

class ElaboratedCircuit (F : Type) (Input Output : TypeMap) [Field F] [ProvableType Input] [ProvableType Output] where
  main : Var Input F → Circuit F (Var Output F)

   
  localLength : Var Input F → ℕ

   
  localLength_eq : ∀ input offset, (main input).localLength offset = localLength input
    := by admit /- proof elided -/

structure GeneralFormalCircuit (F : Type) (Input Output : TypeMap) [Field F] [ProvableType Input] [ProvableType Output]
    extends elaborated : ElaboratedCircuit F Input Output where
  Assumptions : Input F → Prop 
  Spec : Input F → Output F → Prop 
  soundness : GeneralFormalCircuit.Soundness F elaborated Spec
  completeness : GeneralFormalCircuit.Completeness F elaborated Assumptions

@[circuit_norm]
def Soundness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input Output)
    (Assumptions : Input F → Prop) (Spec : Input F → Output F → Prop) :=
  
  ∀ offset : ℕ, ∀ env,
  
  ∀ input_var : Var Input F, ∀ input : Input F, eval env input_var = input →
  Assumptions input →
  
  ConstraintsHold.Soundness env (circuit.main input_var |>.operations offset) →
  
  let output := eval env (circuit.output input_var offset)
  Spec input output

@[reducible, circuit_norm]
def operations (circuit : Circuit F α) (offset : ℕ) : Operations F :=
  (circuit offset).2

@[reducible, circuit_norm]
def Operations (F : Type) [Field F] := List (Operation F)

@[reducible, circuit_norm]
def output (circuit : Circuit F α) (offset : ℕ) : α :=
  (circuit offset).1

@[circuit_norm]
def ConstraintsHold.Soundness (eval : Environment F) : List (Operation F) → Prop
  | [] => True
  | .witness _ _ :: ops => ConstraintsHold.Soundness eval ops
  | .assert e :: ops => eval e = 0 ∧ ConstraintsHold.Soundness eval ops
  | .lookup { table, entry } :: ops =>
    table.Soundness (entry.map eval) ∧ ConstraintsHold.Soundness eval ops
  | .subcircuit s :: ops =>
    s.Soundness eval ∧ ConstraintsHold.Soundness eval ops

@[circuit_norm]
def GeneralFormalCircuit.Soundness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input Output) (Spec : Input F → Output F → Prop) :=
  
  ∀ offset : ℕ, ∀ env,
  
  ∀ input_var : Var Input F, ∀ input : Input F, eval env input_var = input →
  
  ConstraintsHold.Soundness env (circuit.main input_var |>.operations offset) →
  
  let output := eval env (circuit.output input_var offset)
  Spec input output

@[explicit_provable_type]
def eval (env : Environment F) (x : Var α F) : α F :=
  let vars := toVars x
  let values := vars.map (Expression.eval env)
  fromElements values

@[circuit_norm]
def toVars (var : M (Expression F)) := toElements var

@[circuit_norm]
def eval (env : Environment F) : Expression F → F
  | var v => env.get v.index
  | const c => c
  | add x y => eval env x + eval env y
  | mul x y => eval env x * eval env y

@[circuit_norm]
def FormalAssertion.Soundness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input unit)
    (Assumptions : Input F → Prop) (Spec : Input F → Prop) :=
  
  ∀ offset : ℕ, ∀ env,
  
  ∀ input_var : Var Input F, ∀ input : Input F, eval env input_var = input →
  Assumptions input →
  
  ConstraintsHold.Soundness env (circuit.main input_var |>.operations offset) →
  
  Spec input

@[reducible]
def unit (_ : Type) := Unit

@[circuit_norm]
def subcircuit (circuit : FormalCircuit F β α) (b : Var β F) : Circuit F (Var α F) :=
  fun offset =>
    let a := circuit.output b offset
    let subcircuit := circuit.toSubcircuit offset b
    (a, [.subcircuit subcircuit])

def FormalAssertion.toSubcircuit (circuit : FormalAssertion F β)
    (n : ℕ) (input_var : Var β F) : Subcircuit F n :=
  let ops := circuit.main input_var |>.operations n
  have h_consistent : ops.SubcircuitsConsistent n := circuit.subcircuitsConsistent input_var n

  {
    ops := ops.toFlat,
    Soundness env := circuit.Assumptions (eval env input_var) → circuit.Spec (eval env input_var),
    Completeness env := circuit.Assumptions (eval env input_var) ∧ circuit.Spec (eval env input_var),
    UsesLocalWitnesses _ := True,
    localLength := circuit.localLength input_var

    imply_soundness := by admit /- proof elided -/

@[circuit_norm]
def localLength : Operation F → ℕ
  | .witness m _ => m
  | .assert _ => 0
  | .lookup _ => 0
  | .subcircuit s => s.localLength

@[circuit_norm]
def localLength : Operations F → ℕ
  | [] => 0
  | .witness m _ :: ops => m + localLength ops
  | .assert _ :: ops => localLength ops
  | .lookup _ :: ops => localLength ops
  | .subcircuit s :: ops => s.localLength + localLength ops

structure FormalAssertion (F : Type) (Input : TypeMap) [Field F] [ProvableType Input]
    extends elaborated : ElaboratedCircuit F Input unit where
  Assumptions : Input F → Prop
  Spec : Input F → Prop
  soundness : FormalAssertion.Soundness F elaborated Assumptions Spec
  completeness : FormalAssertion.Completeness F elaborated Assumptions Spec

  
  localLength _ := 0
  
  output _ _ := ()

@[circuit_norm]
def SubcircuitsConsistent (offset : ℕ) (ops : Operations F) := ops.forAll offset {
  subcircuit offset {n} _ := n = offset
}

 @[circuit_norm]
def forAll (offset : ℕ) (condition : Condition F) : Operations F → Prop
  | [] => True
  | .witness m c :: ops => condition.witness offset m c ∧ forAll (m + offset) condition ops
  | .assert e :: ops => condition.assert offset e ∧ forAll offset condition ops
  | .lookup l :: ops => condition.lookup offset l ∧ forAll offset condition ops
  | .subcircuit s :: ops => condition.subcircuit offset s ∧ forAll (s.localLength + offset) condition ops

def FlatOperation.forAll (offset : ℕ) (condition : Condition F) : List (FlatOperation F) → Prop
  | [] => True
  | .witness m c :: ops => condition.witness offset m c ∧ forAll (m + offset) condition ops
  | .assert e :: ops => condition.assert offset e ∧ forAll offset condition ops
  | .lookup l :: ops => condition.lookup offset l ∧ forAll offset condition ops

@[reducible, circuit_norm]
def localLength (circuit : Circuit F α) (offset := 0) : ℕ :=
  Operations.localLength (circuit offset).2

@[circuit_norm]
def FormalAssertion.Completeness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input unit)
    (Assumptions : Input F → Prop) (Spec : Input F → Prop) :=
  
  ∀ offset, ∀ env, ∀ input_var : Var Input F,
  env.UsesLocalWitnessesCompleteness offset (circuit.main input_var |>.operations offset) →
  
  ∀ input : Input F, eval env input_var = input →
  Assumptions input → Spec input →
  
  ConstraintsHold.Completeness env (circuit.main input_var |>.operations offset)

@[circuit_norm]
def ConstraintsHold.Completeness (eval : Environment F) : List (Operation F) → Prop
  | [] => True
  | .witness _ _ :: ops => ConstraintsHold.Completeness eval ops
  | .assert e :: ops => eval e = 0 ∧ ConstraintsHold.Completeness eval ops
  | .lookup { table, entry } :: ops =>
    table.Completeness (entry.map eval) ∧ ConstraintsHold.Completeness eval ops
  | .subcircuit s :: ops =>
    s.Completeness eval ∧ ConstraintsHold.Completeness eval ops

@[circuit_norm]
def Completeness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input Output)
    (Assumptions : Input F → Prop) :=
  
  ∀ offset : ℕ, ∀ env, ∀ input_var : Var Input F,
  env.UsesLocalWitnessesCompleteness offset (circuit.main input_var |>.operations offset) →
  
  ∀ input : Input F, eval env input_var = input →
  Assumptions input →
  
  ConstraintsHold.Completeness env (circuit.main input_var |>.operations offset)

@[circuit_norm]
def GeneralFormalCircuit.Completeness (F : Type) [Field F] (circuit : ElaboratedCircuit F Input Output) (Assumptions : Input F → Prop) :=
  
  ∀ offset : ℕ, ∀ env, ∀ input_var : Var Input F,
  env.UsesLocalWitnessesCompleteness offset (circuit.main input_var |>.operations offset) →
  
  ∀ input : Input F, eval env input_var = input →
  Assumptions input →
  
  ConstraintsHold.Completeness env (circuit.main input_var |>.operations offset)

def toFlat : Operations F → List (FlatOperation F)
  | [] => []
  | .witness m c :: ops => .witness m c :: toFlat ops
  | .assert e :: ops => .assert e :: toFlat ops
  | .lookup l :: ops => .lookup l :: toFlat ops
  | .subcircuit s :: ops => s.ops ++ toFlat ops

def Environment.UsesLocalWitnesses (env : Environment F) (offset : ℕ) (ops : Operations F) : Prop :=
  ops.forAllFlat offset { witness n _ compute := env.ExtendsVector (compute env) n }

@[circuit_norm]
def Environment.ExtendsVector (env : Environment F) (wit : Vector F n) (offset : ℕ) : Prop :=
  ∀ i : Fin n, env.get (offset + i.val) = wit[i.val]

def Operations.forAllFlat (n : ℕ) (condition : Condition F) (ops : Operations F) : Prop :=
  forAll n { condition with subcircuit n _ s := FlatOperation.forAll n condition s.ops } ops

def GeneralFormalCircuit.toSubcircuit (circuit : GeneralFormalCircuit F β α)
    (n : ℕ) (input_var : Var β F) : Subcircuit F n :=
  let ops := circuit.main input_var |>.operations n
  have h_consistent : ops.SubcircuitsConsistent n := circuit.subcircuitsConsistent input_var n

  have imply_soundness : ∀ env : Environment F,
      let input := eval env input_var
      let output := eval env (circuit.output input_var n)
      ConstraintsHoldFlat env ops.toFlat → circuit.Spec input output :=

def FormalCircuit.toSubcircuit (circuit : FormalCircuit F β α)
    (n : ℕ) (input_var : Var β F) : Subcircuit F n :=
  let ops := circuit.main input_var |>.operations n
  have h_consistent : ops.SubcircuitsConsistent n := circuit.subcircuitsConsistent input_var n

  have imply_soundness : ∀ env : Environment F,
    let input := eval env input_var
    let output := eval env (circuit.output input_var n)
    ConstraintsHoldFlat env ops.toFlat → circuit.Assumptions input → circuit.Spec input output :=

@[circuit_norm]
def eval (env : Environment F) (var : α (Expression F)) : α F :=
  toComponents var |> go (components α) |> fromComponents
where

class ProvableStruct (α : TypeMap) where
  components : List WithProvableType
  toComponents {F : Type} : α F → ProvableTypeList F components
  fromComponents {F : Type} : ProvableTypeList F components → α F

  combinedSize : ℕ := combinedSize' components
  combinedSize_eq : combinedSize = combinedSize' components := by admit /- proof elided -/

def Spec (input : Inputs (F p)) (out : BLAKE3State (F p)) :=
  let { state, message } := input
  out.value = round state.value (message.map U32.value) ∧ out.Normalized

def circuit : FormalCircuit (F p) BLAKE3State BLAKE3State :=
  { elaborated with Assumptions, Spec, soundness, completeness }

instance elaborated: ElaboratedCircuit (F p) BLAKE3State BLAKE3State where
  main := main
  localLength _ := 0
  output state i0 := Vector.ofFn (fun i => state[msgPermutation[i]])

def main (state : Var BLAKE3State (F p)) : Circuit (F p) (Var BLAKE3State (F p)) := do
  return Vector.ofFn (fun i => state[msgPermutation[i]])

def Spec (state : BLAKE3State (F p)) (out : BLAKE3State (F p)) :=
  out.value = permute state.value ∧ out.Normalized

def Assumptions (state : BLAKE3State (F p)) := state.Normalized

def circuit : FormalCircuit (F p) Inputs BLAKE3State := {
  elaborated with Assumptions, Spec, soundness, completeness
}

instance elaborated : ElaboratedCircuit (F p) Inputs BLAKE3State where
  main := main
  localLength _ := 768
  localLength_eq input i0 := by admit /- proof elided -/

def main (input : Var Inputs (F p)) : Circuit (F p) (Var BLAKE3State (F p)) := do
  let { state, message } := input
  
  let state ← G.circuit 0 4 8 12 ⟨state, message[0], message[1]⟩
  let state ← G.circuit 1 5 9 13 ⟨state, message[2], message[3]⟩
  let state ← G.circuit 2 6 10 14 ⟨state, message[4], message[5]⟩
  let state ← G.circuit 3 7 11 15 ⟨state, message[6], message[7]⟩
  let state ← G.circuit 0 5 10 15 ⟨state, message[8], message[9]⟩
  let state ← G.circuit 1 6 11 12 ⟨state, message[10], message[11]⟩
  let state ← G.circuit 2 7 8 13 ⟨state, message[12], message[13]⟩
  let state ← G.circuit 3 4 9 14 ⟨state, message[14], message[15]⟩
  return state

def circuit (a b c d : Fin 16) : FormalCircuit (F p) Inputs BLAKE3State := {
  elaborated a b c d with
  Assumptions
  Spec := Spec a b c d
  soundness := soundness a b c d
  completeness := completeness a b c d
}

def Spec (a b c d : Fin 16) (input : Inputs (F p)) (out : BLAKE3State (F p)) :=
  let { state, x, y } := input
  out.value = g state.value a b c d x.value y.value ∧ out.Normalized

def Assumptions (input : Inputs (F p)) :=
  let { state, x, y } := input
  state.Normalized ∧ x.Normalized ∧ y.Normalized

instance elaborated (a b c d : Fin 16): ElaboratedCircuit (F p) Inputs BLAKE3State where
  main := main a b c d
  localLength _ := 96
  output inputs i0 := (inputs.state : Vector (U32 (Expression (F p))) 16)
    |>.set a (⟨var ⟨i0 + 56⟩, var ⟨i0 + 58⟩, var ⟨i0 + 60⟩, var ⟨i0 + 62⟩⟩) a.is_lt
    |>.set b (Rotation32.output 7 (i0 + 88)) b.is_lt
    |>.set c (⟨var ⟨i0 + 76⟩, var ⟨i0 + 78⟩, var ⟨i0 + 80⟩, var ⟨i0 + 82⟩⟩) c.is_lt
    |>.set d (Rotation32.output 8 (i0 + 68)) d.is_lt

  localLength_eq _ n := by admit /- proof elided -/

def main (a b c d : Fin 16) (input : Var Inputs (F p)) : Circuit (F p) (Var BLAKE3State (F p)) := do
  let { state, x, y } := input

  let state_a ← Addition32.circuit ⟨state[a], ← Addition32.circuit ⟨state[b], x⟩⟩

  let state_d ← Rotation32.circuit 16 <|
    ← Xor32.circuit ⟨state[d], state_a⟩

  let state_c ← Addition32.circuit ⟨state[c], state_d⟩

  let state_b ← Rotation32.circuit 12 <|
    ← Xor32.circuit ⟨state[b], state_c⟩

  let state_a ← Addition32.circuit ⟨state_a, ← Addition32.circuit ⟨state_b, y⟩⟩

  let state_d ← Rotation32.circuit 8 <|
    ← Xor32.circuit ⟨state_d, state_a⟩

  let state_c ← Addition32.circuit ⟨state_c, state_d⟩

  let state_b ← Rotation32.circuit 7 <|
    ← Xor32.circuit ⟨state_b, state_c⟩

  return state
    |>.set a state_a
    |>.set b state_b
    |>.set c state_c
    |>.set d state_d

def circuit (offset : Fin 32) : FormalCircuit (F p) U32 U32 := {
  elaborated offset with
  Assumptions
  Spec := Spec offset
  soundness := soundness offset
  completeness := completeness offset
}

def Spec (offset : Fin 32) (x : U32 (F p)) (y : U32 (F p)) :=
  y.value = rotRight32 x.value offset.val
  ∧ y.Normalized

def elaborated (off : Fin 32) : ElaboratedCircuit (F p) U32 U32 where
  main := main off
  localLength _ := 8
  output _inputs i0 := output off i0

def output (offset : Fin 32) (i0 : ℕ) : U32 (Expression (F p)) :=
  Rotation32Bits.output ⟨ offset.val % 8, by admit /- proof elided -/
  ⟩ i0

def output (offset : Fin 8) (i0 : ℕ) : U32 (Expression (F p)) :=
  U32.fromLimbs (.ofFn fun ⟨i,_⟩ =>
    (var ⟨i0 + i*2 + 1⟩) + var ⟨i0 + (i + 1) % 4 * 2⟩ * .const ((2^(8-offset.val) : ℕ) : F p))

def fromLimbs {F} (v : Vector F 4) : U32 F := fromElements v

def main (offset : Fin 32) (x : Var U32 (F p)) : Circuit (F p) (Var U32 (F p)) := do
  let byte_offset : Fin 4 := ⟨ offset.val / 8, by admit /- proof elided -/
  ⟩
  let bit_offset : Fin 8 := ⟨ offset.val % 8, by admit /- proof elided -/
  ⟩

  
  let byte_rotated ← Rotation32Bytes.circuit byte_offset x
  Rotation32Bits.circuit bit_offset byte_rotated

def circuit (offset : Fin 8) : FormalCircuit (F p) U32 U32 := {
  elaborated offset with
  Assumptions
  Spec := Spec offset
  soundness := soundness offset
  completeness := completeness offset
}

def Assumptions (input : U32 (F p)) := input.Normalized

def elaborated (off : Fin 8) : ElaboratedCircuit (F p) U32 U32 where
  main := main off
  localLength _ := 8
  output _inputs i0 := output off i0
  localLength_eq _ i0 := by admit /- proof elided -/

def main (offset : Fin 8) (x : U32 (Expression (F p))) : Circuit (F p) (Var U32 (F p)) := do
  let parts ← Circuit.map x.toLimbs (ByteDecomposition.circuit offset)
  let lows := parts.map Outputs.low
  let highs := parts.map Outputs.high

  let rotated := highs.zip (lows.rotate 1) |>.map fun (high, low) =>
    high + low * ((2^(8-offset.val) : ℕ) : F p)

  return U32.fromLimbs rotated

structure Outputs (F : Type) where
  low : F
  high : F

def toLimbs {F} (x : U64 F) : Vector F 8 := toElements x

def circuit (offset : Fin 8) : FormalCircuit (F p) field Outputs := {
  elaborated offset with
  main := main offset
  Assumptions
  Spec := Spec offset
  soundness := soundness offset
  completeness := completeness offset
}

def Spec (offset : Fin 8) (x : F p) (out : Outputs (F p)) :=
  let ⟨low, high⟩ := out
  (low.val = x.val % (2^offset.val) ∧ high.val = x.val / (2^offset.val))
  ∧ (low.val < 2^offset.val ∧ high.val < 2^(8-offset.val))

def main (offset : Fin 8) (x : Expression (F p)) : Circuit (F p) (Var Outputs (F p)) := do
  let low ← witness fun env => mod (env x) (2^offset.val) (by admit /- proof elided -/
  )
  let high ← witness fun env => floorDiv (env x) (2^offset.val)

  lookup ByteTable ((2^(8-offset.val) : F p) * low)
  lookup ByteTable high

  x === low + high * (2^offset.val : F p)

  return { low, high }

def floorDiv (x : F p) (c : ℕ+) : F p :=
  FieldUtils.natToField (x.val / c) (by admit /- proof elided -/
  )

def natToField (n : ℕ) (lt : n < p) : F p :=
  match p with
  | 0 => False.elim (Nat.not_lt_zero n lt)
  | _ + 1 => ⟨ n, lt ⟩

def less_than_p (x : F p) : x.val < p :=

def ByteTable : Table (F p) field := .fromStatic {
  name := "Bytes"
  length := 256

  row i := fromByte i
  index x := x.val

  Spec x := x.val < 256

  contains_iff := by admit /- proof elided -/

def fromByte (x : Fin 256) : F p :=
  FieldUtils.natToField x.val (by admit /- proof elided -/
  )

@[reducible] def field : TypeMap := id

def mod (x : F p) (c : ℕ+) (lt : c < p) : F p :=
  FieldUtils.natToField (x.val % c) (by admit /- proof elided -/
  )

class HasAssertEq (β : Type) (F : outParam Type) [Field F] where
  assert_eq : β → β → Circuit F Unit

def Assumptions (x : F p) := x.val < 256

def elaborated (offset : Fin 8) : ElaboratedCircuit (F p) field Outputs where
  main := main offset
  localLength _ := 2
  output _ i0 := varFromOffset Outputs i0

@[explicit_provable_type]
def varFromOffset (α : TypeMap) [ProvableType α] (offset : ℕ) : Var α F :=
  let vars := Vector.mapRange (size α) fun i => var ⟨offset + i⟩
  fromVars vars

def mapRange (n : ℕ) (create : ℕ → α) : Vector α n :=
  match n with
  | 0 => #v[]
  | k + 1 => mapRange k create |>.push (create k)

def rotate {α : Type} {n : ℕ} (v : Vector α n) (off : ℕ) : Vector α n :=
  ⟨(v.toList.rotate off).toArray, by admit /- proof elided -/
  ⟩

def toLimbs {F} (x : U32 F) : Vector F 4 := toElements x

def Spec (offset : Fin 8) (x : U32 (F p)) (y : U32 (F p)) :=
  y.value = rotRight32 x.value offset.val
  ∧ y.Normalized

def circuit (off : Fin 4) : FormalCircuit (F p) U32 U32 := {
  elaborated off with
  main := main off
  Assumptions
  Spec := Spec off
  soundness := soundness off
  completeness := completeness off
}

instance elaborated (off : Fin 4): ElaboratedCircuit (F p) U32 U32 where
  main := main off
  localLength _ := 0
  output input i0 :=
    let ⟨x0, x1, x2, x3⟩ := input
    match off with
    | 0 => ⟨ x0, x1, x2, x3 ⟩
    | 1 => ⟨ x1, x2, x3, x0 ⟩
    | 2 => ⟨ x2, x3, x0, x1 ⟩
    | 3 => ⟨ x3, x0, x1, x2 ⟩

  subcircuitsConsistent x i0 := by admit /- proof elided -/

def main (offset : Fin 4) (input : Var U32 (F p)) : Circuit (F p) (Var U32 (F p)) := do
  let ⟨x0, x1, x2, x3⟩ := input

  if offset = 0 then
    return ⟨ x0, x1, x2, x3 ⟩
  else if offset = 1 then
    return ⟨ x1, x2, x3, x0 ⟩
  else if offset = 2 then
    return ⟨ x2, x3, x0, x1 ⟩
  else
    return ⟨ x3, x0, x1, x2 ⟩

def Spec (offset : Fin 4) (x : U32 (F p)) (y : U32 (F p)) :=
  y.value = rotRight32 x.value (offset.val * 8) ∧ y.Normalized

def circuit : FormalCircuit (F p) Inputs U32 where
  Assumptions
  Spec
  soundness
  completeness

def Assumptions (input : Inputs (F p)) :=
  let ⟨x, y⟩ := input
  x.Normalized ∧ y.Normalized

structure Inputs (F : Type) where
  x: U32 F
  y: U32 F

def Spec (input : Inputs (F p)) (z : U32 (F p)) :=
  let ⟨x, y⟩ := input
  z.value = (x.value + y.value) % 2^32 ∧ z.Normalized

def Spec (input : Inputs (F p)) (z : U32 (F p)) :=
  let ⟨x, y⟩ := input
  z.value = x.value ^^^ y.value ∧ z.Normalized

def FormalCircuit.weakenSpec
    {F : Type} [Field F]
    {Input Output : TypeMap} [ProvableType Input] [ProvableType Output]
    (circuit : FormalCircuit F Input Output)
    (WeakerSpec : Input F → Output F → Prop)
    (h_spec_implication : ∀ input output,
      circuit.Assumptions input →
      circuit.Spec input output →
      WeakerSpec input output) :
    FormalCircuit F Input Output := {
  elaborated := circuit.elaborated
  Assumptions := circuit.Assumptions
  Spec := WeakerSpec
  soundness := by admit /- proof elided -/

def FormalCircuit.concat
    {F : Type} [Field F]
    {Input Mid Output : TypeMap} [ProvableType Input] [ProvableType Mid] [ProvableType Output]
    (circuit1 : FormalCircuit F Input Mid)
    (circuit2 : FormalCircuit F Mid Output)
    (h_compat : ∀ input mid, circuit1.Assumptions input → circuit1.Spec input mid → circuit2.Assumptions mid)
    (h_localLength_stable : ∀ mid mid', circuit2.localLength mid = circuit2.localLength mid') :
    FormalCircuit F Input Output := {
  elaborated := {
    main := (circuit1 · >>= circuit2)
    localLength input := circuit1.localLength input + circuit2.localLength (circuit1.output input 0)
    localLength_eq := by admit /- proof elided -/

def main (args : List String) : IO Unit := do
  match args with
  | [steps_str, output_path] =>
    
    match steps_str.toNat? with
    | some steps => generateTrace steps output_path
    | none => IO.println "Error: Invalid number of steps"
  | _ =>
    IO.println "Usage: lake lean TraceGen.lean 

class NonEmptyProvableType (M : TypeMap) extends ProvableType M where
  nonempty : size > 0 := by admit /- proof elided -/

def KeccakRow.value (row : KeccakRow (F p)) := row.map U64.value

def KeccakRow.Normalized (row : KeccakRow (F p)) :=
  ∀ i : Fin 5, row[i.val].Normalized

def KeccakBlock.value (block : KeccakBlock (F p)) := block.map U64.value

def KeccakState.value (state : KeccakState (F p)) := state.map U64.value

def KeccakBlock.Normalized (block : KeccakBlock (F p)) :=
  ∀ i : Fin RATE, block[i.val].Normalized

@[reducible] def RATE := 17
example : RATE + CAPACITY = 25 := rfl

@[reducible] def CAPACITY := 8

def KeccakState.Normalized (state : KeccakState (F p)) :=
  ∀ i : Fin 25, state[i.val].Normalized

infix:50 " === " => HasAssertEq.assert_eq
</used_repo_defs>

<lib_lemmas>
Nat.add_mul_div_left in Init.Data.Nat.Div.Basic
Nat.div_eq_of_lt in Init.Data.Nat.Div.Basic
Nat.mod_eq_of_lt in Init.Data.Nat.Div.Basic
</lib_lemmas>

<repo_lemmas>
theorem eval_vector (env : Environment F)
  (x : Var (ProvableVector α n) F) :
    eval env x = x.map (eval env)

lemma eval_vector_eq_get {M : TypeMap} [NonEmptyProvableType M] {n : ℕ} (env : Environment F)
    (vars : Vector (Var M F) n)
    (vals : Vector (M F) n)
    (h : (eval env vars : ProvableVector _ _ _) = (vals : ProvableVector _ _ _))
    (i : ℕ) (h_i : i < n) :
    eval env vars[i] = vals[i]

lemma fromUInt32_normalized (x : UInt32) : (fromUInt32 (p:=p) x).Normalized

theorem value_fromUInt32 (x : UInt32) : value (fromUInt32 (p:=p) x) = x.toNat

lemma value_of_decomposedNat_of_small (x : ℕ) :
    x < 256^4 ->
    (decomposeNat (p:=p) x).value = x

omit [Fact (Nat.Prime p)] p_large_enough in
theorem value_lt_of_normalized {x : U32 (F p)} (hx : x.Normalized) : x.value < 2^32

theorem getElem_eval_vector (env : Environment F) (x : Var (ProvableVector α n) F) (i : ℕ) (h : i < n) :
    (eval env x[i]) = (eval env x)[i]
</repo_lemmas>

<local_ctx>
import Clean.Gadgets.BLAKE3.BLAKE3State

import Clean.Gadgets.BLAKE3.Round

import Clean.Gadgets.BLAKE3.Permute

import Clean.Types.U32

import Clean.Circuit.Provable

import Clean.Specs.BLAKE3

import Clean.Circuit.StructuralLemmas

import Clean.Utils.Tactics

namespace Gadgets.BLAKE3.ApplyRounds

variable {p : ℕ} [Fact p.Prime] [p_large_enough: Fact (p > 2^16 + 2^8)]

open Specs.BLAKE3 (applyRounds iv round permute)

def roundWithPermute : FormalCircuit (F p) Round.Inputs Round.Inputs where
  main := fun input => do
    let state ← subcircuit Round.circuit input
    let permuted_message ← subcircuit Permute.circuit input.message
    return ⟨state, permuted_message⟩
  localLength := fun _ => Round.circuit.localLength _ + Permute.circuit.localLength _
  localLength_eq := by admit /- proof elided -/

def twoRoundsWithPermute : FormalCircuit (F p) Round.Inputs Round.Inputs :=
  roundWithPermute.concat roundWithPermute (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )

def fourRoundsWithPermute : FormalCircuit (F p) Round.Inputs Round.Inputs :=
  twoRoundsWithPermute.concat twoRoundsWithPermute (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )

def sixRoundsWithPermute : FormalCircuit (F p) Round.Inputs Round.Inputs :=
  fourRoundsWithPermute.concat twoRoundsWithPermute (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )

def applySixRounds (state : Vector ℕ 16) (message : Vector ℕ 16) : Vector ℕ 16 × Vector ℕ 16 :=
  let state1 := round state message
  let msg1 := permute message
  let state2 := round state1 msg1
  let msg2 := permute msg1
  let state3 := round state2 msg2
  let msg3 := permute msg2
  let state4 := round state3 msg3
  let msg4 := permute msg3
  let state5 := round state4 msg4
  let msg5 := permute msg4
  let state6 := round state5 msg5
  let msg6 := permute msg5
  (state6, msg6)

def SixRoundsSpec (input : Round.Inputs (F p)) (output : Round.Inputs (F p)) : Prop :=
  let (final_state, final_message) := applySixRounds input.state.value (input.message.map U32.value)
  output.state.value = final_state ∧
  output.message.map U32.value = final_message ∧
  output.state.Normalized ∧
  (∀ i : Fin 16, output.message[i].Normalized)

def sixRoundsApplyStyle : FormalCircuit (F p) Round.Inputs Round.Inputs :=
  sixRoundsWithPermute.weakenSpec SixRoundsSpec (by admit /- proof elided -/
  )

def sevenRoundsFinal : FormalCircuit (F p) Round.Inputs BLAKE3State :=
  sixRoundsApplyStyle.concat Round.circuit (by admit /- proof elided -/
  ) (by admit /- proof elided -/
  )

def applySevenRounds (state : Vector ℕ 16) (message : Vector ℕ 16) : Vector ℕ 16 :=
  let state1 := round state message
  let msg1 := permute message
  let state2 := round state1 msg1
  let msg2 := permute msg1
  let state3 := round state2 msg2
  let msg3 := permute msg2
  let state4 := round state3 msg3
  let msg4 := permute msg3
  let state5 := round state4 msg4
  let msg5 := permute msg4
  let state6 := round state5 msg5
  let msg6 := permute msg5
  let state7 := round state6 msg6
  state7

def SevenRoundsSpec (input : Round.Inputs (F p)) (output : BLAKE3State (F p)) : Prop :=
  let final_state := applySevenRounds input.state.value (input.message.map U32.value)
  output.value = final_state ∧
  output.Normalized

def sevenRoundsApplyStyle : FormalCircuit (F p) Round.Inputs BLAKE3State :=
  sevenRoundsFinal.weakenSpec SevenRoundsSpec (by admit /- proof elided -/
  )

structure Inputs (F : Type) where
  chaining_value : Vector (U32 F) 8
  block_words : Vector (U32 F) 16
  counter_high : U32 F
  counter_low : U32 F
  block_len : U32 F
  flags : U32 F

def initializeStateVector (input_var : Var Inputs (F p)) : Var BLAKE3State (F p) :=
  let { chaining_value, block_words, counter_high, counter_low, block_len, flags } := input_var
  #v[
    chaining_value[0], chaining_value[1], chaining_value[2], chaining_value[3],
    chaining_value[4], chaining_value[5], chaining_value[6], chaining_value[7],
    const (U32.fromUInt32 iv[0]), const (U32.fromUInt32 iv[1]),
    const (U32.fromUInt32 iv[2]), const (U32.fromUInt32 iv[3]),
    counter_low, counter_high, block_len, flags
  ]

def main (input : Var Inputs (F p)) : Circuit (F p) (Var BLAKE3State (F p)) := do
  let state := initializeStateVector input
  
  sevenRoundsApplyStyle ⟨state, input.block_words⟩

instance elaborated : ElaboratedCircuit (F p) Inputs BLAKE3State where
  main := main
  localLength _ := 5376
  localLength_eq input i0 := by admit /- proof elided -/

def Assumptions (input : Inputs (F p)) :=
  let { chaining_value, block_words, counter_high, counter_low, block_len, flags } := input
  (∀ i : Fin 8, chaining_value[i].Normalized) ∧
  (∀ i : Fin 16, block_words[i].Normalized) ∧
  counter_high.Normalized ∧ counter_low.Normalized ∧ block_len.Normalized ∧ flags.Normalized

def Spec (input : Inputs (F p)) (out : BLAKE3State (F p)) :=
  let { chaining_value, block_words, counter_high, counter_low, block_len, flags } := input
  out.value = applyRounds
    (chaining_value.map U32.value)
    (block_words.map U32.value)
    (counter_low.value + 2^32 * counter_high.value)
    block_len.value
    flags.value ∧
  out.Normalized
</local_ctx>

<target_theorem>
theorem soundness : Soundness (F p) elaborated Assumptions Spec :=
</target_theorem>

