import Clean.Circuit.Provable
import Clean.Circuit.Subcircuit
import Clean.Gadgets.Equality
import Clean.Utils.Primes
import Clean.Utils.Bitwise
import Clean.Circuit.Extensions
import Clean.Gadgets.ByteLookup

syntax "infer_constant_length" : tactic

syntax "infer_constant_length" : tactic

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

KeccakRow.value

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

map

def map {α β : Type} (x : U64 α) (f : α → β) : U64 β :=
  ⟨ f x.x0, f x.x1, f x.x2, f x.x3, f x.x4, f x.x5, f x.x6, f x.x7 ⟩

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

value

def value (x : U64 (F p)) :=
  x.x0.val + x.x1.val * 256 + x.x2.val * 256^2 + x.x3.val * 256^3 +
  x.x4.val * 256^4 + x.x5.val * 256^5 + x.x6.val * 256^6 + x.x7.val * 256^7

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

KeccakRow.Normalized

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

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

KeccakBlock.value

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

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

KeccakState.value

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

KeccakBlock.Normalized

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

RATE

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

CAPACITY

@[reducible] def CAPACITY := 8

KeccakState.Normalized

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

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

private lemma reorganize_value (a b c d : ℕ) :
  a + 256 * (b + 256 * (c + 256 * d)) =
  2^8 * (2^8 * (2^8 * d + c) + b) + a

private lemma reorganize_value' (a b c d : ℕ) :
  a + b * 256 + c * 256 ^ 2 + d * 256 ^ 3 =
  2^8 * (2^8 * (2^8 * d + c) + b) + a

import Clean.Gadgets.ByteLookup

import Clean.Circuit.Extensions

import Clean.Utils.Bitwise

import Clean.Circuit.Provable

import Clean.Utils.Primes

import Clean.Circuit.Subcircuit

import Clean.Gadgets.Equality

section

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

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

namespace U32

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

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

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

end U32

namespace U32.AssertNormalized

open Gadgets (ByteTable)

end U32.AssertNormalized

namespace U32.ByteVector

end ByteVector

section Bitwise

omit [Fact (Nat.Prime p)] p_large_enough in
lemma bitwise_componentwise (f : Bool → Bool → Bool)
    {x y : U32 (F p)} (x_norm : x.Normalized) (y_norm : y.Normalized) :
    f false false = false →
    Nat.bitwise f x.value y.value =
      Nat.bitwise f x.x0.val y.x0.val + 256 *
        (Nat.bitwise f x.x1.val y.x1.val + 256 *
          (Nat.bitwise f x.x2.val y.x2.val + 256 * Nat.bitwise f x.x3.val y.x3.val)) :=

:= by
  intro _
  simp only [value]

  have ⟨_, _, _, _⟩ := x_norm
  have ⟨_, _, _, _⟩ := y_norm
  apply Nat.eq_of_testBit_eq
  intro i
  simp only [reorganize_value, reorganize_value']
  rw [Nat.testBit_bitwise] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val x.x0)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val y.x0)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=Nat.bitwise f (ZMod.val x.x0) (ZMod.val y.x0))] <;> try (apply Nat.bitwise_lt_two_pow <;> assumption)
  split
  · simp_all only [Nat.testBit_bitwise]
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val x.x1)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val y.x1)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=Nat.bitwise f (ZMod.val x.x1) (ZMod.val y.x1))] <;> try (apply Nat.bitwise_lt_two_pow <;> assumption)
  split
  · simp_all only [not_lt, Nat.testBit_bitwise]
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val x.x2)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=ZMod.val y.x2)] <;> try assumption
  rw [Nat.testBit_two_pow_mul_add (i:=8) (b:=Nat.bitwise f (ZMod.val x.x2) (ZMod.val y.x2))] <;> try (apply Nat.bitwise_lt_two_pow <;> assumption)
  aesop


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

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.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.filterAddress_sorted_from_addressTimestampSorted
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isAddressTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (addr : ℕ) :
    (accesses.filterAddress addr).isTimestampSorted

theorem MemoryAccessList.isConsistentSingleAddress_filterAddress_forall_of_cons
    (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: tail))
    (h_nodup : Notimestampdup (head :: tail))
    (h : ∀ addr, isConsistentSingleAddress (filterAddress (head :: tail) addr)
         (filterAddress_sorted_from_addressTimestampSorted (head :: tail) h_sorted h_nodup addr)) :
    ∀ addr, isConsistentSingleAddress (filterAddress tail addr)
      (filterAddress_sorted_from_addressTimestampSorted tail
        (isAddressTimestampSorted_of_cons head tail h_sorted)
        (noTimestampDup_of_cons head tail h_nodup) addr)

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)

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)

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_iff_all_single_addresses (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted) (h_nodup : accesses.Notimestampdup) :
    MemoryAccessList.isConsistentOffline accesses h_sorted ↔
    ∀ addr, MemoryAccessList.isConsistentSingleAddress (MemoryAccessList.filterAddress accesses addr) (filterAddress_sorted_from_addressTimestampSorted accesses h_sorted h_nodup addr) :=

:= by
  constructor
  · intro h_offline addr
    exact isConsistentOffline_implies_single_address accesses h_sorted h_nodup h_offline addr
  · induction accesses
    · simp [isConsistentOffline]
    rename_i hd tl h_ih
    intro h
    rcases tl
    · rcases hd with ⟨ hd_t, hd_a, hd_r, hd_w ⟩
      simp only [isConsistentOffline]
      specialize h hd_a
      simp only [filterAddress, List.filter, decide_true, isConsistentSingleAddress] at h
      assumption
    rename_i snd tl
    rcases hd with ⟨ hd_t, hd_a, hd_r, hd_w ⟩
    rcases snd with ⟨ snd_t, snd_a, snd_r, snd_w ⟩
    simp only [isConsistentOffline]
    and_intros
    · split
      · rename_i addr_eq
        subst addr_eq
        specialize h snd_a
        simp only [filterAddress, List.filter, decide_true, isConsistentSingleAddress] at h
        aesop
      · -- addresstimestampsorted, and seeing two different addresses, then the first address will never appear again
        -- Since hd_a ≠ snd_a and the list is address-timestamp sorted, hd_a won't appear in (snd :: tl)
        -- Therefore filterAddress ((hd :: snd :: tl)) hd_a = [hd]
        -- And since isConsistentSingleAddress [hd] must hold, we get hd_r = 0
        rename_i h_addr_ne
        -- Use the lemma to show filterAddress (snd :: tl) hd_a is empty
        have h_empty := filterAddress_empty_when_address_changes (hd_t, hd_a, hd_r, hd_w) (snd_t, snd_a, snd_r, snd_w) tl h_sorted (by simp; intro h_eq; exact h_addr_ne h_eq.symm)
        specialize h hd_a
        simp only [filterAddress, List.filter_cons, decide_true] at h
        have h_snd_ne : decide (snd_a = hd_a) = false := by
          simp only [decide_eq_false_iff_not]
          exact h_addr_ne
        simp only [h_snd_ne, ↓reduceIte] at h
        simp only [filterAddress, List.filter_cons, h_snd_ne] at h_empty
        simp only [h_empty, isConsistentSingleAddress] at h
        exact h
    apply h_ih
    · exact isConsistentSingleAddress_filterAddress_forall_of_cons (hd_t, hd_a, hd_r, hd_w) ((snd_t, snd_a, snd_r, snd_w) :: tl) h_sorted h_nodup h
    · exact noTimestampDup_of_cons (hd_t, hd_a, hd_r, hd_w) ((snd_t, snd_a, snd_r, snd_w) :: tl) h_nodup


import Clean.Circuit.Provable
import Clean.Circuit.Basic

syntax "infer_constant_length" : tactic

syntax "infer_constant_length" : tactic

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

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

FlatOperation.dynamicWitness

def FlatOperation.dynamicWitness (op : FlatOperation F) (acc : List F) : List F := match op with
  | .witness _ compute => (compute (.fromList acc)).toList
  | .assert _ => []
  | .lookup _ => []

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

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

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

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

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

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

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

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)

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

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

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]

Environment.AgreesBelow

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

Environment.UsesLocalWitnessesFlat

def Environment.UsesLocalWitnessesFlat (env : Environment F) (n : ℕ) (ops : List (FlatOperation F)) : Prop :=
  FlatOperation.forAll n { witness n _ compute := env.ExtendsVector (compute env) n } ops

Condition.applyFlat

def Condition.applyFlat (condition : Condition F) (offset : ℕ) : FlatOperation F → Prop
  | .witness m c => condition.witness offset m c
  | .assert e => condition.assert offset e
  | .lookup l => condition.lookup offset l

Environment.fromList

def Environment.fromList (witnesses : List F) : Environment F :=
  .mk fun i => witnesses[i]?.getD 0

FlatOperation.proverEnvironment

def FlatOperation.proverEnvironment (ops : List (FlatOperation F)) (init : List F) : Environment F :=
  .fromList (FlatOperation.dynamicWitnesses ops init)

FlatOperation.dynamicWitnesses

def FlatOperation.dynamicWitnesses (ops : List (FlatOperation F)) (init : List F) : List F :=
  ops.foldl (fun (acc : List F) (op : FlatOperation F) =>
    acc ++ op.dynamicWitness acc
  ) init

foldl

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

ConstantLength

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

ConstantOutput

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

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

FlatOperation.singleLocalLength

def FlatOperation.singleLocalLength : FlatOperation F → ℕ
  | .witness m _ => m
  | .assert _ => 0
  | .lookup _ => 0

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

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

lemma dynamicWitness_length {op : FlatOperation F} {init : List F} :
    (op.dynamicWitness init).length = op.singleLocalLength

lemma dynamicWitnesses_cons {op : FlatOperation F} {ops : List (FlatOperation F)} {acc : List F} :
    dynamicWitnesses (op :: ops) acc = dynamicWitnesses ops (acc ++ op.dynamicWitness acc)

lemma getElem?_dynamicWitnesses_of_lt {ops : List (FlatOperation F)} {acc : List F} {i : ℕ} (hi : i < acc.length) :
    (dynamicWitnesses ops acc)[i]?.getD 0 = acc[i]

lemma getElem?_dynamicWitnesses_cons_right {op : FlatOperation F} {ops : List (FlatOperation F)} {init : List F} {i : ℕ} (hi : i < op.singleLocalLength) :
    (dynamicWitnesses (op :: ops) init)[init.length + i]?.getD 0 = (op.dynamicWitness init)[i]'(dynamicWitness_length (F:=F) ▸ hi)

import Clean.Circuit.Basic

import Clean.Circuit.Provable

variable {F : Type} [Field F] {α β : Type}

namespace Operations

end Operations

namespace Circuit

end Circuit

namespace FlatOperation

end FlatOperation

namespace Environment

open FlatOperation (localLength localWitnesses)

end Environment

namespace Circuit

end Circuit

namespace Circuit

variable {α β : Type} {n : ℕ} {prop : Condition F} {env : Environment F}

end Circuit

namespace FlatOperation

end FlatOperation

namespace Operations

end Operations

namespace FlatOperation

theorem proverEnvironment_usesLocalWitnesses {ops : List (FlatOperation F)} (init : List F) :
  (∀ (env env' : Environment F),
    forAll init.length { witness n _ c := env.AgreesBelow n env' → c env = c env' } ops) →
    (proverEnvironment ops init).UsesLocalWitnessesFlat init.length ops :=

:= by
  simp only [proverEnvironment, Environment.UsesLocalWitnessesFlat, Environment.ExtendsVector]
  intro h_computable
  induction ops generalizing init with
  | nil => trivial
  | cons op ops ih =>
    simp only [forAll_cons] at h_computable ⊢
    cases op with
    | assert | lookup  =>
      simp_all [dynamicWitnesses_cons, Condition.applyFlat, singleLocalLength, dynamicWitness]
    | witness m compute =>
      simp_all only [Condition.applyFlat, singleLocalLength, Environment.AgreesBelow]
      -- get rid of ih first
      constructor; case right =>
        specialize ih (init ++ (compute (.fromList init)).toList)
        simp only [List.length_append, Vector.length_toList] at ih
        ring_nf at *
        exact ih fun _ _ => (h_computable ..).right
      clear ih
      replace h_computable := fun env env' => (h_computable env env').left
      intro i
      simp only [Environment.fromList]
      rw [getElem?_dynamicWitnesses_cons_right i.is_lt]
      simp only [dynamicWitness, Vector.getElem_toList]
      congr 1
      apply h_computable
      intro j hj
      simp [Environment.fromList, hj, getElem?_dynamicWitnesses_of_lt]


Utils.StateTransition.Transition

def Transition (S : Type*) := S × S

Utils.StateTransition.Run

def Run (S : Type*) := Transition S → ℕ

Utils.StateTransition.Run.isLeaf

lemma finset_ssubset_univ_of_not_mem {α : Type*} [Fintype α] (s : Finset α) (x : α)
    (h : x ∉ s) :
    s ⊂ Finset.univ

lemma zip_drop_sublist (l : List S) (n : ℕ) :
    ((l.drop n).zip (l.drop (n + 1))).Sublist (l.zip l.tail)

lemma containsPath_drop (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.drop n)

lemma zip_take_sublist (l1 l2 : List S) (n m : ℕ) :
    ((l1.take n).zip (l2.take m)).Sublist (l1.zip l2)

lemma tail_take {α : Type*} (l : List α) (n : ℕ) :
    (l.take n).tail = (l.tail).take (n - 1)

lemma containsPath_take (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.take n)

lemma drop_take_length_ge_two {α : Type*} (path : List α) (n m : Fin path.length)
    (h_n_lt_m : n < m) :
    ((path.drop n.val).take (m.val - n.val + 1)).length ≥ 2

lemma getLast_drop_take {α : Type*} (path : List α) (n k : ℕ)
    (h_n_lt : n < path.length)
    (h_bound : n + k ≤ path.length)
    (h_k_pos : k > 0) :
    ((path.drop n).take k).getLast? = path[n + k - 1]?

lemma drop_take_cycle_same_endpoints (path : List S) (x : S) (n m : Fin path.length)
    (h_n_lt_m : n < m)
    (h_x_at_n : path[n] = x)
    (h_x_at_m : path[m] = x) :
    ((path.drop n.val).take (m.val - n.val + 1)).head? =
    ((path.drop n.val).take (m.val - n.val + 1)).getLast?

lemma containsPath_drop_take (R : Run S) (path : List S) (n m : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath ((path.drop n).take m)

lemma acyclic_containsPath_nodup (R : Run S) (path : List S)
    (h_acyclic : R.isAcyclic)
    (h_contains : R.containsPath path) :
    path.Nodup

lemma countTransitionInPath_append_singleton (path : List S) (x y : S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_not_in : (x, y) ∉ path.zip path.tail) :
    countTransitionInPath (x, y) (path ++ [y]) = 1

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

lemma acyclic_no_self_loop (R : Run S) (s : S) (h_acyclic : R.isAcyclic) (h_edge : R (s, s) > 0) : False

lemma getLast_mem {α : Type*} (l : List α) (x : α) (h_last : l.getLast? = some x) :
    x ∈ l

lemma last_not_in_zip_tail {α : Type*} [DecidableEq α] (l : List α) (x : α)
    (h_nodup : l.Nodup)
    (h_last : l.getLast? = some x) :
    ∀ y : α, (x, y) ∉ l.zip l.tail

lemma drop_of_lt_length_nonempty {α : Type*} (path : List α) (i : ℕ)
    (h_i_lt : i < path.length) :
    path.drop i ≠ []

lemma cycle_from_suffix_contains (R : Run S) (suffix : List S) (current y : S)
    (h_suffix_nodup : suffix.Nodup)
    (h_contains_suffix : R.containsPath suffix)
    (h_suffix_nonempty : suffix ≠ [])
    (h_suffix_last : suffix.getLast? = some current)
    (h_edge : R (current, y) > 0) :
    ∀ t : Transition S, countTransitionInPath t (suffix ++ [y]) ≤ R t

lemma path_with_back_edge_creates_cycle (R : Run S) (path : List S) (current y : S)
    (h_acyclic : R.isAcyclic)
    (h_end : path.getLast? = some current)
    (h_contains : R.containsPath path)
    (h_y_in_path : y ∈ path)
    (h_edge : R (current, y) > 0) :
    R.hasCycle

lemma acyclic_edge_not_in_path (R : Run S) (path : List S) (current y : S)
    (h_acyclic : R.isAcyclic)
    (h_end : path.getLast? = some current)
    (h_contains : R.containsPath path)
    (h_edge : R (current, y) > 0)
    (h_y_in_path : y ∈ path) :
    False

lemma not_mem_implies_transition_not_in_zip_tail {α : Type*} (path : List α) (x y : α)
    (h_y_not_in : y ∉ path) :
    (x, y) ∉ path.zip path.tail

lemma containsPath_append_singleton (R : Run S) (path : List S) (x y : S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_contains : R.containsPath path)
    (h_y_not_in_path : y ∉ path)
    (h_edge : R (x, y) > 0) :
    R.containsPath (path ++ [y])

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Fintype.Basic

import Mathlib.Data.Fintype.Prod

import Mathlib.Data.List.Basic

import Mathlib.Algebra.BigOperators.Group.Finset.Basic

import Mathlib.Algebra.Order.BigOperators.Group.Finset

import Mathlib.Algebra.BigOperators.Group.Finset.Piecewise

import Mathlib.Algebra.BigOperators.Ring.Finset

namespace Utils.StateTransition

variable {S : Type*} [DecidableEq S] [Fintype S]

def Transition (S : Type*) := S × S

def Run (S : Type*) := Transition S → ℕ

def countTransitionInPath [DecidableEq S] (t : Transition S) (path : List S) : ℕ :=
  (path.zip path.tail).count t

def Run.containsPath [DecidableEq S] (R : Run S) (path : List S) : Prop :=
  ∀ t : Transition S, countTransitionInPath t path ≤ R t

def Run.hasCycle [DecidableEq S] (R : Run S) : Prop :=
  ∃ (cycle : List S), cycle.length ≥ 2 ∧
    cycle.head? = cycle.getLast? ∧
    R.containsPath cycle

def Run.isAcyclic [DecidableEq S] (R : Run S) : Prop :=
  ¬R.hasCycle

def Run.reachable [DecidableEq S] (R : Run S) (start finish : S) : Prop :=
  ∃ (path : List S), path.head? = some start ∧ path.getLast? = some finish ∧
    path ≠ [] ∧ R.containsPath path

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

lemma acyclic_has_leaf_aux (R : Run S) (root current : S)
    (path : List S)
    (h_acyclic : R.isAcyclic)
    (h_start : path.head? = some root)
    (h_end : path.getLast? = some current)
    (h_nonempty : path ≠ [])
    (h_contains : R.containsPath path)
    (h_has_out : ∃ y, y ∉ path ∧ R (current, y) > 0) :
    ∃ leaf, R.isLeaf root leaf :=

:= by
  -- Get a successor not in path
  obtain ⟨y, h_y_not_in_path, h_edge⟩ := h_has_out

  -- Check if y has any outgoing edges to states not in path ++ [y]
  by_cases h_y_has_out : ∃ z, z ∉ path ∧ z ≠ y ∧ R (y, z) > 0
  case neg =>
    -- y has no outgoing edges to states not in path ++ [y]
    -- We'll show y is actually a leaf (has NO outgoing edges at all)
    use y
    constructor
    · -- Show y is reachable from root
      -- Extend the path by adding y
      use path ++ [y]
      constructor
      · aesop
      constructor
      · simp
      constructor
      · aesop
      · exact containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge
    · -- Show y has no outgoing edges
      intro z
      by_contra h_pos
      push_neg at h_y_has_out
      -- If R(y,z) > 0, then by h_y_has_out, either z ∈ path or z = y
      have h_z_pos : R (y, z) > 0 := by omega
      have h_z_in_path_or_y : z ∈ path ∨ z = y := by grind

      -- Derive contradiction from cycle
      cases h_z_in_path_or_y with
      | inl h_z_in_path =>
        -- z ∈ path, so we can construct a cycle
        -- path ends with current, current → y, y → z, and z ∈ path
        -- This creates a cycle: (suffix of path from z to current) → y → z
        have h_z_in_extended : z ∈ path ++ [y] := by aesop
        exact acyclic_edge_not_in_path R (path ++ [y]) y z h_acyclic (by simp) (containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge) h_z_pos h_z_in_extended
      | inr h_z_eq_y =>
        -- z = y, so we have a self-loop y → y
        rw [h_z_eq_y] at h_z_pos
        exact acyclic_no_self_loop R y h_acyclic h_z_pos

  case pos =>
    -- y has an outgoing edge to some z not in path and z ≠ y
    -- Recurse with path ++ [y]
    obtain ⟨z, h_z_not_in_path, h_z_ne_y, h_y_z_edge⟩ := h_y_has_out

    -- Construct new path
    let new_path := path ++ [y]

    -- Show properties of new_path
    have h_new_start : new_path.head? = some root := by aesop
    have h_new_end : new_path.getLast? = some y := by aesop
    have h_new_nonempty : new_path ≠ [] := by aesop
    have h_new_contains : R.containsPath new_path :=
      containsPath_append_singleton R path current y h_nonempty h_end h_contains h_y_not_in_path h_edge

    have h_new_has_out : ∃ w, w ∉ new_path ∧ R (y, w) > 0 := by grind
    exact acyclic_has_leaf_aux R root y new_path h_acyclic h_new_start h_new_end h_new_nonempty h_new_contains h_new_has_out
termination_by Fintype.card S - path.toFinset.card
decreasing_by
  simp_wf
  -- new_path = path ++ [y], and y ∉ path, so toFinset increases by 1
  have h_y_not_mem : y ∉ path.toFinset := by
    simp
    exact h_y_not_in_path
  have h_card_increase : (insert y path.toFinset).card = path.toFinset.card + 1 := by
    rw [Finset.card_insert_of_notMem h_y_not_mem]
  -- path.toFinset is a strict subset of univ
  have h_path_subset := finset_ssubset_univ_of_not_mem path.toFinset y h_y_not_mem
  have h_card_bound : path.toFinset.card < Fintype.card S := by
    rw [← Finset.card_univ]
    exact Finset.card_lt_card h_path_subset
  rw [h_card_increase]
  omega


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.addressStrictTimestampSorted_of_AddressTimestampSorted_noTimestampDup
    (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted)
    (h_no_timestamp_dup : accesses.Notimestampdup) :
    accesses.isAddressStrictTimestampSorted

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)

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

:= by
  have h_strict := addressStrictTimestampSorted_of_AddressTimestampSorted_noTimestampDup accesses h_sorted h_nodup
  simp only [isAddressStrictTimestampSorted, filterAddress, isTimestampSorted] at h_strict ⊢
  induction accesses with
  | nil => simp
  | cons head tail ih =>
    obtain ⟨t, a, r, w⟩ := head
    simp only [List.filter_cons]
    split_ifs with h_addr
    · simp only [List.sorted_cons]
      constructor
      · intro z hz
        simp only [List.mem_filter] at hz
        obtain ⟨hz_mem, hz_addr⟩ := hz
        obtain ⟨t_z, a_z, r_z, w_z⟩ := z
        simp only [decide_eq_true_eq] at hz_addr
        simp only [List.sorted_cons] at h_strict
        have h_ord := h_strict.1 (t_z, a_z, r_z, w_z) hz_mem
        simp only [address_strict_timestamp_ordering] at h_ord
        simp only [decide_eq_true_eq] at h_addr
        rw [hz_addr, h_addr] at h_ord
        simp only [↓reduceIte] at h_ord
        simp only [timestamp_ordering]
        exact h_ord
      · apply ih
        · simp only [isAddressTimestampSorted] at h_sorted ⊢
          exact List.Sorted.of_cons h_sorted
        · simp only [Notimestampdup] at h_nodup ⊢
          exact List.Pairwise.of_cons h_nodup
        · exact List.Sorted.of_cons h_strict
    · apply ih
      · simp only [isAddressTimestampSorted] at h_sorted ⊢
        exact List.Sorted.of_cons h_sorted
      · simp only [Notimestampdup] at h_nodup ⊢
        exact List.Pairwise.of_cons h_nodup
      · exact List.Sorted.of_cons h_strict


MemoryAccess

def MemoryAccess := ℕ × ℕ × ℕ × ℕ

MemoryAccessList.isConsistentOffline

theorem MemoryAccessList.addressTimestampSort_sorted (accesses : MemoryAccessList) :
    (MemoryAccessList.addressTimestampSort accesses).Sorted address_timestamp_ordering

theorem MemoryAccessList.addressTimestampSort_perm (accesses : MemoryAccessList) :
    (MemoryAccessList.addressTimestampSort accesses).Perm accesses

theorem MemoryAccessList.noTimestampDup_perm (l1 l2 : MemoryAccessList)
    (h_l1_nodup : l1.Notimestampdup) (h_perm : l1.Perm l2) :
    l2.Notimestampdup

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.noTimestampDup_of_TimestampSorted
    (accesses : MemoryAccessList) (h_sorted : accesses.isTimestampSorted) :
    accesses.Notimestampdup

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

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

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

theorem MemoryAccessList.isConsistentSingleAddress_filterAddress_forall_of_cons
    (head : MemoryAccess) (tail : MemoryAccessList)
    (h_sorted : isAddressTimestampSorted (head :: tail))
    (h_nodup : Notimestampdup (head :: tail))
    (h : ∀ addr, isConsistentSingleAddress (filterAddress (head :: tail) addr)
         (filterAddress_sorted_from_addressTimestampSorted (head :: tail) h_sorted h_nodup addr)) :
    ∀ addr, isConsistentSingleAddress (filterAddress tail addr)
      (filterAddress_sorted_from_addressTimestampSorted tail
        (isAddressTimestampSorted_of_cons head tail h_sorted)
        (noTimestampDup_of_cons head tail h_nodup) addr)

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)

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)

theorem MemoryAccessList.isConsistentOffline_iff_all_single_addresses (accesses : MemoryAccessList) (h_sorted : accesses.isAddressTimestampSorted) (h_nodup : accesses.Notimestampdup) :
    MemoryAccessList.isConsistentOffline accesses h_sorted ↔
    ∀ addr, MemoryAccessList.isConsistentSingleAddress (MemoryAccessList.filterAddress accesses addr) (filterAddress_sorted_from_addressTimestampSorted accesses h_sorted h_nodup addr)

theorem MemoryAccessList.addressTimestampSort_noTimestampDup
    (accesses : MemoryAccessList)
    (h_nodup : accesses.Notimestampdup) :
    accesses.addressTimestampSort.Notimestampdup

theorem MemoryAccessList.filterAddress_addressTimestampSort_eq
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isTimestampSorted)
    (h_nodup : accesses.Notimestampdup)
    (addr : ℕ) :
    (accesses.filterAddress addr).isConsistentSingleAddress
      (filterAddress_sorted accesses h_sorted addr) ↔
    (accesses.addressTimestampSort.filterAddress addr).isConsistentSingleAddress
      (filterAddress_sorted_from_addressTimestampSorted accesses.addressTimestampSort
        (addressTimestampSort_sorted accesses)
        (addressTimestampSort_noTimestampDup accesses h_nodup) addr)

theorem MemoryAccessList.isConsistentOnline_iff_sorted_isConsistentOffline
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isTimestampSorted)
    (h_nodup : accesses.Notimestampdup) :
    MemoryAccessList.isConsistentOnline accesses h_sorted ↔
    MemoryAccessList.isConsistentOffline (MemoryAccessList.addressTimestampSort accesses) (MemoryAccessList.addressTimestampSort_sorted accesses)

lemma MemoryAccessList.eq_of_perm_of_sorted {l1 l2 l3 : MemoryAccessList} (h_l1_sorted: l1.isTimestampSorted)
    (h_l2_sorted : l2.isAddressTimestampSorted) (h_l3_sorted : l3.isAddressTimestampSorted)
    (h_perm1 : l1.Perm l2) (h_perm2 : l1.Perm l3) : l2 = l3

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 AddressSortedMemoryAccessList := {accesses : MemoryAccessList // accesses.isAddressTimestampSorted}

def MemoryAccessList.addressTimestampSort (accesses : MemoryAccessList) : MemoryAccessList :=
  List.insertionSort address_timestamp_ordering accesses

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)

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.isConsistentOnline_iff_isConsistentOffline
    (accesses : MemoryAccessList)
    (h_sorted : accesses.isTimestampSorted)
    (h_nodup : accesses.Notimestampdup) :
    MemoryAccessList.isConsistentOnline accesses h_sorted ↔
    ∃ permuted : AddressSortedMemoryAccessList,
      permuted.val.Perm accesses ∧
      MemoryAccessList.isConsistentOffline permuted.val permuted.property :=

:= by
  constructor
  · intro h
    use ⟨MemoryAccessList.addressTimestampSort accesses, MemoryAccessList.addressTimestampSort_sorted accesses⟩
    constructor
    · simp only
      apply MemoryAccessList.addressTimestampSort_perm
    · simp only
      have h' := MemoryAccessList.isConsistentOnline_iff_sorted_isConsistentOffline accesses h_sorted h_nodup
      rw [←h']
      assumption
  · rintro ⟨⟨permuted, h_permuted_sorted⟩, h_perm, h_offline⟩
    simp_all only
    rw [List.perm_comm] at h_perm
    have h_eq := MemoryAccessList.eq_of_perm_of_sorted h_sorted h_permuted_sorted (MemoryAccessList.addressTimestampSort_sorted accesses)
      h_perm (by rw [List.perm_comm]; apply MemoryAccessList.addressTimestampSort_perm)
    simp only [h_eq] at h_offline
    rw [←MemoryAccessList.isConsistentOnline_iff_sorted_isConsistentOffline accesses h_sorted h_nodup] at h_offline
    assumption


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

Utils.StateTransition.Transition

def Transition (S : Type*) := S × S

Utils.StateTransition.Run

def Run (S : Type*) := Transition S → ℕ

Utils.StateTransition.Run.isAcyclic

lemma zip_drop_sublist (l : List S) (n : ℕ) :
    ((l.drop n).zip (l.drop (n + 1))).Sublist (l.zip l.tail)

lemma containsPath_drop (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.drop n)

lemma zip_take_sublist (l1 l2 : List S) (n m : ℕ) :
    ((l1.take n).zip (l2.take m)).Sublist (l1.zip l2)

lemma tail_take {α : Type*} (l : List α) (n : ℕ) :
    (l.take n).tail = (l.tail).take (n - 1)

lemma containsPath_take (R : Run S) (path : List S) (n : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath (path.take n)

lemma drop_take_length_ge_two {α : Type*} (path : List α) (n m : Fin path.length)
    (h_n_lt_m : n < m) :
    ((path.drop n.val).take (m.val - n.val + 1)).length ≥ 2

lemma getLast_drop_take {α : Type*} (path : List α) (n k : ℕ)
    (h_n_lt : n < path.length)
    (h_bound : n + k ≤ path.length)
    (h_k_pos : k > 0) :
    ((path.drop n).take k).getLast? = path[n + k - 1]?

lemma drop_take_cycle_same_endpoints (path : List S) (x : S) (n m : Fin path.length)
    (h_n_lt_m : n < m)
    (h_x_at_n : path[n] = x)
    (h_x_at_m : path[m] = x) :
    ((path.drop n.val).take (m.val - n.val + 1)).head? =
    ((path.drop n.val).take (m.val - n.val + 1)).getLast?

lemma containsPath_drop_take (R : Run S) (path : List S) (n m : ℕ)
    (h_contains : R.containsPath path) :
    R.containsPath ((path.drop n).take m)

lemma acyclic_containsPath_nodup (R : Run S) (path : List S)
    (h_acyclic : R.isAcyclic)
    (h_contains : R.containsPath path) :
    path.Nodup

lemma countTransitionInPath_append_singleton (path : List S) (x y : S)
    (h_nonempty : path ≠ [])
    (h_last : path.getLast? = some x)
    (h_not_in : (x, y) ∉ path.zip path.tail) :
    countTransitionInPath (x, y) (path ++ [y]) = 1

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

lemma getLast_mem {α : Type*} (l : List α) (x : α) (h_last : l.getLast? = some x) :
    x ∈ l

lemma last_not_in_zip_tail {α : Type*} [DecidableEq α] (l : List α) (x : α)
    (h_nodup : l.Nodup)
    (h_last : l.getLast? = some x) :
    ∀ y : α, (x, y) ∉ l.zip l.tail

lemma drop_of_lt_length_nonempty {α : Type*} (path : List α) (i : ℕ)
    (h_i_lt : i < path.length) :
    path.drop i ≠ []

lemma cycle_from_suffix_contains (R : Run S) (suffix : List S) (current y : S)
    (h_suffix_nodup : suffix.Nodup)
    (h_contains_suffix : R.containsPath suffix)
    (h_suffix_nonempty : suffix ≠ [])
    (h_suffix_last : suffix.getLast? = some current)
    (h_edge : R (current, y) > 0) :
    ∀ t : Transition S, countTransitionInPath t (suffix ++ [y]) ≤ R t

import Mathlib.Data.Finset.Basic

import Mathlib.Data.Fintype.Basic

import Mathlib.Data.Fintype.Prod

import Mathlib.Data.List.Basic

import Mathlib.Algebra.BigOperators.Group.Finset.Basic

import Mathlib.Algebra.Order.BigOperators.Group.Finset

import Mathlib.Algebra.BigOperators.Group.Finset.Piecewise

import Mathlib.Algebra.BigOperators.Ring.Finset

namespace Utils.StateTransition

variable {S : Type*} [DecidableEq S] [Fintype S]

def Transition (S : Type*) := S × S

def Run (S : Type*) := Transition S → ℕ

def countTransitionInPath [DecidableEq S] (t : Transition S) (path : List S) : ℕ :=
  (path.zip path.tail).count t

def Run.containsPath [DecidableEq S] (R : Run S) (path : List S) : Prop :=
  ∀ t : Transition S, countTransitionInPath t path ≤ R t

def Run.hasCycle [DecidableEq S] (R : Run S) : Prop :=
  ∃ (cycle : List S), cycle.length ≥ 2 ∧
    cycle.head? = cycle.getLast? ∧
    R.containsPath cycle

def Run.isAcyclic [DecidableEq S] (R : Run S) : Prop :=
  ¬R.hasCycle

lemma path_with_back_edge_creates_cycle (R : Run S) (path : List S) (current y : S)
    (h_acyclic : R.isAcyclic)
    (h_end : path.getLast? = some current)
    (h_contains : R.containsPath path)
    (h_y_in_path : y ∈ path)
    (h_edge : R (current, y) > 0) :
    R.hasCycle :=

:= by
  -- Find y's position in path
  rw [List.mem_iff_getElem] at h_y_in_path
  obtain ⟨i, h_i_lt, h_y_eq⟩ := h_y_in_path
  let suffix := path.drop i
  have h_suffix_nonempty : suffix ≠ [] := drop_of_lt_length_nonempty path i h_i_lt
  have h_suffix_head : suffix.head? = some y := by aesop
  have h_suffix_last : suffix.getLast? = some current := by grind
  let cycle := suffix ++ [y]
  have h_cycle_head : cycle.head? = some y := by grind
  have h_cycle_last : cycle.getLast? = some y := by grind
  have h_cycle_len : cycle.length ≥ 2 := by grind
  use cycle
  constructor
  · exact h_cycle_len
  constructor
  · rw [h_cycle_head, h_cycle_last]
  · unfold cycle
    have h_suffix_contains : R.containsPath suffix := by
      unfold suffix
      exact containsPath_drop R path i h_contains
    have h_path_nodup := acyclic_containsPath_nodup R path h_acyclic h_contains
    have h_suffix_nodup : suffix.Nodup := by grind
    exact cycle_from_suffix_contains R suffix current y h_suffix_nodup h_suffix_contains h_suffix_nonempty h_suffix_last h_edge


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

Vector.toChunks

theorem toList_cons {a : α} {v : Vector α n} : (cons a v).toList = a :: v.toList

import Mathlib.Analysis.Normed.Ring.Lemmas

import Mathlib.Combinatorics.Enumerative.Composition

import Init.Data.List.Find

variable {α β : Type} {n m : ℕ}

open Vector (finRange)

namespace Vector

def cons (a : α) (v : Vector α n) : Vector α (n + 1) :=
  ⟨ .mk (a :: v.toList), by admit /- proof elided -/
  ⟩

def induct {motive : {n : ℕ} → Vector α n → Sort u}
    (nil : motive #v[])
    (cons: ∀ {n : ℕ} (a : α) (as : Vector α n), motive as → motive (cons a as))
    {n : ℕ} (v : Vector α n) : motive v :=
  match v with
  | ⟨ .mk [], h ⟩ => by admit /- proof elided -/
  | ⟨ .mk (a :: as), h ⟩ => by admit /- proof elided -/

end Vector

def Composition.ofProductLength (m : ℕ+) {α : Type} {l : List α} (hl : l.length = n * m.val) : Composition l.length := {
  blocks := List.replicate n m.val
  blocks_pos hi := (List.mem_replicate.mp hi).right ▸ m.pos
  blocks_sum := hl ▸ List.sum_replicate_nat
}

namespace Vector

def toChunks (m : ℕ+) {α : Type} (v : Vector α (n*m)) : Vector (Vector α m) n :=
  let comp := Composition.ofProductLength m v.length_toList
  let list : List (Vector α m) := v.toList.splitWrtComposition comp
    |>.attachWith (List.length · = m) (comp.ofProductLength_mem_length rfl)
    |>.map fun ⟨ l, hl ⟩ => .mk (.mk l) hl
  .mk (.mk list) (by admit /- proof elided -/
  )

theorem flatten_toChunks {α : Type} (m : ℕ+) (v : Vector (Vector α m) n) :
    v.flatten.toChunks m = v :=

:= by
  simp only [toChunks]
  rw [←Vector.toArray_inj,←Array.toList_inj]
  simp only
  let v_list_list := v.toList.map (Array.toList ∘ toArray)
  have h_flatten : v.flatten.toList = v_list_list.flatten := by
    rw [Vector.flatten_mk, Vector.toList_mk, Array.toList_flatten, Array.toList_map, List.map_map]
    congr
  have h_length : v.flatten.toList.length = n * ↑m := by rw [length_toList]
  have h_flatten_length : v_list_list.flatten.length = n * ↑m := by rw [←h_flatten, h_length]
  have h' : (v.flatten.toList.splitWrtComposition (Composition.ofProductLength m h_length)) = v_list_list := by
    rw [← v_list_list.splitWrtComposition_flatten (Composition.ofProductLength m h_flatten_length)]
    congr 1
    · rw [h_length, h_flatten_length]
    congr
    · simp [h_flatten]
    simp only [List.map_map, Composition.ofProductLength, v_list_list]
    clear *-
    induction v using Vector.induct
    case nil => rfl
    case cons xs x hi => rw [List.replicate_succ, Vector.toList_cons, List.map_cons, hi,
      Function.comp_apply, Function.comp_apply, Array.length_toList, size_toArray]
  simp_all only [List.length_flatten, List.map_map, List.map_attachWith, v_list_list]
  rw [List.map_attach_eq_pmap, List.pmap_map]
  simp only [Function.comp_apply, Array.toArray_toList, mk_toArray, List.pmap_eq_map,
    List.map_id_fun', id_eq]
  congr


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

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

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

syntax "infer_constant_length" : tactic

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

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

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

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

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

CellOffset

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

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

F

def F p := ZMod p

TwoRowsConstraint

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

TableConstraint

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

TableContext

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

CellAssignment

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

  vars : Vector (Cell W S) offset

offset

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

ConstantLength

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

localLength

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

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

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

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

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

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

Environment

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

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

Lookup

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

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

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

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

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)

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

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

Operations

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

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

localLength

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

Cell

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

assign

@[table_norm, table_assignment_norm]
def assign (off : CellOffset W S) : Expression F → TableConstraint W S F Unit

  | .var v => assignVar off v

  | x => do
    let new_var ← witnessVar x.eval
    assertZero (x - var new_var)
    assignVar off new_var

assignVar

@[table_norm, table_assignment_norm]
def assignVar (off : CellOffset W S) (v : Variable F) : TableConstraint W S F Unit :=
  modify fun ctx =>
    let assignment := ctx.assignment.setVarInput off v.index
    { ctx with assignment }

setVarInput

@[table_assignment_norm]
def setVarInput (assignment : CellAssignment W S) (off : CellOffset W S) (var : ℕ) : CellAssignment W S :=
  let vars := assignment.vars.set? var (.input off)

  { assignment with vars }

set?

def set? (v : Vector α n) (i : ℕ) (a : α) : Vector α n :=
  ⟨ .mk <| v.toList.set i a, by admit /- proof elided -/
  ⟩

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

witnessVar

@[circuit_norm]
def witnessVar (compute : Environment F → F) : Circuit F (Variable F) :=
  fun (offset : ℕ) =>
    let var : Variable F := ⟨ offset ⟩
    (var, [.witness 1 fun env => #v[compute env]])

eval

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

ProvableStruct

class ProvableStruct (α : TypeMap) where
  components : List WithProvableType
  toComponents {F : Type} : α F → ProvableTypeList F components
  fromComponents {F : Type} : ProvableTypeList F components → α F

  combinedSize : ℕ := combinedSize' components
  combinedSize_eq : combinedSize = combinedSize' components := by admit /- proof elided -/

assertZero

@[circuit_norm]
def assertZero (e : Expression F) : Circuit F Unit := fun _ =>
  ((), [.assert e])

getCurrRow

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

getRow

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

Row.get

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

pushRow

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

mapFinRange

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

get

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

TraceOfLength

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

Trace

inductive Trace (F : Type) (S : Type → Type) [ProvableType S] where

  | empty : Trace F S

  | cons (rest : Trace F S) (row : Row F S) : Trace F S

empty

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

Row

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

mapRange

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

varFromOffset

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

HasAssignEq

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

getNextRow

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

HasAssertEq

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

circuit

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

Assumptions

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

Inputs

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

Normalized

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

Spec

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

value

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

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

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

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

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

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

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)

map

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

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

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

RowType

structure RowType (F : Type) where
  x: F
  y: F

Inputs

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

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

windowEnv

def windowEnv (table : TableConstraint W S F Unit)
  (window : TraceOfLength F S W) (aux_env : Environment F) : Environment F :=
  let assignment := table.finalAssignment
  .mk fun i =>
    if hi : i < assignment.offset then
      match assignment.vars[i] with
      | .input ⟨i, j⟩ => window.get i j
      | .aux k => aux_env.get k
    else aux_env.get (i + assignment.aux_length)

finalAssignment

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

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

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

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

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

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

theorem mapRange_zero {create : ℕ → α} : mapRange 0 create = #v[]

theorem mapFinRange_succ {n : ℕ} {create : Fin (n + 1) → α} :
    mapFinRange (n + 1) create = (mapFinRange n (fun i => create i.castSucc)).push (create (.last n))

theorem mapRange_succ {n} {create : ℕ → α} :
    mapRange (n + 1) create = (mapRange n create).push (create n)

theorem mapFinRange_zero {create : Fin 0 → α} : mapFinRange 0 create = #v[]

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

@[reducible]
def nextRowOff : RowType (CellOffset 2 RowType) := {
  x := ⟨.next 0, .next 1, .next 2, .next 3⟩,
  y := ⟨.next 4, .next 5, .next 6, .next 7⟩
}

lemma fib_assignment : (recursiveRelation (p:=p)).finalAssignment.vars =
   #v[.input ⟨0, 0⟩, .input ⟨0, 1⟩, .input ⟨0, 2⟩, .input ⟨0, 3⟩, .input ⟨0, 4⟩, .input ⟨0, 5⟩, .input ⟨0, 6⟩,
      .input ⟨0, 7⟩, .input ⟨1, 0⟩, .input ⟨1, 1⟩, .input ⟨1, 2⟩, .input ⟨1, 3⟩, .input ⟨1, 4⟩, .input ⟨1, 5⟩,
      .input ⟨1, 6⟩, .input ⟨1, 7⟩, .input ⟨1, 4⟩, .aux 1, .input ⟨1, 5⟩, .aux 3, .input ⟨1, 6⟩, .aux 5,
      .input ⟨1, 7⟩, .aux 7]

lemma fib_vars (curr next : Row (F p) RowType) (aux_env : Environment (F p)) :
    let env := recursiveRelation.windowEnv ⟨<+> +> curr +> next, rfl⟩ aux_env;
    eval env (varFromOffset U32 0) = curr.x ∧
    eval env (varFromOffset U32 4) = curr.y ∧
    eval env (varFromOffset U32 8) = next.x ∧
    eval env (U32.mk (var ⟨16⟩) (var ⟨18⟩) (var ⟨20⟩) (var ⟨22⟩)) = next.y

import Clean.Utils.Vector

import Clean.Circuit.Basic

import Clean.Table.Theorems

import Clean.Gadgets.Addition32.Addition32

import Clean.Gadgets.Equality

import Clean.Types.U32

namespace Tables.Fibonacci32

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

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

@[reducible]
def nextRowOff : RowType (CellOffset 2 RowType) := {
  x := ⟨.next 0, .next 1, .next 2, .next 3⟩,
  y := ⟨.next 4, .next 5, .next 6, .next 7⟩
}

def assignU32 (offs : U32 (CellOffset 2 RowType)) (x : Var U32 (F p)) : TwoRowsConstraint RowType (F p) := do
  assign offs.x0 x.x0
  assign offs.x1 x.x1
  assign offs.x2 x.x2
  assign offs.x3 x.x3

def recursiveRelation : TwoRowsConstraint RowType (F p) := do
  let curr ← TableConstraint.getCurrRow
  let next ← TableConstraint.getNextRow

  let z ← Gadgets.Addition32.circuit { x := curr.x, y := curr.y }

  assignU32 nextRowOff.y z
  curr.y === next.x

variable {α : Type}

lemma fib_constraints (curr next : Row (F p) RowType) (aux_env : Environment (F p))
  : recursiveRelation.ConstraintsHoldOnWindow ⟨<+> +> curr +> next, rfl⟩ aux_env →
  curr.y = next.x ∧
  (curr.x.Normalized → curr.y.Normalized → next.y.value = (curr.x.value + curr.y.value) % 2^32 ∧ next.y.Normalized) :=

:= by
  simp only [table_norm]
  obtain ⟨ hcurr_x, hcurr_y, hnext_x, hnext_y ⟩ := fib_vars curr next aux_env
  set env := recursiveRelation.windowEnv  ⟨<+> +> curr +> next, rfl⟩ aux_env
  simp only [table_norm, circuit_norm, recursiveRelation,
    assignU32, Gadgets.Addition32.circuit]
  rintro ⟨ h_add, h_eq ⟩
  simp only [table_norm, circuit_norm, Nat.reduceAdd, zero_add] at h_add
  simp only [circuit_norm] at hnext_y
  rw [hcurr_x, hcurr_y, hnext_y] at h_add
  rw [hcurr_y, hnext_x] at h_eq
  clear hcurr_x hcurr_y hnext_x hnext_y
  constructor
  · exact h_eq
  rw [Gadgets.Addition32.Assumptions, Gadgets.Addition32.Spec] at h_add
  intro h_norm_x h_norm_y
  specialize h_add ⟨ h_norm_x, h_norm_y ⟩
  obtain ⟨ h_add_mod, h_norm_next_y ⟩ := h_add
  exact ⟨h_add_mod, h_norm_next_y⟩