-- Clean/Tables/KeccakInductive.lean

import Clean.Table.Inductive

import Clean.Circuit.Extensions

import Clean.Gadgets.Keccak.AbsorbBlock

import Clean.Specs.Keccak256

open Specs.Keccak256

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

namespace Tables.KeccakInductive

open Gadgets.Keccak256

def table : InductiveTable (F p) KeccakState KeccakBlock where
  step state block := do
    KeccakBlock.normalized block
    AbsorbBlock.circuit { state, block }

  Spec _ blocks i _ state : Prop :=
    state.Normalized
    ∧ state.value = absorbBlocks (blocks.map KeccakBlock.value)

  InputAssumptions i block := block.Normalized

  soundness := by admit /- proof elided -/

def initialState : KeccakState (F p) := .fill 25 (U64.fromByte 0)

def formalTable (output : KeccakState (F p)) := table.toFormal initialState output

-- Repository definitions used:
content: syntax "infer_constant_length" : tactic

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

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

content: @[inherit_doc] notation:67 "<+>" => Trace.empty

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

content: macro_rules
  | `(tactic|simp_assign_row) =>
    `(tactic|(
    simp only [assignCurrRow, assignNextRow, size]
    rw [List.finRange, List.ofFn]
    repeat rw [Fin.foldr_succ]
    rw [Fin.foldr_zero]
    repeat rw [List.forM_cons]
    rw [List.forM_nil, bind_pure_unit]
    simp only [seval, toVars, toElements, Fin.cast_eq_self, Fin.val_zero, Fin.val_one, Fin.isValue,
      List.getElem_toArray, List.getElem_cons_zero, List.getElem_cons_succ, Fin.succ_zero_eq_one]))

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

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

content: def F p := ZMod p

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

content: @[reducible] def CAPACITY := 8

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

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

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

content: 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 ⟩

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

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

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

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

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

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

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

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

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

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

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

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

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

content: @[reducible]
def TypeMap := Type → Type

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

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

content: structure Variable (F : Type) where
  index : ℕ

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

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

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

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

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

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

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

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

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

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

content: 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)

   
  Spec : (initialState : State F) → (xs : List (Input F)) → (i : ℕ) → (xs.length = i) → (currentState : State F) → Prop

   
  InputAssumptions : ℕ → Input F → Prop := fun _ _ => True
  InitialStateAssumptions : State F → Prop := fun _ => True

  soundness : InductiveTable.Soundness F State Input Spec step

  completeness : InductiveTable.Completeness F State Input InputAssumptions InitialStateAssumptions Spec step

  subcircuitsConsistent : ∀ acc x, ((step acc x).operations ((size State) + (size Input))).SubcircuitsConsistent ((size State) + (size Input))
    := by admit /- proof elided -/

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

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

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

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

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

content: @[reducible] def Spec (input : Input (F p)) (out_state : KeccakState (F p)) :=
  out_state.Normalized ∧
  out_state.value = absorbBlock input.state.value input.block.value

content: structure Input (F : Type) where
  state : KeccakState F
  block : KeccakBlock F

content: @[reducible] def KeccakBlock := ProvableVector U64 RATE

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

content: @[reducible] def KeccakState := ProvableVector U64 25

content: structure Inputs (F : Type) where
  state : KeccakState F
  d : KeccakRow F

content: def state : Vector (U64 ℕ) 25 := #v[
  ⟨67, 168, 144, 181, 2, 173, 144, 47⟩,
  ⟨114, 52, 107, 105, 171, 22, 114, 75⟩,
  ⟨196, 118, 22, 253, 100, 162, 87, 52⟩,
  ⟨50, 65, 171, 81, 229, 6, 172, 155⟩,
  ⟨178, 167, 68, 225, 82, 73, 216, 194⟩,
  ⟨193, 5, 52, 193, 148, 168, 64, 147⟩,
  ⟨212, 142, 107, 244, 55, 237, 100, 203⟩,
  ⟨101, 34, 195, 62, 133, 216, 64, 34⟩,
  ⟨240, 214, 204, 27, 17, 231, 66, 179⟩,
  ⟨136, 37, 228, 137, 64, 208, 27, 90⟩,
  ⟨177, 229, 130, 4, 191, 7, 25, 117⟩,
  ⟨124, 168, 245, 7, 222, 138, 168, 16⟩,
  ⟨115, 130, 213, 74, 217, 123, 172, 109⟩,
  ⟨128, 149, 137, 6, 45, 133, 77, 101⟩,
  ⟨104, 90, 153, 237, 72, 44, 164, 84⟩,
  ⟨129, 177, 235, 28, 82, 30, 150, 201⟩,
  ⟨52, 55, 32, 241, 142, 211, 246, 68⟩,
  ⟨149, 124, 124, 204, 34, 220, 229, 69⟩,
  ⟨215, 168, 47, 96, 70, 5, 220, 2⟩,
  ⟨53, 224, 38, 18, 110, 66, 70, 9⟩,
  ⟨213, 122, 200, 196, 186, 122, 207, 42⟩,
  ⟨141, 103, 32, 88, 244, 160, 37, 76⟩,
  ⟨99, 242, 138, 24, 4, 30, 100, 196⟩,
  ⟨141, 253, 136, 54, 8, 21, 204, 152⟩,
  ⟨93, 161, 29, 12, 44, 252, 49, 57⟩
]

content: def absorbBlock (state : Vector ℕ 25) (block : Vector ℕ RATE) : Vector ℕ 25 :=
  
  let state' := Vector.mapFinRange 25 fun i => state[i] ^^^ (if _ : i.val < RATE then block[i] else 0)
  
  keccakPermutation state'

content: def keccakPermutation (state : Vector ℕ 25): Vector ℕ 25 :=
  roundConstants.foldl keccakRound state

content: def keccakRound (state : Vector ℕ 25) (rc : UInt64) : Vector ℕ 25 :=
  let theta_state := theta state
  let rho_pi_state := rhoPi theta_state
  let chi_state := chi rho_pi_state
  iota chi_state rc

content: def theta (state : Vector ℕ 25) : Vector ℕ 25 :=
  let c := thetaC state
  let d := thetaD c
  thetaXor state d

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

content: 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]
  ]

content: def thetaD (c : Vector ℕ 5) : Vector ℕ 5 :=
  #v[
    c[4] ^^^ (rotLeft64 c[1] 1),
    c[0] ^^^ (rotLeft64 c[2] 1),
    c[1] ^^^ (rotLeft64 c[3] 1),
    c[2] ^^^ (rotLeft64 c[4] 1),
    c[3] ^^^ (rotLeft64 c[0] 1)
  ]

content: def rotLeft64 (value : ℕ) (left : Fin 64) : ℕ:=
  let right := (64 - left) % 64
  rotRight64 value right

content: def rotRight64 (x : ℕ) (offset : ℕ) : ℕ :=
  let offset := offset % 64
  let low := x % (2^offset)
  let high := x / (2^offset)
  low * (2^(64-offset)) + high

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

content: def iota (state : Vector ℕ 25) (rc : UInt64) : Vector ℕ 25 :=
  state.set 0 ((state[0]) ^^^ rc.toFin)

content: def rhoPi (state : Vector ℕ 25) : Vector ℕ 25 :=
  #v[
    rotLeft64 state[0] 0,
    rotLeft64 state[15] 28,
    rotLeft64 state[5] 1,
    rotLeft64 state[20] 27,
    rotLeft64 state[10] 62,
    rotLeft64 state[6] 44,
    rotLeft64 state[21] 20,
    rotLeft64 state[11] 6,
    rotLeft64 state[1] 36,
    rotLeft64 state[16] 55,
    rotLeft64 state[12] 43,
    rotLeft64 state[2] 3,
    rotLeft64 state[17] 25,
    rotLeft64 state[7] 10,
    rotLeft64 state[22] 39,
    rotLeft64 state[18] 21,
    rotLeft64 state[8] 45,
    rotLeft64 state[23] 8,
    rotLeft64 state[13] 15,
    rotLeft64 state[3] 41,
    rotLeft64 state[24] 14,
    rotLeft64 state[14] 61,
    rotLeft64 state[4] 18,
    rotLeft64 state[19] 56,
    rotLeft64 state[9] 2
  ]

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

content: def not64 (a : ℕ) : ℕ := a ^^^ 0xffffffffffffffff

content: def roundConstants : Vector UInt64 24 := #v[
  0x0000000000000001, 0x0000000000008082,
  0x800000000000808a, 0x8000000080008000,
  0x000000000000808b, 0x0000000080000001,
  0x8000000080008081, 0x8000000000008009,
  0x000000000000008a, 0x0000000000000088,
  0x0000000080008009, 0x000000008000000a,
  0x000000008000808b, 0x800000000000008b,
  0x8000000000008089, 0x8000000000008003,
  0x8000000000008002, 0x8000000000000080,
  0x000000000000800a, 0x800000008000000a,
  0x8000000080008081, 0x8000000000008080,
  0x0000000080000001, 0x8000000080008008
]

content: def mapFinRange (n : ℕ) (create : Fin n → α) : Vector α n := finRange n |>.map create

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

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

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

content: @[reducible] def Assumptions (input : Input (F p)) :=
  input.state.Normalized ∧ input.block.Normalized

content: def InductiveTable.Completeness (F : Type) [Field F] (State Input : Type → Type) [ProvableType State] [ProvableType Input]
    (InputAssumptions : ℕ → Input F → Prop) (InitialStateAssumptions : State F → Prop)
    (Spec : (initialState : State F) → (xs : List (Input F)) → (i : ℕ) → (xs.length = i) → (currentState : State F) → Prop)
    (step : Var State F → Var Input F → Circuit F (Var State F)) :=
  ∀ (initialState : State F) (row_index : ℕ) (env : Environment F),
  
  ∀ (acc_var : Var State F) (x_var : Var Input F)
    (acc : State F) (x : Input F) (xs : List (Input F)) (xs_len : xs.length = row_index),
    (eval env acc_var = acc) ∧ (eval env x_var = x) →
  
  env.UsesLocalWitnessesCompleteness ((size State) + (size Input)) (step acc_var x_var |>.operations ((size State) + (size Input))) →
  
  InitialStateAssumptions initialState ∧
  Spec initialState xs row_index xs_len acc ∧ InputAssumptions row_index x →
  
  Circuit.ConstraintsHold.Completeness env (step acc_var x_var |>.operations ((size State) + (size Input)))

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

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

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

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

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

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

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

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

content: 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 _ _ := ()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: @[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]

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

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

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

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

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

content: instance elaborated : ElaboratedCircuit (F p) Input KeccakState where
  main
  localLength _ := 31048
  output _ i0 := Permutation.stateVar (i0 + 136) 23

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

content: def main (input : Var Input (F p)) : Circuit (F p) (Var KeccakState (F p)) := do
  let { state, block } := input
  
  let state_rate ← Circuit.mapFinRange RATE fun i => Xor64.circuit ⟨state[i.val], block[i.val]⟩

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

  
  Permutation.circuit state'

content: def circuit : FormalCircuit (F p) KeccakState KeccakState := {
  elaborated with
  Assumptions, Spec, soundness
  
  
  completeness := by admit /- proof elided -/

content: def Spec (state : KeccakState (F p)) (out_state : KeccakState (F p)) :=
  out_state.Normalized
  ∧ out_state.value = keccakPermutation state.value

content: instance elaborated : ElaboratedCircuit (F p) KeccakState KeccakState where
  main
  localLength _ := 30912
  output _ i0 := stateVar i0 23

  localLength_eq state i0 := by admit /- proof elided -/

content: def main (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakState (F p)) :=
  .foldl roundConstants state
    fun state rc => KeccakRound.circuit rc state

content: def circuit (rc : UInt64) : FormalCircuit (F p) KeccakState KeccakState := {
  elaborated rc with
  Spec := Spec rc
  Assumptions
  soundness := soundness rc,
  completeness := completeness rc,
}

content: def Spec (rc : UInt64) (state : KeccakState (F p)) (out_state : KeccakState (F p)) :=
  out_state.Normalized
  ∧ out_state.value = keccakRound state.value rc

content: def Assumptions (state : KeccakState (F p)) := state.Normalized

content: instance elaborated (rc : UInt64) : ElaboratedCircuit (F p) KeccakState KeccakState where
  main := main rc
  localLength _ := 1288
  output _ i0 := (Vector.mapRange 25 fun i => varFromOffset U64 (i0 + i*16 + 888) ).set 0 (varFromOffset U64 (i0 + 1280))

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

content: def main (rc : UInt64) (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakState (F p)) := do
  let state ← Theta.circuit state
  let state ← RhoPi.circuit state
  let state ← Chi.circuit state

  
  let s0 ← Xor64.circuit ⟨state[0], const (U64.fromUInt64 rc)⟩
  return state.set 0 s0

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

content: structure Inputs (F : Type) where
  x: U64 F
  y: U64 F

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

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

content: def fromUInt64 (x : UInt64) : U64 (F p) :=
  decomposeNat x.toFin

content: def decomposeNat (x : ℕ) : U64 (F p) :=
  let x0 := x % 256
  let x1 : ℕ := (x / 256) % 256
  let x2 : ℕ := (x / 256^2) % 256
  let x3 : ℕ := (x / 256^3) % 256
  let x4 : ℕ := (x / 256^4) % 256
  let x5 : ℕ := (x / 256^5) % 256
  let x6 : ℕ := (x / 256^6) % 256
  let x7 : ℕ := (x / 256^7) % 256
  ⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩

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

content: def Assumptions := KeccakState.Normalized (p:=p)

content: instance elaborated : ElaboratedCircuit (F p) KeccakState KeccakState where
  main
  localLength _ := 400
  localLength_eq _ _ := by admit /- proof elided -/

content: def main (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakState (F p)) :=
  .map rhoPiConstants fun (i, s) =>
    Rotation64.circuit (-s) state[i.val]

content: def rhoPiConstants := rhoPiIndices.zip rhoPiShifts

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

content: def rhoPiShifts : Vector (Fin 64) 25 := #v[
  0, 28, 1, 27, 62, 44, 20, 6, 36, 55, 43, 3, 25, 10, 39, 21, 45, 8, 15, 41, 14, 61, 18, 56, 2
]

content: def circuit (offset : Fin 64) : FormalCircuit (F p) U64 U64 := {
  elaborated offset with
  Assumptions
  Spec := Spec offset
  soundness := soundness offset
  completeness := completeness offset
}

content: def Spec (offset : Fin 64) (x : U64 (F p)) (y : U64 (F p)) :=
  y.value = rotRight64 x.value offset.val
  ∧ y.Normalized

content: def elaborated (off : Fin 64) : ElaboratedCircuit (F p) U64 U64 where
  main := main off
  localLength _ := 16
  output _ i0 := output off i0

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

  
  let byte_rotated ← Rotation64Bytes.circuit byte_offset x
  Rotation64Bits.circuit bit_offset byte_rotated

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

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

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

content: def fromLimbs {F} (v : Vector F 8) : U64 F := fromElements v

content: def main (offset : Fin 8) (x : Var U64 (F p)) : Circuit (F p) (Var U64 (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 U64.fromLimbs rotated

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

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

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

content: 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))

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

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

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

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

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

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

content: @[reducible] def field : TypeMap := id

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

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

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

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

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

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

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

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

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

content: def Spec (offset : Fin 8) (x : U64 (F p)) (y : U64 (F p)) :=
  y.value = rotRight64 x.value offset.val
  ∧ y.Normalized

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

content: def main (offset : Fin 8) (input : Var U64 (F p)) : Circuit (F p) (Var U64 (F p)) := do
  let ⟨x0, x1, x2, x3 , x4, x5, x6, x7⟩ := input

  if offset = 0 then
    return ⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩
  else if offset = 1 then
    return ⟨ x1, x2, x3, x4, x5, x6, x7, x0 ⟩
  else if offset = 2 then
    return ⟨ x2, x3, x4, x5, x6, x7, x0, x1 ⟩
  else if offset = 3 then
    return ⟨ x3, x4, x5, x6, x7, x0, x1, x2 ⟩
  else if offset = 4 then
    return ⟨ x4, x5, x6, x7, x0, x1, x2, x3 ⟩
  else if offset = 5 then
    return ⟨ x5, x6, x7, x0, x1, x2, x3, x4 ⟩
  else if offset = 6 then
    return ⟨ x6, x7, x0, x1, x2, x3, x4, x5 ⟩
  else
    return ⟨ x7, x0, x1, x2, x3, x4, x5, x6 ⟩

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

content: def Spec (offset : Fin 8) (x : U64 (F p)) (y : U64 (F p)) :=
  y.value = rotRight64 x.value (offset.val * 8) ∧ y.Normalized

content: def output (offset : Fin 64) (i0 : ℕ) : U64 (Expression (F p)) :=
  Rotation64Bits.output ⟨ offset.val % 8, by admit /- proof elided -/
  ⟩ i0

content: def Spec (state : KeccakState (F p)) (out_state : KeccakState (F p)) :=
  out_state.Normalized
  ∧ out_state.value = Specs.Keccak256.rhoPi state.value

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

content: def Spec (state : KeccakState (F p)) (out_state : KeccakState (F p)) : Prop :=
  out_state.Normalized
  ∧ out_state.value = Specs.Keccak256.theta state.value

content: instance elaborated : ElaboratedCircuit (F p) KeccakState KeccakState where
  main
  localLength _ := 480

content: def main (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakState (F p)) := do
  let c ← ThetaC.circuit state
  let d ← ThetaD.circuit c
  ThetaXor.circuit ⟨state, d⟩

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

content: instance elaborated : ElaboratedCircuit (F p) KeccakRow KeccakRow where
  main
  localLength _ := 120

content: 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 c1 ← Rotation64.circuit (64 - 1) row[2]
  let c1 ← Xor64.circuit ⟨row[0], c1⟩

  let c2 ← Rotation64.circuit (64 - 1) row[3]
  let c2 ← Xor64.circuit ⟨row[1], c2⟩

  let c3 ← Rotation64.circuit (64 - 1) row[4]
  let c3 ← Xor64.circuit ⟨row[2], c3⟩

  let c4 ← Rotation64.circuit (64 - 1) row[0]
  let c4 ← Xor64.circuit ⟨row[3], c4⟩

  return #v[c0, c1, c2, c3, c4]

content: @[reducible] def KeccakRow := ProvableVector U64 5

content: def Spec (row : KeccakRow (F p)) (out : KeccakRow (F p)) : Prop :=
  out.Normalized
  ∧ out.value = Specs.Keccak256.thetaD row.value

content: def Assumptions (state : KeccakRow (F p)) := state.Normalized

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

content: instance elaborated : ElaboratedCircuit (F p) Inputs KeccakState where
  main
  localLength _ := 200

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

content: def main : Var Inputs (F p) → Circuit (F p) (Var KeccakState (F p))
  | { state, d } => .mapFinRange 25 fun i =>
    Xor64.circuit ⟨state[i.val], d[i.val / 5]⟩

content: def mapFinRange (m : ℕ) [NeZero m] (body : Fin m → Circuit F β)
    (_constant : ConstantLength body := by admit /- proof elided -/
    ) : Circuit F (Vector β m) :=
  Vector.mapFinRangeM m body

content: def mapFinRangeM (n : ℕ) {m : Type → Type} [Monad m] (f : Fin n → m β) : m (Vector β n) := (finRange n).mapM f

content: def Spec (inputs : Inputs (F p)) (out : KeccakState (F p)) : Prop :=
  let ⟨state, d⟩ := inputs
  out.Normalized
  ∧ out.value = Specs.Keccak256.thetaXor state.value d.value

content: def Assumptions (inputs : Inputs (F p)) : Prop :=
  let ⟨state, d⟩ := inputs
  state.Normalized ∧ d.Normalized

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

content: instance elaborated : ElaboratedCircuit (F p) KeccakState KeccakRow where
  main
  localLength _ := 160
  localLength_eq _ _ := by admit /- proof elided -/

content: def main (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakRow (F p)) :=
  .mapFinRange 5 fun i => do
    let c ← Xor64.circuit ⟨state[5*i.val], state[5*i.val + 1]⟩
    let c ← Xor64.circuit ⟨c, state[5*i.val + 2]⟩
    let c ← Xor64.circuit ⟨c, state[5*i.val + 3]⟩
    let c ← Xor64.circuit ⟨c, state[5*i.val + 4]⟩
    return c

content: def Spec (state : KeccakState (F p)) (out : KeccakRow (F p)) :=
  out.Normalized
  ∧ out.value = Specs.Keccak256.thetaC state.value

content: instance elaborated : ElaboratedCircuit (F p) KeccakState KeccakState where
  main
  localLength _ := 400
  output _ i0 := Vector.mapRange 25 fun i => varFromOffset U64 (i0 + i*16 + 8)

  localLength_eq state i0 := by admit /- proof elided -/

content: def main (state : Var KeccakState (F p)) : Circuit (F p) (Var KeccakState (F p)) :=
  .mapFinRange 25 fun i => do
    let state_not ← Not.circuit (state[i + 5])
    let state_and ← And.And64.circuit ⟨state_not, state[i + 10]⟩
    Xor64.circuit ⟨state[i], state_and⟩

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

content: def circuit : FormalCircuit (F p) U64 U64 where
  main x := pure (not64_bytewise x)
  Assumptions x := x.Normalized
  Spec x z := z.value = not64 x.value ∧ z.Normalized

  localLength _ := 0
  output x _ := not64_bytewise x

  soundness := by admit /- proof elided -/

content: def not64_bytewise (x : Var U64 (F p)) : Var U64 (F p) := U64.map x (fun x => 255 - x)

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

content: def Spec (state : KeccakState (F p)) (out_state : KeccakState (F p)) :=
  out_state.Normalized
  ∧ out_state.value = Specs.Keccak256.chi state.value

content: def stateVar (n : ℕ) (i : ℕ) : Var KeccakState (F p) :=
  Vector.mapRange 25 (fun j => varFromOffset U64 (n + i * 1288 + j * 16 + 888))
  |>.set 0 (varFromOffset U64 (n + i * 1288 + 1280))

content: instance ExplicitCircuits.from_pure {f : α → β} : ExplicitCircuits (fun a => pure (f a) : α → Circuit F β) where
  output a _ := f a
  localLength _ _ := 0
  operations _ _ := []

content: def KeccakBlock.normalized : FormalAssertion (F p) KeccakBlock where
  main block := .forEach block (assertion U64.AssertNormalized.circuit)
  Assumptions _ := True
  Spec block := block.Normalized
  localLength_eq _ _ := by admit /- proof elided -/

content: def circuit : FormalAssertion (F p) U64 where
  main

  Assumptions _ := True
  Spec inputs := inputs.Normalized

  soundness := by admit /- proof elided -/

content: def main (inputs : Var U64 (F p)) : Circuit (F p) Unit  := do
  let ⟨x0, x1, x2, x3, x4, x5, x6, x7⟩ := inputs
  lookup ByteTable x0
  lookup ByteTable x1
  lookup ByteTable x2
  lookup ByteTable x3
  lookup ByteTable x4
  lookup ByteTable x5
  lookup ByteTable x6
  lookup ByteTable x7

content: @[circuit_norm]
def assertion (circuit : FormalAssertion F β) (b : Var β F) : Circuit F Unit :=
  fun offset =>
    let subcircuit := circuit.toSubcircuit offset b
    ((), [.subcircuit subcircuit])

content: def absorbBlocks (blocks : List (Vector ℕ RATE)) : Vector ℕ 25 :=
  blocks.foldl absorbBlock initialState

content: def initialState : Vector ℕ 25 := .fill 25 0

content: def foldl {m : ℕ} [Inhabited β] [Inhabited α] (xs : Vector α m) (init : β) (body : β → α → Circuit F β)
  (_const_out : ConstantOutput (fun (s, a) => body s a) := by admit /- proof elided -/
    )
  (_constant : ConstantLength (fun (s, a) => body s a) := by admit /- proof elided -/
  )
   : Circuit F β :=
  xs.foldlM body init

content: @[circuit_norm]
def ConstantOutput (circuit : α → Circuit F β) [Inhabited α] :=
  ∀ (x : α) (n : ℕ), (circuit x).output n = (circuit default).output n

content: def InductiveTable.Soundness (F : Type) [Field F] (State Input : Type → Type) [ProvableType State] [ProvableType Input]
    (Spec : (initialState : State F) → (xs : List (Input F)) → (i : ℕ) → (xs.length = i) → (currentState : State F) → Prop)
    (step : Var State F → Var Input F → Circuit F (Var State F)) :=
  ∀ (initialState : State F) (row_index : ℕ) (env : Environment F),
  
  ∀ (acc_var : Var State F) (x_var : Var Input F)
    (acc : State F) (x : Input F) (xs : List (Input F)) (xs_len : xs.length = row_index),
      (eval env acc_var = acc) ∧ (eval env x_var = x) →
    
    Circuit.ConstraintsHold.Soundness env (step acc_var x_var |>.operations ((size State) + (size Input))) →
    
    Spec initialState xs row_index xs_len acc →
    
    Spec initialState (xs.concat x) (row_index + 1) (xs_len ▸ List.length_concat) (eval env (step acc_var x_var |>.output ((size State) + (size Input))))

content: def fromByte (x : Fin 256) : U64 (F p) :=
  ⟨ x.val, 0, 0, 0, 0, 0, 0, 0 ⟩

content: def traceInputs {N : ℕ} (trace : TraceOfLength F (ProvablePair State Input) N) : List (Input F) :=
  trace.val.toList.map Prod.snd

content: def toList : Operations F → List (Operation F) := id

content: def toList : Trace F S → List (Row F S)
  | <+> => []
  | rest +> row => rest.toList.concat row

content: inductive Trace (F : Type) (S : Type → Type) [ProvableType S] where
   
  | empty : Trace F S
   
  | cons (rest : Trace F S) (row : Row F S) : Trace F S

content: structure CellOffset (W : ℕ+) (S : Type → Type) [ProvableType S]  where
  row: Fin W
  column: Fin (size S)

content: @[reducible, table_norm, table_assignment_norm]
def empty : TableContext W S F where
  circuit := []
  assignment := .empty W

content: structure TableContext (W : ℕ+) (S : Type → Type) (F : Type) [Field F] [ProvableType S] where
  circuit : Operations F
  assignment : CellAssignment W S
deriving Repr

content: structure CellAssignment (W : ℕ+) (S : Type → Type) [ProvableType S] where
  offset : ℕ 
  aux_length : ℕ 

   
  vars : Vector (Cell W S) offset

content: @[reducible, table_norm, table_assignment_norm]
def offset (table : TableContext W S F) : ℕ := table.circuit.localLength

content: inductive Cell (W : ℕ+) (S : Type → Type) [ProvableType S] where
  | input : CellOffset W S → Cell W S
  | aux : ℕ → Cell W S

content: @[reducible]
def Row (F : Type) (S : Type → Type) [ProvableType S] := S F

content: instance ProvablePair.instance {α β: TypeMap} [ProvableType α] [ProvableType β] : ProvableType (ProvablePair α β) where
  size := size α + size β
  toElements := fun (a, b) => toElements a ++ toElements b
  fromElements {F} v :=
    let a : α F := v.take (size α) |>.cast Nat.min_add_right_self |> fromElements
    let b : β F := v.drop (size α) |>.cast (Nat.add_sub_self_left _ _) |> fromElements
    (a, b)
  fromElements_toElements x := by admit /- proof elided -/

content: def TraceOfLength (F : Type) (S : Type → Type) [ProvableType S] (N : ℕ) : Type :=
  { env : Trace F S // env.len = N }

content: def toList {N : ℕ} (trace : TraceOfLength F S N) : List.Vector (Row F S) N :=
  ⟨ trace.val.toList, by admit /- proof elided -/
  ⟩

content: def FormalTable.statement (table : FormalTable F S) (N : ℕ) (trace : TraceOfLength F S N) : Prop :=
  table.Assumption N → table.Spec trace

content: structure FormalTable (F : Type) [Field F] (S : Type → Type) [ProvableType S] where
   
  constraints : List (TableOperation S F)

   
  Assumption : ℕ → Prop := fun _ => True

   
  Spec {N : ℕ} : TraceOfLength F S N → Prop

   
  soundness :
    ∀ (N : ℕ) (trace : TraceOfLength F S N) (env : ℕ → ℕ → Environment F),
    Assumption N →
    TableConstraintsHold constraints trace env →
    Spec trace

   
  offset_consistent :
    constraints.Forall fun cs =>
      match cs with
      | .boundary _ constraint => constraint.OffsetConsistent
      | .everyRow constraint => constraint.OffsetConsistent
      | .everyRowExceptLast constraint => constraint.OffsetConsistent
    := by admit /- proof elided -/

content: def tail : Trace F S → Trace F S
  | <+>  => <+>
  | rest +> _ => rest

content: def tail {M : ℕ} (trace : TraceOfLength F S M) : TraceOfLength F S (M - 1) :=
  ⟨ trace.val.tail, by admit /- proof elided -/
  ⟩

content: def toFormal (table : InductiveTable F State Input) (input output : State F) : FormalTable F (ProvablePair State Input) where
  constraints := table.tableConstraints input output
  Assumption N := N > 0 ∧ table.Spec input [] 0 rfl input
  Spec {N} trace := table.Spec input (traceInputs trace.tail) (N-1) (traceInputs_length trace.tail) output

  soundness N trace env assumption constraints :=
    table.table_soundness input output ⟨N, assumption.left⟩ trace env assumption.right constraints

  offset_consistent := by admit /- proof elided -/

content: def equalityConstraint (Input : TypeMap) [ProvableType Input] (target : State F) : SingleRowConstraint (ProvablePair State Input) F := do
  let (actual, _) ← getCurrRow
  actual === (const target)

content: @[table_norm, table_assignment_norm]
def getCurrRow : TableConstraint W S F (Var S F) := getRow 0

content: @[reducible, table_norm, table_assignment_norm]
def TableConstraint (W : ℕ+) (S : Type → Type) (F : Type) [Field F] [ProvableType S] :=
  StateM (TableContext W S F)

content: @[table_norm, table_assignment_norm]
def getRow (row : Fin W) : TableConstraint W S F (Var S F) :=
  modifyGet fun ctx =>
    let ctx' : TableContext W S F := {
      circuit := ctx.circuit ++ [.witness (size S) fun env => .mapRange _ fun i => env.get (ctx.offset + i)],
      assignment := ctx.assignment.pushRow row
    }
    (varFromOffset S ctx.offset, ctx')

content: @[table_norm, table_assignment_norm]
def Row.get (row : Row F S) (i : Fin (size S)) : F :=
  (toElements row)[i.val]

content: @[table_assignment_norm]
def pushRow (assignment : CellAssignment W S) (row : Fin W) : CellAssignment W S :=
  let row_vars : Vector (Cell W S) (size S) := .mapFinRange _ fun col => .input ⟨ row, col ⟩
  {
    offset := assignment.offset + size S
    aux_length := assignment.aux_length
    vars := assignment.vars ++ row_vars
  }

content: @[table_assignment_norm]
def get {M : ℕ} :
    (env : TraceOfLength F S M) → (i : Fin M) → (j : Fin (size S)) → F
  | ⟨env, h⟩, i, j => env.getLeFromBottom ⟨
      M - 1 - i,
      by admit /- proof elided -/
    ⟩ j

content: @[reducible]
def SingleRowConstraint (S : Type → Type) (F : Type) [Field F] [ProvableType S] := TableConstraint 1 S F Unit

content: def inductiveConstraint (table : InductiveTable F State Input) : TableConstraint 2 (ProvablePair State Input) F Unit := do
  let (acc, x) ← getCurrRow
  let output ← table.step acc x
  let (output', _) ← getNextRow
  
  output' === output

content: @[table_norm, table_assignment_norm]
def getNextRow : TableConstraint W S F (Var S F) := getRow 1

content: def tableConstraints (table : InductiveTable F State Input) (input_state output_state : State F) :
  List (TableOperation (ProvablePair State Input) F) := [
    .everyRowExceptLast table.inductiveConstraint,
    .boundary (.fromStart 0) (equalityConstraint Input input_state),
    .boundary (.fromEnd 0) (equalityConstraint Input output_state),
  ]

content: inductive TableOperation (S : Type → Type) (F : Type) [Field F] [ProvableType S] where
   
  | boundary: RowIndex → SingleRowConstraint S F → TableOperation S F

   
  | everyRow: SingleRowConstraint S F → TableOperation S F

   
  | everyRowExceptLast: TwoRowsConstraint S F → TableOperation S F

content: inductive RowIndex where
  | fromStart : ℕ → RowIndex
  | fromEnd : ℕ → RowIndex

content: @[reducible]
def TwoRowsConstraint (S : Type → Type) (F : Type) [Field F] [ProvableType S] := TableConstraint 2 S F Unit

content: @[table_norm]
def TableConstraintsHold {N : ℕ} (constraints : List (TableOperation S F))
  (trace : TraceOfLength F S N) (env : ℕ → ℕ → Environment F) : Prop :=
  let constraints_and_envs := constraints.mapIdx (fun i cs => (cs, env i))
  foldl N constraints_and_envs trace.val constraints_and_envs
  where

content: @[table_assignment_norm]
def OffsetConsistent (table : TableConstraint W S F α) : Prop :=
  table.finalOffset = table.finalAssignment.offset

content: @[reducible, table_norm, table_assignment_norm]
def finalOffset (table : TableConstraint W S F α) : ℕ :=
  table .empty |>.snd.circuit.localLength

content: @[table_assignment_norm]
def finalAssignment (table : TableConstraint W S F α) : CellAssignment W S :=
  table .empty |>.snd.assignment

content: def fill (n : ℕ) (a : α) : Vector α n :=
  match n with
  | 0 => #v[]
  | k + 1 => (fill k a).push a

content: @[table_norm, table_assignment_norm]
def len : Trace F S → ℕ
  | <+> => 0
  | rest +> _ => rest.len + 1

content: infix:50 " === " => HasAssertEq.assert_eq

content: @[inherit_doc] notation:67 "<+>" => Trace.empty

content: @[inherit_doc] infixl:67 " +> " => Trace.cons

-- Repository lemmas used:
content: lemma fromByte_value {x : Fin 256} : (fromByte x).value (p:=p) = x

content: theorem getElem_fill {n} {a : α} {i : ℕ} {hi : i < n} :
    (fill n a)[i] = a

content: lemma fromByte_normalized {x : Fin 256} : (fromByte x).Normalized (p:=p)

content: lemma toList_length : (trace : Trace F S) → trace.toList.length = trace.len
  | <+> => rfl

-- Local definitions used:
content: def table : InductiveTable (F p) KeccakState KeccakBlock where
  step state block := do
    KeccakBlock.normalized block
    AbsorbBlock.circuit { state, block }

  Spec _ blocks i _ state : Prop :=
    state.Normalized
    ∧ state.value = absorbBlocks (blocks.map KeccakBlock.value)

  InputAssumptions i block := block.Normalized

  soundness := by admit /- proof elided -/

content: def initialState : KeccakState (F p) := .fill 25 (U64.fromByte 0)

content: def formalTable (output : KeccakState (F p)) := table.toFormal initialState output

-- Local lemmas used:
content: lemma initialState_value : (initialState (p:=p)).value = .fill 25 0

content: lemma initialState_normalized : (initialState (p:=p)).Normalized

-- Target theorem and proof:
theorem tableStatement (output : KeccakState (F p)) : ∀ n > 0, ∀ trace, ∃ blocks, blocks.length = n - 1 ∧
  (formalTable output).statement n trace →
    output.Normalized ∧ output.value = absorbBlocks blocks
:= by
  intro n hn trace
  use (InductiveTable.traceInputs trace.tail).map KeccakBlock.value
  intro Spec
  simp only [formalTable, FormalTable.statement, table, InductiveTable.toFormal] at Spec
  simp only [List.length_map, Trace.toList_length, trace.tail.prop, InductiveTable.traceInputs, hn] at Spec
  simp only [initialState_value, initialState_normalized, absorbBlocks, Specs.Keccak256.initialState, true_and] at Spec
  exact Spec rfl