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:
* <all_available_defs>: Definitions available from imported modules (either from the current Lean repository or from Mathlib)
* <all_available_lemmas>: Lemmas/theorems available that may assist the proof.
* <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 available definitions and lemmas (all_available_defs, all_available_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:
<all_available_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)
</all_available_defs>

<all_available_lemmas>
theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) : (l₁ ++ l₂)[i] = l₁[i]
theorem getElem_eq_testBit_toNat (x : BitVec w) (i : Nat) (h : i < w) :
  x[i] = x.toNat.testBit i := rfl
@[simp, grind =] theorem toNat_cons (b : Bool) (x : BitVec w) :
    (cons b x).toNat = (b.toNat <<< w) ||| x.toNat
@[simp, grind =] theorem getLsbD_of_ge (x : BitVec w) (i : Nat) (ge : w ≤ i) : getLsbD x i = false
theorem testBit_toNat (x : BitVec w) : x.toNat.testBit i = x.getLsbD i
</all_available_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:
<all_available_defs>
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 Assumptions (input : U32 (F p)) := input.Normalized

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

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

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

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

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 :=

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 -/

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 :=

@[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])

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

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

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

@[reducible] def CAPACITY := 8

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

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

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 -/
  ⟩

contains_iff := by admit /- proof elided -/

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

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

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

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

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

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

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

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

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

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

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

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

imply_soundness := by admit /- proof elided -/

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

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 Assumptions (x : F p) := x.val < 256

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 elaborated (offset : Fin 8) : ElaboratedCircuit (F p) field Outputs where
  main := main offset
  localLength _ := 2
  output _ i0 := varFromOffset Outputs i0

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 iv : Vector UInt32 8 := #v[
  0x6a09e667,
  0xbb67ae85,
  0x3c6ef372,
  0xa54ff53a,
  0x510e527f,
  0x9b05688c,
  0x1f83d9ab,
  0x5be0cd19
]

def msgPermutation : Vector (Fin 16) 16 :=

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

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 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 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 permute (state : Vector ℕ 16) : Vector ℕ 16 :=
  Vector.ofFn (fun i => state[msgPermutation[i]])

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

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

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

@[reducible] def BLAKE3State := ProvableVector U32 16

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

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

@[reducible]
def TypeMap := Type → Type

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

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 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 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

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

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

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 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 circuit (offset : Fin 8) : FormalCircuit (F p) U32 U32 := {
  elaborated offset with
  Assumptions
  Spec := Spec offset
  soundness := soundness offset
  completeness := completeness offset
}

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

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 Lookup (F : Type) where
  table : RawTable F
  entry : Vector (Expression F) table.arity

@[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 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

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

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

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

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 Assumptions (input : Inputs (F p)) :=
  let { state, x, y } := input
  state.Normalized ∧ x.Normalized ∧ y.Normalized

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 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 main (state : Var BLAKE3State (F p)) : Circuit (F p) (Var BLAKE3State (F p)) := do
  return Vector.ofFn (fun i => state[msgPermutation[i]])

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

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

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

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

structure Variable (F : Type) where
  index : ℕ

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 Environment (F : Type) where
  get : ℕ → F

@[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

@[reducible]
def TypeMap := Type → Type

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

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

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

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

@[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

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

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

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

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

@[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 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 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

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

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 -/

@[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

@[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)

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

@[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

@[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)

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 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

@[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)

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

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 -/

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

@[reducible]
def TypeMap := Type → Type

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

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

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

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

@[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

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

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

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

@[reducible] def field : TypeMap := id

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 -/

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 -/

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

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 -/

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

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

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

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

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

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) Inputs BLAKE3State := {
  elaborated with Assumptions, Spec, soundness, completeness
}

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

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

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

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

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

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

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 -/

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

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

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 circuit : FormalCircuit (F p) Inputs U32 where
  Assumptions
  Spec
  soundness
  completeness

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

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

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

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

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

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

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 fromUInt32 (x : UInt32) : U32 (F p) :=
  decomposeNat x.toFin

def F p := ZMod p

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 mod (x : F p) (c : ℕ+) (lt : c < p) : F p :=
  FieldUtils.natToField (x.val % c) (by admit /- proof elided -/
  )

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

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

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

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

@[circuit_norm]
def localLength : List (FlatOperation F) → ℕ
  | [] => 0
  | witness m _ :: ops => m + localLength ops
  | assert _ :: ops | lookup _ :: ops => 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 Environment.ExtendsVector (env : Environment F) (wit : Vector F n) (offset : ℕ) : Prop :=
  ∀ i : Fin n, env.get (offset + i.val) = wit[i.val]

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

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

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

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

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

@[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 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

@[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

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

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

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

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

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

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 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 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

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

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 Assumptions (input : U32 (F p)) := input.Normalized

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

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 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
}

Assumptions input :=
    let ⟨c, s⟩ := input
    IsBool s

let ctx ← getLCtx
      for localDecl in ctx do
        if localDecl.isImplementationDetail then
          continue

AliasCheck.circuit bits

replaceMainGoal [currentGoal]

structure InductiveTable (F : Type) [Field F] (State Input : Type → Type) [ProvableType State] [ProvableType Input] where
   
  step : Var State F → Var Input F → Circuit F (Var State F)

soundness := by admit /- proof elided -/

def bind {α β} (f : Circuit F α) (g : α → Circuit F β) : Circuit F β := fun (n : ℕ) =>
  
  let (b, ops') := g (f n).1 (n + Operations.localLength (f n).2)
  (b, (f n).2 ++ ops')

ZMod in Mathlib.Data.ZMod.Defs

HDiv in Init.Prelude

let out <== c.map fun ci =>
    let a10 := (ci[3] - ci[2] - ci[1] + ci[0]) * s10
    let a1 := (ci[2] - ci[0]) * s[1]
    let a0 := (ci[1] - ci[0]) * s[0]
    let a := ci[0]
    a10 + a1 + a0 + a

def circuit : FormalAssertion (F p) BlockInput where
  main
  localLength_eq := by admit /- proof elided -/

def Environment.AgreesBelow (n : ℕ) (env env' : Environment F) :=
  ∀ i < n, env.get i = env'.get i

instance elaborated (off : Fin 8): ElaboratedCircuit (F p) U64 U64 where
  main := main off
  localLength _ := 0
  output input i0 :=
    let ⟨x0, x1, x2, x3, x4, x5, x6, x7⟩ := input
    match off with
    | 0 => ⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩
    | 1 => ⟨ x1, x2, x3, x4, x5, x6, x7, x0 ⟩
    | 2 => ⟨ x2, x3, x4, x5, x6, x7, x0, x1 ⟩
    | 3 => ⟨ x3, x4, x5, x6, x7, x0, x1, x2 ⟩
    | 4 => ⟨ x4, x5, x6, x7, x0, x1, x2, x3 ⟩
    | 5 => ⟨ x5, x6, x7, x0, x1, x2, x3, x4 ⟩
    | 6 => ⟨ x6, x7, x0, x1, x2, x3, x4, x5 ⟩
    | 7 => ⟨ x7, x0, x1, x2, x3, x4, x5, x6 ⟩
  subcircuitsConsistent x i0 := by admit /- proof elided -/

def MemoryAccessList.isConsistentOnline (accesses : MemoryAccessList) (h : accesses.isTimestampSorted) : Prop := match accesses with
  | [] => True 
  | (_timestamp, addr, readValue, _writeValue) :: rest =>
    
    readValue = MemoryAccessList.lastWriteValue rest (List.Sorted.of_cons h) addr
    ∧ MemoryAccessList.isConsistentOnline rest (List.Sorted.of_cons h)

List.tail in Init.Data.List.Basic

instance {x y : MemoryAccess} : Decidable (address_strict_timestamp_ordering x y) :=

instance : ProvableType U32 where
  size := 4
  toElements x := #v[x.x0, x.x1, x.x2, x.x3]
  fromElements v :=
    let ⟨ .mk [x0, x1, x2, x3], _ ⟩ := v
    ⟨ x0, x1, x2, x3 ⟩

Lean.Meta.mkAppM in Lean.Meta.AppBuilder

def main (row : Var KeccakRow (F p)) : Circuit (F p) (Var KeccakRow (F p)) := do
  let c0 ← Rotation64.circuit (64 - 1) row[1]
  let c0 ← Xor64.circuit ⟨row[4], c0⟩

let (instr_type, mode1, mode2, mode3) ← decodeInstruction rawInstrType

def ReadOnlyTableFromFunction
    {n : ℕ} (f : Fin n → (F p)) (h : n < p) [NeZero n] :
    Table (F p) fieldPair
  := .fromStatic {
  name := "ReadOnlyMemory"
  length := n
  row i := (i, f i)
  index := fun (i, _) => i.val
  Spec := fun (i, v) => v = f (Fin.ofNat n i.val) ∧ i.val < n
  contains_iff := by admit /- proof elided -/

BLAKE3ProcessBlocksStateNormalized.circuit state 
  BLAKE3BlockInputNormalized.circuit input

def DecodedInstructionType.isEncodedCorrectly (instrType : DecodedInstructionType (F p)) : Prop :=
  (instrType.isAdd = 1 ∧ instrType.isMul = 0 ∧ instrType.isStoreState = 0 ∧ instrType.isLoadState = 0) ∨
  (instrType.isAdd = 0 ∧ instrType.isMul = 1 ∧ instrType.isStoreState = 0 ∧ instrType.isLoadState = 0) ∨
  (instrType.isAdd = 0 ∧ instrType.isMul = 0 ∧ instrType.isStoreState = 1 ∧ instrType.isLoadState = 0) ∨
  (instrType.isAdd = 0 ∧ instrType.isMul = 0 ∧ instrType.isStoreState = 0 ∧ instrType.isLoadState = 1)

def Spec (n : ℕ) (input : fields n (F p)) (output : F p) : Prop :=
  output.val = (input.map (·.val)).foldl (· &&& ·) 1 ∧ IsBool output

Fin.reduceEq in Lean.Meta.Tactic.Simp.BuiltinSimprocs.Fin

modify in Init.Prelude

let decoded ← subcircuitWithAssertion decodeInstructionCircuit rawInstrType

instance explicit : ExplicitCircuits (Gadgets.Addition32Full.main (p:=pBabybear)) :=

def circuit : FormalAssertion (F p) U64 where
  main

Option in Init.Prelude

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

Assumptions := fun (x, y) => x.val < 2^n ∧ y.val < 2^n
  Spec := fun (x, y) output =>
    output = (if x.val <= y.val then 1 else 0)

def generateFieldBasedNames (baseName : Name) (fieldNames : Array Name) : Array Name :=
  fieldNames.map fun fieldName => baseName.appendAfter s!"_{fieldName}"

instance : ProvableType DecodedAddressingMode where
  size := 4
  toElements := fun { isDoubleAddressing, isApRelative, isFpRelative, isImmediate } => #v[isDoubleAddressing, isApRelative, isFpRelative,
    isImmediate]
  fromElements := fun elements => {
    isDoubleAddressing := elements[0],
    isApRelative := elements[1],
    isFpRelative := elements[2],
    isImmediate := elements[3]
  }

| everyRowExceptLast: TwoRowsConstraint S F → TableOperation S F

Vector.instGetElemNatLt in Init.Data.Vector.Basic

for (eqExpr, _, _) in equalities do
          
          if let some typeExpr := eqExpr.getArg? 0 then
            
            let typeExpr' ← whnf typeExpr
            match typeExpr'.getAppFn with
            | .const typeName _ =>
              
              let mkInjEqName := typeName ++ `mk ++ `injEq
              if env.contains mkInjEqName then
                
                
                if ← hasProvableStructInstance typeExpr' then
                  let lemmaIdent := mkIdent mkInjEqName
                  if !lemmasToApply.any (fun l => l.getId == mkInjEqName) then
                    lemmasToApply := lemmaIdent :: lemmasToApply
            | _ => pure ()

def thetaC (state : Vector ℕ 25) : Vector ℕ 5 :=
  #v[
    state[0] ^^^ state[1] ^^^ state[2] ^^^ state[3] ^^^ state[4],
    state[5] ^^^ state[6] ^^^ state[7] ^^^ state[8] ^^^ state[9],
    state[10] ^^^ state[11] ^^^ state[12] ^^^ state[13] ^^^ state[14],
    state[15] ^^^ state[16] ^^^ state[17] ^^^ state[18] ^^^ state[19],
    state[20] ^^^ state[21] ^^^ state[22] ^^^ state[23] ^^^ state[24]
  ]

def Run.isLeaf (R : Run S) (root leaf : S) : Prop :=
  R.reachable root leaf ∧ ∀ y, R (leaf, y) = 0

def not64_bytewise_value (x : U64 (F p)) : U64 (F p) := x.map (fun x => 255 - x)

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

let state_capacity := Vector.mapFinRange (25 - RATE) fun i => state[RATE + i.val]
  let state' : Vector _ 25 := state_rate ++ state_capacity

Vector.mapM in Init.Data.Vector.Basic

def circuit : FormalCircuit (F p) fieldPair field where
  main
  localLength _ := 2

def DecodedInstructionType.val : DecodedInstructionType (F p) → ℕ := fun instrType =>
  if instrType.isAdd = 1 then 0
  else if instrType.isMul = 1 then 1
  else if instrType.isStoreState = 1 then 2
  else 3

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

return { varNames := customNames.toList }

List.zip in Init.Data.List.Basic
</all_available_defs>

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

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

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]

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

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

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

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

@[circuit_norm ↓]
theorem varFromOffset_field {F} (offset : ℕ) :
  varFromOffset (F:=F) field offset = var ⟨offset⟩

lemma main_output_binary_from_completeness (n : ℕ) (offset : ℕ) (env : Environment (F p))
    (input_var : Var (fields n) (F p)) (input : fields n (F p))
    (h_eval : input = eval env input_var)
    (h_assumptions : Assumptions n input)
    (h_local_witnesses : env.UsesLocalWitnessesCompleteness offset ((main input_var).operations offset))
    (h_completeness : Circuit.ConstraintsHold.Completeness env ((main input_var).operations offset)) :
    let output := env ((main input_var).output offset)
    IsBool output

lemma forAll_flatten_foldl :
  Operations.forAll n prop (List.ofFn fun (i : Fin m) => (circuit (foldlAcc n xs circuit init i) xs[i.val]).operations (n + i * constant.localLength)).flatten
    ↔ ∀ (i : Fin m), (circuit (foldlAcc n xs circuit init i) xs[i.val]).forAll (n + i * constant.localLength) prop

theorem two_and_le_add {x y : ℕ} (hx : x < 256) (hy : y < 256) : 2 * (x &&& y) ≤ x + y

theorem MemoryAccessList.isConsistentOnline_filter_of_consistentOnline (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted)
    (h_consistent : MemoryAccessList.isConsistentOnline accesses h_sorted) (addr : ℕ) :
    MemoryAccessList.isConsistentOnline (MemoryAccessList.filterAddress accesses addr) (MemoryAccessList.filterAddress_sorted accesses h_sorted addr)

lemma inductPush_cons_push {motive : {n : ℕ} → Vector α n → Sort u}
  {nil : motive #v[]}
  {push' : ∀ {n : ℕ} (as : Vector α n) (a : α), motive as → motive (as.push a)}
  {n : ℕ} (xs : Vector α n) (x a : α) :
    inductPush nil push' (cons x (xs.push a)) = push' (cons x xs) a (inductPush nil push' (cons x xs))

@[ext]
lemma ext {x y : U32 (F p)}
    (h0 : x.x0 = y.x0)
    (h1 : x.x1 = y.x1)
    (h2 : x.x2 = y.x2)
    (h3 : x.x3 = y.x3)
    : x = y

lemma last_has_incoming_transition {α : Type*} (l : List α) (x : α)
    (h_len : l.length ≥ 2)
    (h_last : l.getLast? = some x) :
    ∃ y, (y, x) ∈ l.zip l.tail

theorem localLength_eq : (xs.mapM circuit).localLength n = m * constant.localLength

@[circuit_norm]
theorem append_localLength {a b: Operations F} :
    (a ++ b).localLength = a.localLength + b.localLength

lemma reachable_leaf_has_incoming_edge (R : Run S) (root leaf : S)
    (h_leaf : R.isLeaf root leaf)
    (h_ne : root ≠ leaf) :
    ∃ y, R (y, leaf) > 0

theorem and_mul_two_pow {x y n : ℕ} : 2 ^ n * (x &&& y) =  2 ^ n * x &&&  2 ^ n * y

lemma countTransitionInPath_append_singleton_other (path : List S) (x y : S) (t : Transition S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_ne : t ≠ (x, y)) :
    countTransitionInPath t (path ++ [y]) = countTransitionInPath t path

theorem Circuit.subcircuitsConsistent_bind {α β : Type} (f : Circuit (F p) α) (g : α → Circuit (F p) β) (offset : ℕ)
    (hf : Operations.SubcircuitsConsistent offset (f.operations offset))
    (hg : Operations.SubcircuitsConsistent (offset + f.localLength offset)
          ((g (f.output offset)).operations (offset + f.localLength offset))) :
    Operations.SubcircuitsConsistent offset ((f >>= g).operations offset)

theorem localLength_eq (n : ℕ) (input : Var (fields n) (F p)) (offset : ℕ) :
    (main input).localLength offset = n - 1

theorem forAll_cons {condition : Condition F} {offset : ℕ} {op : FlatOperation F} {ops : List (FlatOperation F)} :
  forAll offset condition (op :: ops) ↔
    condition.applyFlat offset op ∧ forAll (op.singleLocalLength + offset) condition ops

lemma rotRight64_bytes_u64_eq (o : ℕ) (x : U64 ℕ) :
  rotRight64_bytes x.toLimbs o = (rotRight64_u64 x o).toLimbs

theorem toBits_injective (n : ℕ) {x y : ℕ} : x < 2^n → y < 2^n →
    toBits n x = toBits n y → x = y
</all_available_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>

