import Clean.Circuit.Theorems
import Clean.Circuit.Provable
import Clean.Utils.Field
import Clean.Circuit
import Mathlib.Data.Nat.Bitwise
import Clean.Gadgets.Boolean
import Clean.Utils.Bitwise
import Clean.Utils.BinaryOps
import Clean.Circuit.Basic
import Clean.Utils.Vector

syntax "infer_constant_length" : tactic

syntax "infer_constant_length" : tactic

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

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

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

macro_rules
  | `(tactic|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 -/

F

def F p := ZMod p

fields

@[reducible]
def fields (n : ℕ) := fun F => Vector F n

IsBool

def IsBool {α : Type*} [Zero α] [One α] (x : α) : Prop := x = 0 ∨ x = 1

map

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

Circuit

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

Operation

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

Condition

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

FlatOperation

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

Subcircuit

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

localWitnesses

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

lookup

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

Table.toRaw

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

RawTable

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

StaticTable

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

Contains

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

ProvableType.fromStruct

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

TypeMap

@[reducible]
def TypeMap := Type → Type

Table

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

Expression

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

Variable

structure Variable (F : Type) where
  index : ℕ

const

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

Var

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

fromVars

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

Lookup

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

Witnessable

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

Environment

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

localLength

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

ConstraintsHoldFlat

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

ConstantLength

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

ConstantLength.fromConstantLength

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

HasAssignEq

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

FormalCircuit

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

ElaboratedCircuit

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

GeneralFormalCircuit

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

Soundness

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

operations

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

Operations

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

output

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

ConstraintsHold.Soundness

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

GeneralFormalCircuit.Soundness

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

eval

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

toVars

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

eval

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

FormalAssertion.Soundness

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

unit

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

subcircuit

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

FormalAssertion.toSubcircuit

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

localLength

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

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

FormalAssertion

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

SubcircuitsConsistent

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

FlatOperation.forAll

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

localLength

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

FormalAssertion.Completeness

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

ConstraintsHold.Completeness

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

Completeness

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

GeneralFormalCircuit.Completeness

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

toFlat

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

Environment.UsesLocalWitnesses

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

Environment.ExtendsVector

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

Operations.forAllFlat

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

GeneralFormalCircuit.toSubcircuit

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

FormalCircuit.toSubcircuit

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

fieldPair

@[reducible]
def fieldPair : TypeMap := fun F => F × F

field

@[reducible] def field : TypeMap := id

ConstraintsHold

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

toList

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

induct

def induct {motive : Operations F → Sort*}
  (empty : motive [])
  (witness : ∀ m c ops, motive ops → motive (.witness m c :: ops))
  (assert : ∀ e ops, motive ops → motive (.assert e :: ops))
  (lookup : ∀ l ops, motive ops → motive (.lookup l :: ops))
  (subcircuit : ∀ {n} (s : Subcircuit F n) ops, motive ops → motive (.subcircuit s :: ops))
    (ops : Operations F) : motive ops :=
  match ops with
  | [] => empty
  | .witness m c :: ops => witness m c ops (induct empty witness assert lookup subcircuit ops)
  | .assert e :: ops => assert e ops (induct empty witness assert lookup subcircuit ops)
  | .lookup l :: ops => lookup l ops (induct empty witness assert lookup subcircuit ops)
  | .subcircuit s :: ops => subcircuit s ops (induct empty witness assert lookup subcircuit ops)

forAll

@[reducible, circuit_norm]
def forAll (circuit : Circuit F α) (n : ℕ) (prop : Condition F) :=
  (circuit.operations n).forAll n prop

theorem toList_length_one {α : Type} (v : Vector α 1) :
    v.toList = [v[0]]

theorem one_land_of_IsBool (a : ℕ) (h : IsBool a) : 1 &&& a = a

theorem land_one_of_IsBool (a : ℕ) (h : IsBool a) : a &&& 1 = a

@[circuit_norm]
theorem pure_def {α} (a : α) : (pure a : Circuit F α) = fun _ => (a, [])

theorem val_of_IsBool {p : ℕ} [Fact p.Prime] {x : F p} (h : IsBool x) : IsBool x.val

one

@[circuit_norm]
theorem one {α : Type*} [Zero α] [One α] : IsBool (1 : α)

zero

@[circuit_norm]
theorem zero {α : Type*} [Zero α] [One α] : IsBool (0 : α)

@[circuit_norm ↓]
theorem eval_fieldPair {F : Type} [Field F] (env : Environment F) (t : Var fieldPair F) :
  ProvableType.eval env t = (match t with | (x, y) => (Expression.eval env x, Expression.eval env y))

theorem toList_length_two {α : Type} (v : Vector α 2) :
    v.toList = [v[0], v[1]]

theorem List.foldl_and_IsBool (l : List ℕ) :
    IsBool (List.foldl (· &&& ·) 1 l : ℕ)

theorem List.and_foldl_eq_foldl (a : ℕ) (orig : ℕ) (l : List ℕ) :
    a &&& List.foldl (· &&& ·) orig l = List.foldl (· &&& ·) (a &&& orig) l

@[circuit_norm ↓]
theorem eval_fields {F : Type} [Field F] (env : Environment F) (x : Var (fields n) F) :
  ProvableType.eval env x = x.map (Expression.eval env)

theorem append_take_drop {v : Vector α (n + m)} :
    (v.take n |>.cast Nat.min_add_right_self) ++ (v.drop n |>.cast (Nat.add_sub_self_left n m)) = v

ext

@[ext]
theorem ext {f g : Circuit F α}
  (h_output : ∀ n, f.output n = g.output n)
  (h_operations : ∀ n, f.operations n = g.operations n) :
    f = g

theorem ext_iff {f g : Circuit F α} :
  (f = g) ↔ (∀ n, (f.output n = g.output n) ∧ (f.operations n = g.operations n))

@[circuit_norm] theorem ConstraintsHold.bind_soundness {f : Circuit F α} {g : α → Circuit F β} (n : ℕ) :
  ConstraintsHold.Soundness env ((f >>= g).operations n)
  ↔ ConstraintsHold.Soundness env (f.operations n) ∧
    ConstraintsHold.Soundness env ((g (f.output n)).operations (n + f.localLength n))

@[circuit_norm]
theorem bind_forAll {f : Circuit F α} {g : α → Circuit F β} :
  ((f >>= g).operations n).forAll n prop ↔
    (f.operations n).forAll n prop ∧ (((g (f.output n)).operations (n + f.localLength n)).forAll (n + f.localLength n)) prop

@[circuit_norm]
theorem forAll_append {condition : Condition F} {offset : ℕ} {as bs: Operations F} :
  forAll offset condition (as ++ bs) ↔
    forAll offset condition as ∧ forAll (as.localLength + offset) condition bs

@[circuit_norm]
theorem forAll_empty {condition : Condition F} {n : ℕ} : forAll n condition [] = True

theorem ConstraintsHold.soundness_iff_forAll (n : ℕ) (env : Environment F) (ops : Operations F) :
  ConstraintsHold.Soundness env ops ↔ ops.forAll n {
    assert _ e := env e = 0,
    lookup _ l := l.table.Soundness (l.entry.map env),
    subcircuit _ _ s := s.Soundness env
  }

Circomlib.AND.main

def main (input : Expression (F p) × Expression (F p)) := do
  let a := input.1
  let b := input.2
  let out <== a*b
  return out

Circomlib.MultiAND.Assumptions

def Assumptions (n : ℕ) (input : fields n (F p)) : Prop :=
  ∀ (i : ℕ) (h : i < n), IsBool input[i]

Circomlib.MultiAND.Spec

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

lemma Vector.foldl_empty' {α β : Type} (init : β) (f : β → α → β) (v : Vector α 0) :
    Vector.foldl f init v = init

lemma Vector.foldl_and_split {n1 n2 n3 : ℕ} (v : Vector ℕ n3)
    (v1 : Vector ℕ n1) (v2 : Vector ℕ n2) (h_sum : n1 + n2 = n3)
    (h_split : v = h_sum ▸ (v1 ++ v2)) :
    Vector.foldl (· &&& ·) 1 v =
    Vector.foldl (· &&& ·) 1 v1 &&& Vector.foldl (· &&& ·) 1 v2

lemma soundness_zero {p : ℕ} [Fact p.Prime]
    (offset : ℕ) (env : Environment (F p)) (input_var : Var (fields 0) (F p))
    (input : fields 0 (F p)) (_h_env : input = eval env input_var)
    (_h_assumptions : Assumptions 0 input)
    (_h_hold : Circuit.ConstraintsHold.Soundness env ((main input_var).operations offset)) :
    Spec 0 input (env ((main input_var).output offset))

lemma soundness_one {p : ℕ} [Fact p.Prime]
    (offset : ℕ) (env : Environment (F p)) (input_var : Var (fields 1) (F p))
    (input : fields 1 (F p)) (h_env : input = eval env input_var)
    (h_assumptions : Assumptions 1 input)
    (_h_hold : Circuit.ConstraintsHold.Soundness env ((main input_var).operations offset)) :
    Spec 1 input (env ((main input_var).output offset))

lemma soundness_two {p : ℕ} [Fact p.Prime]
    (offset : ℕ) (env : Environment (F p)) (input_var : Var (fields 2) (F p))
    (input : fields 2 (F p)) (h_env : input = eval env input_var)
    (h_assumptions : Assumptions 2 input)
    (h_hold : Circuit.ConstraintsHold.Soundness env ((main input_var).operations offset)) :
    Spec 2 input (env ((main input_var).output offset))

import Clean.Circuit

import Clean.Utils.Field

import Clean.Gadgets.Boolean

import Clean.Utils.Bitwise

import Clean.Utils.Vector

import Clean.Utils.BinaryOps

import Clean.Circuit.Theorems

import Mathlib.Data.Nat.Bitwise

open IsBool

namespace Circomlib

variable {p : ℕ} [Fact p.Prime]

open Circuit (bind_output_eq bind_localLength_eq bind_forAll)

open Operations (append_localLength)

open BinaryOps (List.foldl_and_IsBool List.and_foldl_eq_foldl)

namespace XOR

end XOR

namespace AND

def main (input : Expression (F p) × Expression (F p)) := do
  let a := input.1
  let b := input.2
  let out <== a*b
  return out

end AND

namespace OR

end OR

namespace NOT

end NOT

namespace NAND

end NAND

namespace NOR

end NOR

namespace MultiAND

def main : {n : ℕ} → Vector (Expression (F p)) n → Circuit (F p) (Expression (F p))
  | 0, _ =>
    return (1 : F p)
  | 1, input =>
    return input[0]
  | 2, input =>
    AND.circuit.main (input[0], input[1])
  | n + 3, input => do
    let n1 := (n + 3) / 2
    let n2 := (n + 3) - n1

    let input1 : Vector (Expression (F p)) n1 := input.take n1 |>.cast (by admit /- proof elided -/
    )
    let input2 : Vector (Expression (F p)) n2 := input.drop n1 |>.cast (by admit /- proof elided -/
    )

    let out1 ← main input1
    let out2 ← main input2

    AND.circuit.main (out1, out2)

def Assumptions (n : ℕ) (input : fields n (F p)) : Prop :=
  ∀ (i : ℕ) (h : i < n), IsBool input[i]

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

theorem soundness {p : ℕ} [Fact p.Prime] (n : ℕ) :
    ∀ (offset : ℕ) (env : Environment (F p)) (input_var : Var (fields n) (F p))
      (input : fields n (F p)),
    input = eval env input_var →
    Assumptions n input →
    Circuit.ConstraintsHold.Soundness env ((main input_var).operations offset) →
    Spec n input (env ((main input_var).output offset)) :=

:= by
  induction n using Nat.strong_induction_on with
  | _ n IH =>
    intro offset env input_var input h_env h_assumptions h_hold
    match n with
    | 0 => exact soundness_zero offset env input_var input h_env h_assumptions h_hold
    | 1 => exact soundness_one offset env input_var input h_env h_assumptions h_hold
    | 2 => exact soundness_two offset env input_var input h_env h_assumptions h_hold
    | m + 3 =>
      simp only [main] at h_hold ⊢
      simp only [Spec]
      let n1 := (m + 3) / 2
      let n2 := (m + 3) - n1
      have h_sum : n1 + n2 = m + 3 := by unfold n1 n2; omega
      have h_n1_lt : n1 < m + 3 := by unfold n1; omega
      have h_n2_lt : n2 < m + 3 := by unfold n2; omega
      let input1 : fields n1 (F p) := input.take n1 |>.cast (by simp only [Nat.min_def, n1]; split <;> omega)
      let input2 : fields n2 (F p) := input.drop n1 |>.cast (by omega)
      let input_var1 : Var (fields n1) (F p) := input_var.take n1 |>.cast (by simp only [Nat.min_def, n1]; split <;> omega)
      let input_var2 : Var (fields n2) (F p) := input_var.drop n1 |>.cast (by omega)
      have h_eval1 : input1 = eval env input_var1 := by
        simp only [input_var1, input1]
        apply Vector.ext
        intro i hi
        simp only [h_env, ProvableType.eval_fields, Vector.getElem_map, Vector.getElem_cast, Vector.getElem_take]

      have h_eval2 : input2 = eval env input_var2 := by
        simp only [input_var2, input2]
        apply Vector.ext
        intro i hi
        simp only [h_env, ProvableType.eval_fields, Vector.getElem_map, Vector.getElem_cast, Vector.getElem_drop]

      have h_assumptions1 : Assumptions n1 input1 := by
        intro i hi
        -- input1[i] = input[i] since input1 is take of input
        simp only [input1]
        have : (input.take n1 |>.cast (by simp only [Nat.min_def, n1]; split <;> omega))[i]'hi = input[i]'(by omega) := by
          rw [Vector.getElem_cast, Vector.getElem_take]
        rw [this]
        apply h_assumptions i (by omega)
      have h_assumptions2 : Assumptions n2 input2 := by
        intro i hi
        -- input2[i] = input[n1 + i] since input2 is drop of input
        simp only [input2]
        have : (input.drop n1 |>.cast (by omega))[i]'hi = input[n1 + i]'(by omega) := by
          rw [Vector.getElem_cast, Vector.getElem_drop]
        rw [this]
        apply h_assumptions (n1 + i) (by omega)
      have h_spec1 : Spec n1 input1 (env ((main input_var1).output offset)) := by
        apply IH n1 h_n1_lt offset env input_var1 input1 h_eval1 h_assumptions1
        rw [Circuit.ConstraintsHold.bind_soundness] at h_hold
        exact h_hold.1
      have h_spec2 : Spec n2 input2 (env ((main input_var2).output (offset + (main input_var1).localLength offset))) := by
        apply IH n2 h_n2_lt (offset + (main input_var1).localLength offset) env input_var2 input2 h_eval2 h_assumptions2
        rw [Circuit.ConstraintsHold.bind_soundness] at h_hold
        rw [Circuit.ConstraintsHold.bind_soundness] at h_hold
        exact h_hold.2.1
      have h_hold' := h_hold
      rw [Circuit.ConstraintsHold.bind_soundness] at h_hold'
      rw [Circuit.ConstraintsHold.bind_soundness] at h_hold'
      let out1 := (main input_var1).output offset
      let out2 := (main input_var2).output (offset + (main input_var1).localLength offset)
      have h_and_spec := AND.circuit.soundness
        (offset + (main input_var1).localLength offset + (main input_var2).localLength (offset + (main input_var1).localLength offset))
        env
        (out1, out2)
        (env out1, env out2)
        (by simp only [ProvableType.eval_fieldPair])
        ⟨by rcases h_spec1 with ⟨_, h_binary1⟩; exact h_binary1,
         by rcases h_spec2 with ⟨_, h_binary2⟩; exact h_binary2⟩
        h_hold'.2.2

      rcases h_spec1 with ⟨h_val1, h_binary1⟩
      rcases h_spec2 with ⟨h_val2, h_binary2⟩
      rcases h_and_spec with ⟨h_and_val, h_and_binary⟩
      constructor
      · trans (Vector.foldl (fun x1 x2 => x1 &&& x2) 1 (input1.map (·.val)) &&&
               Vector.foldl (fun x1 x2 => x1 &&& x2) 1 (input2.map (·.val)))
        · convert h_and_val using 1
          simp only [out1, out2]
          simp only [h_val1, h_val2]

        have h_append : input1.cast (by omega : n1 = n1) ++ input2.cast (by omega : n2 = n2) =
                          input.cast (by omega : m + 3 = n1 + n2) := by
            simp only [input1, input2]
            have h_eq : n1 + n2 = m + 3 := by omega
            simp only [Vector.cast_cast]
            rw [← Vector.append_take_drop (n := n1) (m := n2) (v := input.cast h_eq.symm)]
            congr 1

        symm
        refine Vector.foldl_and_split (Vector.map (·.val) input) (Vector.map (·.val) input1) (Vector.map (·.val) input2) ?_ ?_
        · exact h_sum
        · have h_map_append : Vector.map (·.val) (input.cast (by omega : m + 3 = n1 + n2)) =
                             Vector.map (·.val) (input1.cast (by omega : n1 = n1) ++ input2.cast (by omega : n2 = n2)) := by
            congr 1
            exact h_append.symm

          simp only [Vector.map_append] at h_map_append

          have h1 : Vector.map (·.val) input = (Vector.map (·.val) (input.cast (by omega : m + 3 = n1 + n2))).cast h_sum := by
            ext i
            simp only [Vector.getElem_map, Vector.getElem_cast]

          have h2 : Vector.map (·.val) (input1.cast (by omega : n1 = n1)) = Vector.map (·.val) input1 := by
            ext i
            simp only [Vector.getElem_map, Vector.getElem_cast]

          have h3 : Vector.map (·.val) (input2.cast (by omega : n2 = n2)) = Vector.map (·.val) input2 := by
            ext i
            simp only [Vector.getElem_map, Vector.getElem_cast]

          rw [h1, h_map_append, h2, h3]

          have h_cast_transport : ∀ {n m : ℕ} (h : n = m) (v : Vector ℕ n),
                                  Vector.cast h v = h ▸ v := by
            intros n m h v
            subst h
            rfl

          rw [h_cast_transport]

      · exact h_and_binary


import Clean.Utils.Tactics
import Clean.Circuit.Provable
import Clean.Gadgets.Equality
import Clean.Utils.Primes
import Mathlib.Data.List.Sort
import Clean.Circuit.Basic
import Clean.Utils.Field

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

MemoryAccess

def MemoryAccess := ℕ × ℕ × ℕ × ℕ

MemoryAccessList.filterAddress

theorem MemoryAccessList.filterAddress_sorted (accesses : MemoryAccessList)
    (h : accesses.isTimestampSorted) (addr : ℕ) :
    (MemoryAccessList.filterAddress accesses addr).isTimestampSorted

theorem MemoryAccessList.filterAddress_cons (head : MemoryAccess) (tail : MemoryAccessList) (addr : ℕ) :
    MemoryAccessList.filterAddress (head :: tail) addr =
    match head with
    | ⟨_t, a, _r, _w⟩ => ((if a = addr then
      (head :: (MemoryAccessList.filterAddress tail addr))
        else (MemoryAccessList.filterAddress tail addr)))

theorem MemoryAccessList.isConsistentSingleAddress_iff (accesses : MemoryAccessList) (addr : ℕ) (h_sorted : accesses.isTimestampSorted)
    (h_eq : accesses.all (fun (_t, addr', _readValue, _writeValue) => addr' = addr)) :
    MemoryAccessList.isConsistentOnline accesses h_sorted ↔
    MemoryAccessList.isConsistentSingleAddress accesses h_sorted

theorem MemoryAccessList.lastWriteValue_filter (accesses : MemoryAccessList)
    (h_sorted : accesses.isTimestampSorted) (addr : ℕ) (h_sorted' : ((MemoryAccessList.filterAddress accesses addr).isTimestampSorted))  :
    MemoryAccessList.lastWriteValue accesses h_sorted addr =
    MemoryAccessList.lastWriteValue (MemoryAccessList.filterAddress accesses addr) h_sorted' addr

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)

theorem MemoryAccessList.isConsistentSingleAddress_cons (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isTimestampSorted (head :: tail)) (h_sorted' : tail.isTimestampSorted)
    (h : isConsistentSingleAddress (head :: tail) h_sorted) :
    isConsistentSingleAddress tail h_sorted'

theorem MemoryAccessList.isConsistentSingleAddress_cons_forall (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isTimestampSorted (head :: tail))
    : (∀ addr : ℕ, (filterAddress (head :: tail) addr).isConsistentSingleAddress (MemoryAccessList.filterAddress_sorted (head :: tail) h_sorted addr)) →
    (∀ addr : ℕ, isConsistentSingleAddress (filterAddress tail addr) (MemoryAccessList.filterAddress_sorted tail (by simp_all only [isTimestampSorted,
      List.sorted_cons]) addr))

import Clean.Circuit.Basic

import Clean.Circuit.Provable

import Clean.Gadgets.Equality

import Clean.Utils.Field

import Clean.Utils.Primes

import Clean.Utils.Tactics

import Mathlib.Data.List.Sort

def MemoryAccess := ℕ × ℕ × ℕ × ℕ

def MemoryAccessList := List MemoryAccess

abbrev timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, _a2, _r2, _w2), (t1, _a1, _r1, _w1) => t1 < t2

def MemoryAccessList.isTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted timestamp_ordering

def MemoryAccessList.lastWriteValue (accesses : MemoryAccessList) (h : accesses.isTimestampSorted) (addr : ℕ) : ℕ := match accesses with

  | [] => 0
  | (_t, addr', _readValue, writeValue) :: rest =>
    if addr' = addr then

      writeValue
    else
      MemoryAccessList.lastWriteValue rest (List.Sorted.of_cons h) addr

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)

example : MemoryAccessList.isConsistentOnline [] (by admit /- proof elided -/
) := by admit /- proof elided -/

def MemoryAccessList.filterAddress (accesses : MemoryAccessList) (addr : ℕ) : MemoryAccessList :=
  accesses.filter (fun (_timestamp, addr', _readValue, _writeValue) => addr' = addr)

def MemoryAccessList.isConsistentSingleAddress (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) : Prop := match accesses with

  | [] => True

  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0

  | (_t2, _addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    readValue2 = writeValue1 ∧
    MemoryAccessList.isConsistentSingleAddress ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

theorem MemoryAccessList.isConsistent_iff_all_single_address (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) :
    MemoryAccessList.isConsistentOnline accesses h_sorted ↔
    (∀ addr : ℕ, MemoryAccessList.isConsistentSingleAddress (MemoryAccessList.filterAddress accesses addr) (MemoryAccessList.filterAddress_sorted accesses h_sorted addr)) :=

:= by
  constructor
  · intro h addr
    have h' := MemoryAccessList.isConsistentSingleAddress_iff (accesses.filterAddress addr) addr (MemoryAccessList.filterAddress_sorted accesses h_sorted addr)
        (by simp only [filterAddress, List.all_filter, Bool.not_or_self, List.all_eq_true,
          implies_true])
    exact h'.mp (MemoryAccessList.isConsistentOnline_filter_of_consistentOnline accesses h_sorted h addr)
  · intro h
    -- by induction on the list of accesses
    induction accesses with
    | nil =>
      simp [isConsistentOnline]
    | cons head tail ih =>
      obtain ⟨t, a, r, w⟩ := head
      have h_sorted' : isTimestampSorted tail := by
        unfold isTimestampSorted at h_sorted
        exact List.Sorted.of_cons h_sorted
      specialize ih h_sorted'
      have h_tail := MemoryAccessList.isConsistentSingleAddress_cons_forall (t, a, r, w) tail h_sorted h
      specialize ih h_tail
      simp only [isConsistentOnline, ih, and_true]

      have h_tail_filter_sorted : (MemoryAccessList.filterAddress tail a).isTimestampSorted := by
        simp only [filterAddress]
        apply List.Sorted.filter
        exact h_sorted'

      have h_filtered_sorted : MemoryAccessList.isTimestampSorted ((t, a, r, w) :: (MemoryAccessList.filterAddress tail a)) := by
        simp only [isTimestampSorted, List.sorted_cons, filterAddress, List.mem_filter,
          and_imp] at ⊢ h_sorted
        simp_all only [implies_true, true_and]
        apply List.Sorted.filter
        assumption

      have h' := MemoryAccessList.lastWriteValue_filter tail h_sorted' a h_tail_filter_sorted

      specialize h a
      simp [filterAddress_cons] at h
      -- rw [h']

      have h_iff := MemoryAccessList.isConsistentSingleAddress_iff ((t, a, r, w) :: filterAddress tail a) a h_filtered_sorted
      simp only [filterAddress, List.all_cons, decide_true, List.all_filter, Bool.not_or_self,
        Bool.true_and, List.all_eq_true, implies_true, forall_const] at h_iff
      have is_consistent_online_full := h_iff.mpr h
      simp only [isConsistentOnline] at is_consistent_online_full
      simp only [is_consistent_online_full.left]
      rw [eq_comm]
      apply MemoryAccessList.lastWriteValue_filter


syntax "field_to_nat" : tactic

syntax "field_to_nat" : tactic

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

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

macro_rules
  | `(tactic|field_to_nat) =>
    `(tactic|(
      intros
      repeat rw [ZMod.val_add]
      repeat rw [ZMod.val_mul]
      repeat rw [val_eq_256]
      try simp only [Nat.add_mod_mod, Nat.mod_add_mod, Nat.mul_mod_mod, Nat.mod_mul_mod]
      rw [Nat.mod_eq_of_lt _]
      repeat linarith [‹Fact (_ > 512)›.elim]))

example [Fact (p > 512)] (x y : F p) (hx : x.val < 256) (hy : y.val < 2) :
    (x + y * 256).val = x.val + y.val * 256 := by admit /- proof elided -/

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

F

def F p := ZMod p

Inputs

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

HasAssignEq

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

IsBool

def IsBool {α : Type*} [Zero α] [One α] (x : α) : Prop := x = 0 ∨ x = 1

Outputs

structure Outputs (F : Type) where
  z: F
  carryOut: F

theorem val_lt_two {p : ℕ} [Fact p.Prime] {x : F p} (h : IsBool x) : x.val < 2

lemma lift_val1 {x y b : (F p)} (x_byte : x.val < 256) (y_byte : y.val < 256) (b_bool : IsBool b) :
    (x + y + b).val = (x.val + y.val + b.val)

lemma lift_val2 {x b : (F p)} (x_byte : x.val < 256) (b_bool : IsBool b) :
    (b * 256 + x).val = (b.val * 256 + x.val)

omit p_large_enough in
lemma zify_bool {b : (F p)} (b_bool : IsBool b) : (↑(b.val) : ℤ) = 0 ∨ (↑(b.val) : ℤ) = 1

import Clean.Gadgets.Addition8.Addition8FullCarry

import Clean.Types.U32

import Clean.Utils.Field

import Clean.Gadgets.Boolean

variable {p : ℕ} [Fact p.Prime]

variable [p_large_enough: Fact (p > 512)]

namespace Gadgets.Addition32.Theorems

theorem add32_soundness {x0 x1 x2 x3 y0 y1 y2 y3 carry_in c0 c1 c2 c3 z0 z1 z2 z3 : F p}
    (x0_byte : x0.val < 256) (x1_byte : x1.val < 256) (x2_byte : x2.val < 256) (x3_byte : x3.val < 256)
    (y0_byte : y0.val < 256) (y1_byte : y1.val < 256) (y2_byte : y2.val < 256) (y3_byte : y3.val < 256)
    (z0_byte : z0.val < 256) (z1_byte : z1.val < 256) (z2_byte : z2.val < 256) (z3_byte : z3.val < 256)
    (carry_in_bool : IsBool carry_in) (c0_bool : IsBool c0)
    (c1_bool : IsBool c1) (c2_bool : IsBool c2) (c3_bool : IsBool c3)
    (h0 : x0 + y0 + carry_in = c0 * 256 + z0)
    (h1 : x1 + y1 + c0 = c1 * 256 + z1)
    (h2 : x2 + y2 + c1 = c2 * 256 + z2)
    (h3 : x3 + y3 + c2 = c3 * 256 + z3) :
    ZMod.val z0 + ZMod.val z1 * 256 + ZMod.val z2 * 256 ^ 2 + ZMod.val z3 * 256 ^ 3 =
      (ZMod.val x0 + ZMod.val x1 * 256 + ZMod.val x2 * 256 ^ 2 + ZMod.val x3 * 256 ^ 3 +
            (ZMod.val y0 + ZMod.val y1 * 256 + ZMod.val y2 * 256 ^ 2 + ZMod.val y3 * 256 ^ 3) +
          ZMod.val carry_in) %
        2 ^ 32 ∧
    ZMod.val c3 =
      (ZMod.val x0 + ZMod.val x1 * 256 + ZMod.val x2 * 256 ^ 2 + ZMod.val x3 * 256 ^ 3 +
            (ZMod.val y0 + ZMod.val y1 * 256 + ZMod.val y2 * 256 ^ 2 + ZMod.val y3 * 256 ^ 3) +
          ZMod.val carry_in) /
        2 ^ 32 :=

:= by

  apply_fun ZMod.val at h0 h1 h2 h3
  rw [lift_val1 x0_byte y0_byte carry_in_bool, lift_val2 z0_byte c0_bool] at h0
  rw [lift_val1 x1_byte y1_byte c0_bool, lift_val2 z1_byte c1_bool] at h1
  rw [lift_val1 x2_byte y2_byte c1_bool, lift_val2 z2_byte c2_bool] at h2
  rw [lift_val1 x3_byte y3_byte c2_bool, lift_val2 z3_byte c3_bool] at h3

  have c0_bool := zify_bool c0_bool
  have c1_bool := zify_bool c1_bool
  have c2_bool := zify_bool c2_bool
  have c3_bool := zify_bool c3_bool
  have carry_in_bool := zify_bool carry_in_bool

  have x0_pos : 0 ≤ x0.val := by exact Nat.zero_le _
  have y0_pos : 0 ≤ y0.val := by exact Nat.zero_le _
  have x1_pos : 0 ≤ x1.val := by exact Nat.zero_le _
  have y1_pos : 0 ≤ y1.val := by exact Nat.zero_le _
  have x2_pos : 0 ≤ x2.val := by exact Nat.zero_le _
  have y2_pos : 0 ≤ y2.val := by exact Nat.zero_le _
  have x3_pos : 0 ≤ x3.val := by exact Nat.zero_le _
  have y3_pos : 0 ≤ y3.val := by exact Nat.zero_le _
  have z0_pos : 0 ≤ z0.val := by exact Nat.zero_le _
  have z1_pos : 0 ≤ z1.val := by exact Nat.zero_le _
  have z2_pos : 0 ≤ z2.val := by exact Nat.zero_le _
  have z3_pos : 0 ≤ z3.val := by exact Nat.zero_le _

  zify at *

  set x0 : ℤ := ↑(ZMod.val x0)
  set x1 : ℤ := ↑(ZMod.val x1)
  set x2 : ℤ := ↑(ZMod.val x2)
  set x3 : ℤ := ↑(ZMod.val x3)
  set y0 : ℤ := ↑(ZMod.val y0)
  set y1 : ℤ := ↑(ZMod.val y1)
  set y2 : ℤ := ↑(ZMod.val y2)
  set y3 : ℤ := ↑(ZMod.val y3)
  set z0 : ℤ := ↑(ZMod.val z0)
  set z1 : ℤ := ↑(ZMod.val z1)
  set z2 : ℤ := ↑(ZMod.val z2)
  set z3 : ℤ := ↑(ZMod.val z3)
  set c0 : ℤ := ↑(ZMod.val c0)
  set c1 : ℤ := ↑(ZMod.val c1)
  set c2 : ℤ := ↑(ZMod.val c2)
  set c3 : ℤ := ↑(ZMod.val c3)
  set carry_in : ℤ := ↑(ZMod.val carry_in)

  -- now everything, assumptions and goal, is over Z

  -- add up all the equations
  set z := z0 + z1*256 + z2*256^2 + z3*256^3
  set x := x0 + x1*256 + x2*256^2 + x3*256^3
  set y := y0 + y1*256 + y2*256^2 + y3*256^3
  let lhs := z + c3*2^32
  let rhs₀ := x0 + y0 + carry_in + -1 * z0 + -1 * (c0 * 256) -- h0 expression
  let rhs₁ := x1 + y1 + c0 + -1 * z1 + -1 * (c1 * 256) -- h1 expression
  let rhs₂ := x2 + y2 + c1 + -1 * z2 + -1 * (c2 * 256) -- h2 expression
  let rhs₃ := x3 + y3 + c2 + -1 * z3 + -1 * (c3 * 256) -- h3 expression

  have h_add := calc z + c3*2^32
    -- substitute equations
    _ = lhs + 0 + 256*0 + 256^2*0 + 256^3*0 := by ring
    _ = lhs + rhs₀ + 256*rhs₁ + 256^2*rhs₂ + 256^3*rhs₃ := by
      simp only [rhs₀, rhs₁, rhs₂, rhs₃]
      simp only [h0, h1, h2, h3]
      ring
    -- simplify
    _ = x + y + carry_in := by ring

  have z_lt_2_32 : z < 2^32 := by dsimp only [z]; linarith
  have z_pos : 0 ≤ z := by dsimp [z]; linarith

  rcases c3_bool with c3_zero | c3_one
  · rw [c3_zero] at h_add
    simp only [Int.reducePow, zero_mul, add_zero] at h_add
    have sum_pos : 0 ≤ x + y + carry_in := by
      rcases carry_in_bool with c_zero | c_one
      · rw [c_zero]; linarith
      · rw [c_one]; linarith
    rw [Int.emod_eq_of_lt sum_pos (by linarith)]
    rw [Int.ediv_eq_zero_of_lt sum_pos (by linarith)]
    simp only [h_add, c3_zero, and_self]
  · simp only [c3_one, one_mul] at h_add
    rw [←h_add, Int.add_emod_right, Int.emod_eq_of_lt z_pos (by linarith)]
    rw [
      show z + 2^32 = z + 2^32 * 1 from rfl,
      Int.add_mul_ediv_left _ _ (by linarith),
      Int.ediv_eq_zero_of_lt z_pos z_lt_2_32
    ]
    simp only [c3_one, zero_add, and_self]


import Clean.Types.U32
import Clean.Utils.Bitwise
import Clean.Types.U64
import Clean.Utils.Field

syntax "field_to_nat" : tactic

syntax "field_to_nat" : tactic

F

def F p := ZMod p

lemma two_pow_val (n : ℕ) (hn : n ≤ 8) : (2^n : F p).val = 2^n

omit p_prime in
lemma two_pow_lt (n : ℕ) (hn : n ≤ 8) : 2^n < p

lemma two_val : (2 : F p).val = 2

omit p_prime in
lemma two_lt : 2 < p

theorem val_lt_p {p : ℕ} (x : ℕ) : (x < p) → (x : F p).val = x

theorem byteDecomposition_lift {low high two_power : F p}
  (h_low : low.val < 2^8) (h_high : high.val < 2^8) (h_two_power : two_power.val ≤ 2^8) :
    (low + high * two_power).val = low.val + high.val * two_power.val

import Clean.Utils.Field

import Clean.Utils.Bitwise

import Clean.Types.U64

import Clean.Types.U32

namespace Gadgets.ByteDecomposition.Theorems

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

open FieldUtils (two_val two_pow_val)

theorem soundness (offset : Fin 8) (x low high : F p)
  (x_lt : x.val < 2^8) (low_lt : low.val < 2^offset.val) (high_lt : high.val < 2^8)
  (h_eq : x = low + high * 2^offset.val) :
    low.val = x.val % 2^offset.val ∧ high.val = x.val / 2^offset.val :=

:= by

  have two_power_lt : 2^offset.val < 2^8 := Nat.pow_lt_pow_of_lt (by linarith) offset.is_lt
  have two_power_val : ((2 : F p)^offset.val).val = 2^offset.val := two_pow_val offset.val (by linarith [offset.is_lt])
  have two_power_le : (2^offset.val : F p).val ≤ 2^8 := by rw [two_power_val]; linarith

  have low_byte : low.val < 256 := by linarith
  have h := byteDecomposition_lift low_byte high_lt two_power_le
  rw [two_power_val, ←h_eq] at h

  set low_b := UInt32.ofNat low.val
  set high_b := UInt32.ofNat high.val
  set x_b := UInt32.ofNat x.val

  -- The heavy lifting is done by using a SAT solver
  have h_decomposition_bv (base : UInt32) :
      base < 256 → low_b < base → high_b < 256 → x_b < 256 →
      x_b = low_b + high_b * base →
      low_b = x_b % base ∧ high_b = x_b / base := by
    bv_decide

  -- now it is left to prove that the bv variant is equivalent
  -- to the field variant of the theorem

  -- TODO: when updating to new Mathlib, all the following lemmas should be much easier to prove
  -- thanks new UInt32 lemmas
  have two_power_mod : (2^offset.val % 2 ^ 32) = 2^offset.val := by
    rw [Nat.mod_eq_iff_lt]
    linarith
    simp only [Nat.reducePow, ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true]

  have two_power_lt_bv : UInt32.ofNat (2^offset.val) < 256 := by
    rw [UInt32.lt_iff_toNat_lt]
    simp only [UInt32.toNat_ofNat', Nat.reducePow, UInt32.toNat_ofNat, Nat.reduceMod]
    rw [Nat.mod_eq_of_lt (by linarith)]
    linarith

  have low_mod : low.val % 2^32 = low.val := by
    rw [Nat.mod_eq_iff_lt]
    linarith
    simp only [Nat.reducePow, ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true]

  have high_mod : high.val % 2^32 = high.val := by
    rw [Nat.mod_eq_iff_lt]
    linarith
    simp only [Nat.reducePow, ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true]

  have x_mod : x.val % 2^32 = x.val := by
    rw [Nat.mod_eq_iff_lt]
    linarith
    simp only [Nat.reducePow, ne_eq, OfNat.ofNat_ne_zero, not_false_eq_true]

  have low_b_lt : low_b < UInt32.ofNat (2^offset.val) := by
    simp only [low_b]
    rw [UInt32.lt_iff_toNat_lt]
    simp only [UInt32.toNat_ofNat']
    rw [low_mod, two_power_mod]
    assumption

  have high_b_lt : high_b < 256 := by
    simp only [high_b]
    rw [UInt32.lt_iff_toNat_lt]
    simp only [UInt32.toNat_ofNat', UInt32.reduceToNat]
    rw [high_mod]
    assumption

  have x_lt : x_b < 256 := by
    simp only [x_b]
    rw [UInt32.lt_iff_toNat_lt]
    simp only [UInt32.toNat_ofNat', UInt32.reduceToNat]
    rw [x_mod]
    assumption

  have eq_holds_bv : x_b = low_b + high_b * UInt32.ofNat (2^offset.val) := by
    simp only [x_b, low_b, high_b]
    rw [←UInt32.toNat_inj]
    simp only [UInt32.toNat_ofNat', UInt32.toNat_add, UInt32.toNat_mul,
      Nat.mul_mod_mod, Nat.mod_mul_mod, Nat.add_mod_mod, Nat.mod_add_mod]
    rw [x_mod]
    have h : (low.val + high.val * (2^offset.val)) % 2^32 = low.val + high.val * (2^offset.val) := by
      apply Nat.mod_eq_of_lt
      linarith [p_large_enough.elim]
    rw [h]
    assumption

  specialize h_decomposition_bv (UInt32.ofNat (2^offset.val))
    two_power_lt_bv low_b_lt high_b_lt x_lt eq_holds_bv

  obtain ⟨h1, h2⟩ := h_decomposition_bv

  constructor
  · apply_fun UInt32.toNat at h1
    simp only [UInt32.toNat_ofNat', UInt32.toNat_mod, low_b, x_b] at h1
    rw [low_mod, x_mod, two_power_mod] at h1
    assumption
  · apply_fun UInt32.toNat at h2
    simp only [UInt32.toNat_ofNat', UInt32.toNat_div, high_b, x_b] at h2
    rw [high_mod, x_mod, two_power_mod] at h2
    assumption


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

MemoryAccess

def MemoryAccess := ℕ × ℕ × ℕ × ℕ

MemoryAccessList.isConsistentOffline

theorem MemoryAccessList.addressStrictTimestampSorted_of_AddressTimestampSorted_noTimestampDup
    (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted)
    (h_no_timestamp_dup : accesses.Notimestampdup) :
    accesses.isAddressStrictTimestampSorted

theorem MemoryAccessList.noTimestampDup_of_cons (head : MemoryAccess) (tail : MemoryAccessList)
    (h : Notimestampdup (head :: tail)) :
    Notimestampdup tail

theorem MemoryAccessList.isAddressTimestampSorted_of_cons (head : MemoryAccess) (tail : MemoryAccessList)
    (h : isAddressTimestampSorted (head :: tail)) :
    isAddressTimestampSorted tail

theorem MemoryAccessList.filterAddress_sorted_from_addressTimestampSorted
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isAddressTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (addr : ℕ) :
    (accesses.filterAddress addr).isTimestampSorted

theorem MemoryAccessList.filterAddress_empty_when_address_changes
    (head : MemoryAccess) (second : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: second :: tail))
    (h_addr_ne : head.2.1 ≠ second.2.1) :
    filterAddress (second :: tail) head.2.1 = []

theorem MemoryAccessList.isConsistentOffline_of_cons
    (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: tail))
    (h_offline : isConsistentOffline (head :: tail) h_sorted) :
    isConsistentOffline tail (isAddressTimestampSorted_of_cons head tail h_sorted)

import Clean.Circuit.Basic

import Clean.Circuit.Provable

import Clean.Gadgets.Equality

import Clean.Utils.Field

import Clean.Utils.Primes

import Clean.Utils.Tactics

import Mathlib.Data.List.Sort

def MemoryAccess := ℕ × ℕ × ℕ × ℕ

def MemoryAccessList := List MemoryAccess

abbrev timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, _a2, _r2, _w2), (t1, _a1, _r1, _w1) => t1 < t2

def MemoryAccessList.isTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted timestamp_ordering

def MemoryAccessList.timestamps_neq (x y: MemoryAccess) : Prop :=
  match x, y with
  | (t_x, _a_x, _r_x, _w_x), (t_y, _a_y, _r_y, _w_y) => t_x ≠ t_y

def MemoryAccessList.Notimestampdup (accesses : MemoryAccessList) : Prop :=
  List.Pairwise timestamps_neq accesses

abbrev address_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 ≤ t2 else a1 < a2

abbrev address_strict_timestamp_ordering (x y : MemoryAccess) := match x, y with
| (t2, a2, _, _), (t1, a1, _, _) => if a1 = a2 then t1 < t2 else a1 < a2

@[reducible]
def MemoryAccessList.isAddressTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_timestamp_ordering

@[reducible]
def MemoryAccessList.isAddressStrictTimestampSorted (accesses : MemoryAccessList) : Prop :=
  accesses.Sorted address_strict_timestamp_ordering

def MemoryAccessList.filterAddress (accesses : MemoryAccessList) (addr : ℕ) : MemoryAccessList :=
  accesses.filter (fun (_timestamp, addr', _readValue, _writeValue) => addr' = addr)

def MemoryAccessList.isConsistentSingleAddress (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) : Prop := match accesses with

  | [] => True

  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0

  | (_t2, _addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    readValue2 = writeValue1 ∧
    MemoryAccessList.isConsistentSingleAddress ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

def MemoryAccessList.isConsistentOffline (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted) : Prop := match accesses with
  | [] => True
  | (_timestamp, _addr, readValue, _writeValue) :: [] => readValue = 0
  | (_t2, addr2, readValue2, _writeValue2) :: (t1, addr1, readValue1, writeValue1) :: rest =>
    (if addr1 = addr2 then readValue2 = writeValue1 else readValue2 = 0) ∧
    MemoryAccessList.isConsistentOffline ((t1, addr1, readValue1, writeValue1) :: rest) (List.Sorted.of_cons h_sorted)

theorem MemoryAccessList.isConsistentOffline_implies_single_address
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isAddressTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (h_offline : accesses.isConsistentOffline h_sorted)
    (addr : ℕ) :
    (accesses.filterAddress addr).isConsistentSingleAddress
      (filterAddress_sorted_from_addressTimestampSorted accesses h_sorted h_nodup addr) :=

:= by
  induction accesses with
  | nil =>
    simp [filterAddress, isConsistentSingleAddress]
  | cons head tail ih =>
    obtain ⟨t_head, a_head, r_head, w_head⟩ := head
    simp only [filterAddress, List.filter_cons]
    split_ifs with h_addr
    · -- head has the target address
      rw [decide_eq_true_eq] at h_addr
      subst h_addr
      -- Now we need to show consistency for (head :: filterAddress tail addr)
      cases tail with
      | nil =>
        simp [isConsistentSingleAddress]
        -- For a singleton list, need to show r_head = 0
        exact h_offline
      | cons second tail_rest =>
        obtain ⟨t_second, a_second, r_second, w_second⟩ := second
        -- Filter the tail
        simp only [List.filter_cons]
        split_ifs with h_second_addr
        · -- second also has address a_head
          rw [decide_eq_true_eq] at h_second_addr
          subst h_second_addr
          -- Now we show consistency for (head :: second :: filterAddress tail_rest a_head)
          simp only [isConsistentSingleAddress]
          constructor
          · -- Show r_second = w_head from offline consistency
            simp only [isConsistentOffline] at h_offline
            simp at h_offline
            exact h_offline.1
          · -- Apply IH to (second :: tail_rest)
            have ih_applied := ih
              (isAddressTimestampSorted_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_sorted)
              (noTimestampDup_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_nodup)
              (isConsistentOffline_of_cons (t_head, a_second, r_head, w_head) ((t_second, a_second, r_second, w_second) :: tail_rest) h_sorted h_offline)
            simp only [filterAddress, List.filter_cons, decide_true] at ih_applied
            exact ih_applied
        · -- second has a different address
          -- Need to show isConsistentSingleAddress for (head :: filterAddress tail_rest a_head)
          -- Key insight: since a_head ≠ a_second and list is address-sorted,
          -- all elements in (second :: tail_rest) have address ≠ a_head
          have h_addr_ne : ¬(a_second = a_head) := by
            simp only [decide_eq_true_eq] at h_second_addr
            exact h_second_addr
          -- Use filterAddress_empty_when_address_changes to show filterAddress (second :: tail_rest) a_head = []
          have h_empty := filterAddress_empty_when_address_changes
            (t_head, a_head, r_head, w_head) (t_second, a_second, r_second, w_second) tail_rest h_sorted
            (by simp; intro h_eq; exact h_addr_ne h_eq.symm)
          simp only [filterAddress, List.filter_cons] at h_empty
          have h_second_ne : decide (a_second = a_head) = false := by
            simp [h_addr_ne]
          simp only [h_second_ne] at h_empty
          have h_empty_simplified : List.filter (fun x => match x with | (_, addr', _, _) => decide (addr' = a_head)) tail_rest = [] := by
            simp at h_empty
            exact h_empty
          simp only [h_empty_simplified, isConsistentSingleAddress]
          simp only [isConsistentOffline] at h_offline
          simp [h_addr_ne] at h_offline
          exact h_offline.1
    · -- head doesn't have the target address
      -- Apply IH to tail
      apply ih (isAddressTimestampSorted_of_cons (t_head, a_head, r_head, w_head) tail h_sorted)
        (noTimestampDup_of_cons (t_head, a_head, r_head, w_head) tail h_nodup)
      exact isConsistentOffline_of_cons (t_head, a_head, r_head, w_head) tail h_sorted h_offline


import Clean.Circuit.Theorems
import Clean.Circuit.Provable
import Clean.Utils.Field
import Clean.Circuit
import Mathlib.Data.Nat.Bitwise
import Clean.Gadgets.Boolean
import Clean.Utils.Bitwise
import Clean.Utils.BinaryOps
import Clean.Utils.Vector

syntax "infer_constant_length" : tactic

syntax "infer_constant_length" : tactic

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

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

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

macro_rules
  | `(tactic|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 -/

F

def F p := ZMod p

fields

@[reducible]
def fields (n : ℕ) := fun F => Vector F n

IsBool

def IsBool {α : Type*} [Zero α] [One α] (x : α) : Prop := x = 0 ∨ x = 1

map

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

Circuit

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

Operation

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

Condition

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

FlatOperation

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

Subcircuit

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

localWitnesses

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

lookup

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

Table.toRaw

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

RawTable

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

StaticTable

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

Contains

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

ProvableType.fromStruct

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

TypeMap

@[reducible]
def TypeMap := Type → Type

Table

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

Expression

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

Variable

structure Variable (F : Type) where
  index : ℕ

const

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

Var

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

fromVars

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

Lookup

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

Witnessable

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

Environment

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

localLength

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

ConstraintsHoldFlat

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

ConstantLength

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

ConstantLength.fromConstantLength

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

HasAssignEq

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

FormalCircuit

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

ElaboratedCircuit

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

GeneralFormalCircuit

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

Soundness

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

operations

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

Operations

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

output

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

ConstraintsHold.Soundness

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

GeneralFormalCircuit.Soundness

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

eval

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

toVars

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

eval

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

FormalAssertion.Soundness

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

unit

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

subcircuit

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

FormalAssertion.toSubcircuit

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

localLength

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

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

FormalAssertion

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

SubcircuitsConsistent

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

FlatOperation.forAll

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

localLength

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

FormalAssertion.Completeness

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

ConstraintsHold.Completeness

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

Completeness

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

GeneralFormalCircuit.Completeness

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

toFlat

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

Environment.UsesLocalWitnesses

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

Environment.ExtendsVector

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

Operations.forAllFlat

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

GeneralFormalCircuit.toSubcircuit

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

FormalCircuit.toSubcircuit

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

fieldPair

@[reducible]
def fieldPair : TypeMap := fun F => F × F

field

@[reducible] def field : TypeMap := id

ConstraintsHold

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

@[circuit_norm ↓]
theorem eval_fieldPair {F : Type} [Field F] (env : Environment F) (t : Var fieldPair F) :
  ProvableType.eval env t = (match t with | (x, y) => (Expression.eval env x, Expression.eval env y))

theorem toList_length_two {α : Type} (v : Vector α 2) :
    v.toList = [v[0], v[1]]

theorem one_land_of_IsBool (a : ℕ) (h : IsBool a) : 1 &&& a = a

theorem land_one_of_IsBool (a : ℕ) (h : IsBool a) : a &&& 1 = a

theorem val_of_IsBool {p : ℕ} [Fact p.Prime] {x : F p} (h : IsBool x) : IsBool x.val

one

@[circuit_norm]
theorem one {α : Type*} [Zero α] [One α] : IsBool (1 : α)

zero

@[circuit_norm]
theorem zero {α : Type*} [Zero α] [One α] : IsBool (0 : α)

Circomlib.AND.main

def main (input : Expression (F p) × Expression (F p)) := do
  let a := input.1
  let b := input.2
  let out <== a*b
  return out

Circomlib.MultiAND.main

def main : {n : ℕ} → Vector (Expression (F p)) n → Circuit (F p) (Expression (F p))
  | 0, _ =>
    return (1 : F p)
  | 1, input =>
    return input[0]
  | 2, input =>
    AND.circuit.main (input[0], input[1])
  | n + 3, input => do
    let n1 := (n + 3) / 2
    let n2 := (n + 3) - n1

    let input1 : Vector (Expression (F p)) n1 := input.take n1 |>.cast (by admit /- proof elided -/
    )
    let input2 : Vector (Expression (F p)) n2 := input.drop n1 |>.cast (by admit /- proof elided -/
    )

    let out1 ← main input1
    let out2 ← main input2

    AND.circuit.main (out1, out2)

Circomlib.MultiAND.Assumptions

def Assumptions (n : ℕ) (input : fields n (F p)) : Prop :=
  ∀ (i : ℕ) (h : i < n), IsBool input[i]

Circomlib.MultiAND.Spec

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

lemma soundness_two {p : ℕ} [Fact p.Prime]
    (offset : ℕ) (env : Environment (F p)) (input_var : Var (fields 2) (F p))
    (input : fields 2 (F p)) (h_env : input = eval env input_var)
    (h_assumptions : Assumptions 2 input)
    (h_hold : Circuit.ConstraintsHold.Soundness env ((main input_var).operations offset)) :
    Spec 2 input (env ((main input_var).output offset)) :=

:= by
  simp only [main] at h_hold ⊢
  simp only [Spec]
  have h_input0 := h_assumptions 0 (by norm_num : 0 < 2)
  have h_input1 := h_assumptions 1 (by norm_num : 1 < 2)
  have h_eval0 : env input_var[0] = input[0] := by simp [h_env, circuit_norm]
  have h_eval1 : env input_var[1] = input[1] := by simp [h_env, circuit_norm]
  have h_and_spec := AND.circuit.soundness offset env (input_var[0], input_var[1])
    (input[0], input[1])
    (by simp only [ProvableType.eval_fieldPair, h_eval0, h_eval1])
    ⟨h_input0, h_input1⟩ h_hold

  rcases h_and_spec with ⟨h_val, h_binary⟩
  constructor
  · -- Prove output.val = fold
    have h_fold_two : Vector.foldl (fun x1 x2 => x1 &&& x2) 1 (input.map (·.val)) = input[0].val &&& input[1].val := by
      rw [Vector.foldl_mk, ← Array.foldl_toList]
      have h_toList : (input.map (·.val)).toList = [input[0].val, input[1].val] := by
        rw [Vector.toList_length_two]
        simp only [Vector.getElem_map]
      rw [Vector.toList_toArray, h_toList]
      simp only [List.foldl_cons, List.foldl_nil]
      rw [one_land_of_IsBool input[0].val (val_of_IsBool h_input0)]
    rw [h_fold_two]
    exact h_val
  · exact h_binary