-- Clean/Gadgets/Xor/Xor64.lean

import Mathlib.Algebra.Field.Basic

import Mathlib.Data.ZMod.Basic

import Clean.Utils.Primes

import Clean.Utils.Vector

import Clean.Circuit.Expression

import Clean.Circuit.Provable

import Clean.Circuit.Basic

import Clean.Utils.Field

import Clean.Types.U64

import Clean.Gadgets.Xor.ByteXorTable

section

variable {p : ℕ} [Fact p.Prime] [p_large_enough: Fact (p > 512)]

namespace Gadgets.Xor64

open Gadgets.Xor

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

def main (input : Var Inputs (F p)) : Circuit (F p) (Var U64 (F p))  := do
  let ⟨x, y⟩ := input
  let z ← witness fun env =>
    let z0 := (env x.x0).val ^^^ (env y.x0).val
    let z1 := (env x.x1).val ^^^ (env y.x1).val
    let z2 := (env x.x2).val ^^^ (env y.x2).val
    let z3 := (env x.x3).val ^^^ (env y.x3).val
    let z4 := (env x.x4).val ^^^ (env y.x4).val
    let z5 := (env x.x5).val ^^^ (env y.x5).val
    let z6 := (env x.x6).val ^^^ (env y.x6).val
    let z7 := (env x.x7).val ^^^ (env y.x7).val
    U64.mk z0 z1 z2 z3 z4 z5 z6 z7

  lookup ByteXorTable (x.x0, y.x0, z.x0)
  lookup ByteXorTable (x.x1, y.x1, z.x1)
  lookup ByteXorTable (x.x2, y.x2, z.x2)
  lookup ByteXorTable (x.x3, y.x3, z.x3)
  lookup ByteXorTable (x.x4, y.x4, z.x4)
  lookup ByteXorTable (x.x5, y.x5, z.x5)
  lookup ByteXorTable (x.x6, y.x6, z.x6)
  lookup ByteXorTable (x.x7, y.x7, z.x7)
  return z

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

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

instance elaborated : ElaboratedCircuit (F p) Inputs U64 where
  main := main
  localLength _ := 8
  output _ i0 := varFromOffset U64 i0

-- Repository definitions used:
content: syntax "let " ident " <== " term : doElem

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

syntax "infer_constant_length" : tactic

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|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: structure U64 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
  x4 : T
  x5 : T
  x6 : T
  x7 : T
deriving DecidableEq

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

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: 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: 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: @[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 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: def F p := ZMod p

content: def ByteXorTable : Table (F p) fieldTriple := .fromStatic {
  name := "ByteXor"
  length := 256*256

  row i :=
    let (x, y) := splitTwoBytes i
    (fromByte x, fromByte y, fromByte (x ^^^ y))

  index := fun (x, y, _) => x.val * 256 + y.val

  Spec := fun (x, y, z) =>
    x.val < 256 ∧ y.val < 256 ∧ z.val = x.val ^^^ y.val

  contains_iff := by admit /- proof elided -/

content: def splitTwoBytes (i : Fin (256 * 256)) : Fin 256 × Fin 256 :=
  let x := i.val / 256
  let y := i.val % 256
  have x_lt : x < 256 := by admit /- proof elided -/

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: def fromByte (x : Fin 256) : F p :=
  FieldUtils.natToField x.val (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: @[reducible]
def fieldTriple : TypeMap := fun F => F × F × F

content: def concatTwoBytes (x y : Fin 256) : Fin (256 * 256) :=
  let i := x.val * 256 + y.val
  have i_lt : i < 256 * 256 := by admit /- proof elided -/

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

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

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: @[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 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: 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: @[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 operations (circuit : Circuit F α) (offset : ℕ) : Operations F :=
  (circuit offset).2

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

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

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 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: @[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 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]
def unit (_ : Type) := Unit

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: @[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: 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: @[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: @[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 map {m : ℕ} (xs : Vector α m) (body : α → Circuit F β)
    (_constant : ConstantLength body := by admit /- proof elided -/
    ) : Circuit F (Vector β m) :=
  xs.mapM body

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: 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 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 KeccakRow.Normalized (row : KeccakRow (F p)) :=
  ∀ i : Fin 5, row[i.val].Normalized

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

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

content: def BLAKE3State.value (state : BLAKE3State (F p)) := state.map U32.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: structure U32 (T : Type) where
  x0 : T
  x1 : T
  x2 : T
  x3 : T
deriving DecidableEq

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

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

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

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

content: @[reducible] def CAPACITY := 8

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

-- Repository lemmas used:
content: omit [Fact (Nat.Prime p)] p_large_enough in
theorem value_xor_horner {x : U64 (F p)} (hx : x.Normalized) : x.value =
    x.x0.val ^^^ 2^8 * (x.x1.val ^^^ 2^8 * (x.x2.val ^^^ 2^8 * (x.x3.val ^^^
      2^8 * (x.x4.val ^^^ 2^8 * (x.x5.val ^^^ 2^8 * (x.x6.val ^^^ 2^8 * x.x7.val))))))

content: omit [Fact (Nat.Prime p)] p_large_enough in
theorem value_horner (x : U64 (F p)) : x.value =
    x.x0.val + 2^8 * (x.x1.val + 2^8 * (x.x2.val + 2^8 * (x.x3.val +
      2^8 * (x.x4.val + 2^8 * (x.x5.val + 2^8 * (x.x6.val + 2^8 * x.x7.val))))))

content: theorem xor_mul_two_pow {x y n : ℕ} : 2 ^ n * (x ^^^ y) =  2 ^ n * x ^^^  2 ^ n * y

-- Local definitions used:
content: structure Inputs (F : Type) where
  x: U64 F
  y: U64 F

content: def main (input : Var Inputs (F p)) : Circuit (F p) (Var U64 (F p))  := do
  let ⟨x, y⟩ := input
  let z ← witness fun env =>
    let z0 := (env x.x0).val ^^^ (env y.x0).val
    let z1 := (env x.x1).val ^^^ (env y.x1).val
    let z2 := (env x.x2).val ^^^ (env y.x2).val
    let z3 := (env x.x3).val ^^^ (env y.x3).val
    let z4 := (env x.x4).val ^^^ (env y.x4).val
    let z5 := (env x.x5).val ^^^ (env y.x5).val
    let z6 := (env x.x6).val ^^^ (env y.x6).val
    let z7 := (env x.x7).val ^^^ (env y.x7).val
    U64.mk z0 z1 z2 z3 z4 z5 z6 z7

  lookup ByteXorTable (x.x0, y.x0, z.x0)
  lookup ByteXorTable (x.x1, y.x1, z.x1)
  lookup ByteXorTable (x.x2, y.x2, z.x2)
  lookup ByteXorTable (x.x3, y.x3, z.x3)
  lookup ByteXorTable (x.x4, y.x4, z.x4)
  lookup ByteXorTable (x.x5, y.x5, z.x5)
  lookup ByteXorTable (x.x6, y.x6, z.x6)
  lookup ByteXorTable (x.x7, y.x7, z.x7)
  return z

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

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

content: instance elaborated : ElaboratedCircuit (F p) Inputs U64 where
  main := main
  localLength _ := 8
  output _ i0 := varFromOffset U64 i0

-- Local lemmas used:
content: omit [Fact (Nat.Prime p)] p_large_enough in
theorem soundness_to_u64 {x y z : U64 (F p)}
  (x_norm : x.Normalized) (y_norm : y.Normalized)
  (h_eq :
    z.x0.val = x.x0.val ^^^ y.x0.val ∧
    z.x1.val = x.x1.val ^^^ y.x1.val ∧
    z.x2.val = x.x2.val ^^^ y.x2.val ∧
    z.x3.val = x.x3.val ^^^ y.x3.val ∧
    z.x4.val = x.x4.val ^^^ y.x4.val ∧
    z.x5.val = x.x5.val ^^^ y.x5.val ∧
    z.x6.val = x.x6.val ^^^ y.x6.val ∧
    z.x7.val = x.x7.val ^^^ y.x7.val) : Spec { x, y } z

-- Target theorem and proof:
theorem soundness : Soundness (F p) elaborated Assumptions Spec
:= by
  intro i0 env input_var input h_input h_as h_holds

  let ⟨⟨ x0_var, x1_var, x2_var, x3_var, x4_var, x5_var, x6_var, x7_var ⟩,
       ⟨ y0_var, y1_var, y2_var, y3_var, y4_var, y5_var, y6_var, y7_var ⟩⟩ := input_var
  let ⟨⟨ x0, x1, x2, x3, x4, x5, x6, x7 ⟩,
       ⟨ y0, y1, y2, y3, y4, y5, y6, y7 ⟩⟩ := input

  simp only [circuit_norm, explicit_provable_type, Inputs.mk.injEq, U64.mk.injEq] at h_input

  simp only [circuit_norm, Assumptions] at h_as
  obtain ⟨ x_norm, y_norm ⟩ := h_as

  simp only [h_input, circuit_norm, main, ByteXorTable,
    varFromOffset, Vector.mapRange] at h_holds

  apply soundness_to_u64 x_norm y_norm
  simp only [circuit_norm, explicit_provable_type]
  simp [h_holds]