-- Clean/Circomlib/Gates.lean

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  
  
  localLength : ℕ

  
  imply_soundness : ∀ env,
    ConstraintsHoldFlat env ops → Soundness env

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

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

  
  localLength_eq : localLength = FlatOperation.localLength ops

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

content: 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: structure GeneralFormalCircuit (F : Type) (Input Output : TypeMap) [Field F] [ProvableType Input] [ProvableType Output]
    extends elaborated : ElaboratedCircuit F Input Output where
  Assumptions : Input F → Prop 
  Spec : Input F → Output F → Prop 
  soundness : GeneralFormalCircuit.Soundness F elaborated Spec
  completeness : GeneralFormalCircuit.Completeness F elaborated Assumptions

content: @[circuit_norm]
def 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 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 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: @[reducible]
def fieldPair : TypeMap := fun F => F × F

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

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

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

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

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

-- Repository lemmas used:
content: theorem toList_length_one {α : Type} (v : Vector α 1) :
    v.toList = [v[0]]

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

-- Target theorem and proof:
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