APSP in $O(n^{2.995561})$: definitions, part 2 of 2

/-
Copyright (c) 2026 Anthropic, PBC. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
SPDX-License-Identifier: Apache-2.0
-/
-- Generated by skeleton subtraction from the ThreeSumApsp library at commit 3b45b374640a.

import Definitions.Def_APSP2995561_Defs1

set_option autoImplicit false
set_option relaxedAutoImplicit false

section
-- Source module: ThreeSumApsp.Batch.OWrite
/-!
# The batched counting oracle: writeSrc

`writeSrc_ok`: the procedure writeSrc (82) meets `WriteSrcSpec`.

It clears the source, then runs through the tuples `(u, a, l, cl)` (the product, the two digits of
the row `R = a k + l`, the vertex of the middle block) and marks one cell for each.  `Cells` says that
the cells of the source hold 1 exactly at the marks of the tuples handled so far (`marks`, with the
lexicographic order `lexBefore`), and the end identifies the marks of all tuples with `srcVal = 1`.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Batch Light.Sec4

/-! ## The tuples and their cells -/

/-- The length of the middle block from `c`: `min s (n - c)`, as the program computes it. -/
def wsLen (n s c : ℕ) : ℕ := if c + s ≤ n then s else n - c

/-- The residue of the tuple `(u, a, l, cl)`. -/
def wsF (n p s j0 : ℕ) (T1 T2 : List ℕ) (u a l cl : ℕ) : ℕ :=
  (T1.getD (a * n + (j0 + u) * s + cl) 0 + T2.getD l 0) % p

/-- The cell (from base) that the tuple `(u, a, l, cl)` marks. -/
def wsPos (n k p s j0 D KK N0 : ℕ) (T1 T2 : List ℕ) (u a l cl : ℕ) : ℕ :=
  srcOf KK N0 u (a * k + l) * D + (cl * p + wsF n p s j0 T1 T2 u a l cl)

/-- The tuples that the loops run through. -/
def WsOk (n k s j0 KKb : ℕ) (u a l cl : ℕ) : Prop :=
  u < KKb ∧ a < n ∧ l < k ∧ cl < wsLen n s ((j0 + u) * s)

/-- The tuple `(u', a', l', cl')` comes before `(u, a, l, cl)`. -/
def lexBefore (u a l cl : ℕ) (u' a' l' cl' : ℕ) : Prop :=
  u' < u ∨ (u' = u ∧ (a' < a ∨ (a' = a ∧ (l' < l ∨ (l' = l ∧ cl' < cl)))))

/-- The cell `i` is the cell of a tuple of `ok` that satisfies `dn`. -/
def marks (ok : ℕ → ℕ → ℕ → ℕ → Prop) (pos : ℕ → ℕ → ℕ → ℕ → ℕ)
    (dn : ℕ → ℕ → ℕ → ℕ → Prop) (i : ℕ) : Prop :=
  ∃ u a l cl, ok u a l cl ∧ dn u a l cl ∧ pos u a l cl = i

section marksLemmas

variable {ok : ℕ → ℕ → ℕ → ℕ → Prop} {pos : ℕ → ℕ → ℕ → ℕ → ℕ}

end marksLemmas

/-! ## The cells of the source -/

/-- The `sz` cells from base hold 1 at the cells of `S` and 0 at the others; no other cell has
changed. -/
def Cells (μ μ' : ℕ → ℤ) (base sz : ℕ) (S : ℕ → Prop) : Prop :=
  (∀ i < sz, (S i → μ' (base + i) = 1) ∧ (¬ S i → μ' (base + i) = 0)) ∧
    SameOutside μ μ' base sz

/-! ## The end: the marks are the entries 1 of the source -/

/-! ## The locals of writeSrc -/

/-- The arguments of writeSrc stand in the locals 0 to 12. -/
def WsArgs (loc : ℕ → ℤ) (base t1 t2 rs n k p D N0 j0 KKb s sz : ℕ) : Prop :=
  loc 0 = base ∧ loc 1 = t1 ∧ loc 2 = t2 ∧ loc 3 = rs ∧ loc 4 = n ∧ loc 5 = k ∧ loc 6 = p ∧
    loc 7 = D ∧ loc 8 = N0 ∧ loc 9 = j0 ∧ loc 10 = KKb ∧ loc 11 = s ∧ loc 12 = sz

section prog

variable {lim : Limits} {P : Program} {d : ℕ} {μ : ℕ → ℤ}
  {base t1 t2 rs n kd p D N0 j0 KKb s KK NP : ℕ} {T1 T2 : List ℕ}

end prog

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.OSetup
/-!
# The batched counting oracle: the layout, the locals and the need

From the free pointer `fr` on, the oracle lays out the tables `RAC` (`n²`), `RNBC` (`n²`), `RNAB`
(`n²`), `NL` (`k`), `HK` (`k`), `RS` (`NP`), `RV` (`NP`), `DK` (`k²`), `ZH` (`n²`), `ZL` (`n²`), the
two sources X and Y (`KK NP D` each), and the data structure of Theorem 30 at `B0`.

`oLocP` are the locals after the parameters, `oLocA` after the addresses.  The need `oNeedAt`:
the word bound `oWordSum`, the cells up to the top of the data structure, the depth `L + 8`.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp Light.Sec4 Light.Filter2

/-! ## The addresses, relative to the free pointer -/

def oaRNBC (n fr : ℕ) : ℕ := fr + n * n
def oaRNAB (n fr : ℕ) : ℕ := oaRNBC n fr + n * n
def oaNL (n fr : ℕ) : ℕ := oaRNAB n fr + n * n
def oaHK (n p fr : ℕ) : ℕ := oaNL n fr + oDig p
def oaRS (n p fr : ℕ) : ℕ := oaHK n p fr + oDig p
def oaRV (n s p fr : ℕ) : ℕ := oaRS n p fr + oNP n s p
def oaDK (n s p fr : ℕ) : ℕ := oaRV n s p fr + oNP n s p
def oaZH (n s p fr : ℕ) : ℕ := oaDK n s p fr + oDig p * oDig p
def oaZL (n s p fr : ℕ) : ℕ := oaZH n s p fr + n * n
/-- The cells of a source. -/
def oSz (n s p : ℕ) : ℕ := oKK s p * oNP n s p * oD s p
def oaX (n s p fr : ℕ) : ℕ := oaZL n s p fr + n * n
def oaY (n s p fr : ℕ) : ℕ := oaX n s p fr + oSz n s p
/-- The base of the data structure. -/
def oB0 (n s p fr : ℕ) : ℕ := oaY n s p fr + oSz n s p

/-! ## The locals -/

/-- The locals after the parameters (`oParams`): the arguments, then `m, D, L, t, N0, K0, KK, k, nb,
NP, Nv, T, B` at 9 to 21, zeros, and `binom(L, m)` at 41. -/
def oLocP (x : Light.Filter2.OracleInst) (fr : ℕ) : List ℤ :=
  ([x.n, x.U, x.ab, x.bc, x.ac, x.s, x.p, x.cnt, fr] : List ℕ).map (fun z : ℕ => (z : ℤ)) ++
  ([oM x.s x.p, oD x.s x.p, oL x.s x.p, oT x.s x.p, oN0 x.s x.p, oK0 x.s x.p, oKK x.s x.p,
    oDig x.p, oNb x.n x.s x.p, oNP x.n x.s x.p, oNv x.n x.s x.p, oProd x.n x.s,
    oBat x.n x.s x.p] : List ℕ).map (fun z : ℕ => (z : ℤ)) ++
  List.replicate 19 (0 : ℤ) ++ [(((oL x.s x.p).choose (oM x.s x.p) : ℕ) : ℤ)]

/-- The locals after the addresses (`oAddrs`), with the temporary, the batch, its first product and
its number of products. -/
def oLocA (x : Light.Filter2.OracleInst) (fr : ℕ) (tmp bn j0 kkb : ℤ) : List ℤ :=
  ([x.n, x.U, x.ab, x.bc, x.ac, x.s, x.p, x.cnt, fr] : List ℕ).map (fun z : ℕ => (z : ℤ)) ++
  ([oM x.s x.p, oD x.s x.p, oL x.s x.p, oT x.s x.p, oN0 x.s x.p, oK0 x.s x.p, oKK x.s x.p,
    oDig x.p, oNb x.n x.s x.p, oNP x.n x.s x.p, oNv x.n x.s x.p, oProd x.n x.s,
    oBat x.n x.s x.p] : List ℕ).map (fun z : ℕ => (z : ℤ)) ++
  ([fr, oaRNBC x.n fr, oaRNAB x.n fr, oaNL x.n fr, oaHK x.n x.p fr, oaRS x.n x.p fr,
    oaRV x.n x.s x.p fr, oaDK x.n x.s x.p fr, oaZH x.n x.s x.p fr, oaZL x.n x.s x.p fr,
    oaX x.n x.s x.p fr, oaY x.n x.s x.p fr, oB0 x.n x.s x.p fr] : List ℕ).map
      (fun z : ℕ => (z : ℤ)) ++
  [tmp, bn, j0, kkb] ++
  ([oKK x.s x.p * oNP x.n x.s x.p, oSz x.n x.s x.p, (oL x.s x.p).choose (oM x.s x.p)] :
    List ℕ).map (fun z : ℕ => (z : ℤ))

/-! ## The need -/

/-- Everything that the oracle forms or passes on is at most this. -/
def oWordSum (n U s p : ℕ) : ℕ :=
  100 + 10 ^ oM s p * (7 ^ oL s p * 7 ^ oL s p) + 7 ^ (oL s p + 1) + 10 ^ (oL s p + 1)
    + 2 ^ oL s p + 4 * (oL s p).choose (oM s p) + 4 + 4 * s * p + oD s p + 56569 * oM s p + 500000
    + 3 ^ (oL s p - oM s p) + oL s p + 1 + 4 * U + 16 * p + 16 + n * oDig p + oN0 s p
    + oNb n s p + oNP n s p + oNv n s p + n + s + oProd n s + oKK s p + oBat n s p
    + (oNb n s p + 1) * (oKK s p * oN0 s p + oN0 s p + 1) + (oProd n s + oKK s p + 1) * (s + 1)
    + n * n + 2 * p + 4 + oSz n s p + oK0 s p * oN0 s p

/-- The cells from the free pointer on: the tables, the sources, the data structure, and room for
the table of binomials and the table of emod. -/
def oCells (n U s p : ℕ) : ℕ :=
  oB0 n s p 0 + Light.Sec4.top (oPar n s p) (oT s p) 0 + oL s p + Nat.log 2 (U + 3 * p + 1) + 10

/-- The need of the oracle, for `n`, `U`, `s`, `p`. -/
def oNeedAt (n U s p : ℕ) : Need := ⟨oWordSum n U s p, oCells n U s p, oL s p + 8⟩

/-- **The need of the oracle** as a function of the parameters `[n, U, s, p]`. -/
def oNeed (ps : List ℕ) : Need := Light.Batch.oNeedAt (ps.getD 0 0) (ps.getD 1 0) (ps.getD 2 0) (ps.getD 3 0)

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Filter2.Asymp
/-!
# Asymptotics of the filtering parameters

With `k = Nat.size n`, `ℓ = ⌊√k⌋ + 1`, `m = ⌊√ℓ⌋` (so `b = ℓ (m + 1)`), every parameter is a power
of two whose exponent is a polynomial in `m`, `ℓ`, `k`, and `m → ∞` with `n`.

* `rpow_le_two_mul_sOf`, `sOf_le_two_mul_rpow`: `n^σ / 2 ≤ s ≤ 2 n^σ`, `σ = 0.0336175`.
* `RR_le_rpow`, `QOf_le_rpow`, `calls_le_rpow`, `PmaxOf_le`: `R, Q, rounds · (r (Q + 1) + 2)` are
  `n^{o(1)}`, and `Pmax ≤ s n^{o(1)}`.
* `guard_eventually`: eventually `Guard n U` for all `U ≤ c n^κ`.
-/

@[expose] public section

open Filter

namespace Light.Filter2

/-! ## Basic facts -/

/-- `m = ⌊√ℓ⌋`. -/
def mmOf (n : ℕ) : ℕ := Nat.sqrt (Light.Filter2.ellOf n)

/-- The exponent of `s`: `s = 2^{jOf n}`. -/
def jOf (n : ℕ) : ℕ := 336175 * Light.Filter2.kOf n / 10 ^ 7

/-! ## Polynomial inequalities, eventually -/

/-! ## The parameters are `n^{o(1)}` -/

/-! ## The block size and the number of blocks -/

/-! ## The guard -/

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Batch.OBound
/-!
# An eventual bound on the time of the batched oracle

`oTimeF_le`: eventually, for `U ≤ c n^κ` and `p ≤ Pmax`, the batched oracle at `s = sOf n` takes at
most `n^{2.9955609}` steps.
-/

@[expose] public section

open Filter

namespace Light.Batch

open ThreeSumApsp Light.Sec4

/-! ## The parameters, in the naturals -/

section nat

variable {n s p : ℕ}

end nat

/-! ## Exponents -/

section real

/-- The constraints on `ν = ln n`, `x = ln s`, `y = ln p`, `d = ln D`. -/
structure Ctx (ν x y d : ℝ) : Prop where
  ν_ge : 20000 ≤ ν
  x_nonneg : 0 ≤ x
  x_ge : 0.0336175 * ν ≤ Real.log 2 + x
  x_le : x ≤ Real.log 2 + 0.0336175 * ν
  y_nonneg : 0 ≤ y
  y_le : y ≤ x + 0.000000001 * ν
  d_ge : 2 * Real.log 2 + x + y ≤ d
  d_le : d ≤ 4 * Real.log 2 + x + y

namespace Ctx

variable {ν x y d : ℝ} (h : Ctx ν x y d)
include h

end Ctx

end real

/-! ## The parameters, in the reals -/

section bridge

end bridge

/-! ## The time, for one `n`, `s`, `p` -/

section assembly

/-- The bookkeeping part of the time (without the constant `A`). -/
def oBook (n U s p : ℕ) : ℕ :=
  n * n * (Nat.log 2 (U + 1) + 2) + p + oDig p * oDig p + oNP n s p + n * oDig p
    + (oL s p).choose (oM s p) + 4 ^ oL s p
    + oBat n s p * (oKK s p * oNP n s p * oD s p + oKK s p * n * oDig p * (s + 1)
        + oKK s p * n * n + 1) + 1

/-- The constant `ct` of the binomial tail. -/
noncomputable def ctB : ℝ := 1 / (1 - rhoC 20) * (2 / (1 - 56569 / 500000))

/-- The constant of `oTimeF_le_of`. -/
noncomputable def Kt (A C1 C2 Cq a0 κ' : ℝ) : ℝ :=
  (A * (a0 + κ' + 65) + C1 * (34560 * ctB + 76 * 72) + C2 * Cq * 12) * 10 ^ 18 * Real.exp 100

end assembly

/-! ## The eventual bound -/

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.ORunBatch
/-!
# The batched counting oracle: one batch

`oBatch` for the batch `bt`: the first product `j0 = bt KK`, the number `KKb` of products, the two
sources (writeSrc), the data structure (preCore), the queries.  Afterwards the `KKb n²` counts of
the batch stand at `cnt + j0 n²` (`oBatch_ends`).  The tables that the batch reads are `BTabs`.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Batch Light.Sec4 Light.Filter2 OS

/-- The tables that a batch reads. -/
structure BTabs (x : Light.Filter2.OracleInst) (fr : ℕ) (μ : ℕ → ℤ) : Prop where
  rac : SegN μ fr (x.AC.map fun w => resid x.p (1 * w))
  rnbc : SegN μ (oaRNBC x.n fr) (x.BC.map fun w => resid x.p (-1 * w))
  nl : SegN μ (oaNL x.n fr) ((List.range (oDig x.p)).map fun h : ℕ => resid x.p ((x.p : ℤ) - h))
  hk : SegN μ (oaHK x.n x.p fr)
    ((List.range (oDig x.p)).map fun h : ℕ => resid x.p ((h * oDig x.p : ℕ) : ℤ))
  rs : SegN μ (oaRS x.n x.p fr) ((List.range (oNP x.n x.s x.p)).map fun R =>
    R / oN0 x.s x.p * (oKK x.s x.p * oN0 x.s x.p) + R % oN0 x.s x.p)
  rv : SegN μ (oaRV x.n x.s x.p fr) ((List.range (oNP x.n x.s x.p)).map fun R =>
    R / oN0 x.s x.p * (oK0 x.s x.p * oN0 x.s x.p) + R % oN0 x.s x.p)
  zh : SegN μ (oaZH x.n x.s x.p fr) ((x.AB.map fun w => resid x.p (-1 * w)).map (· / oDig x.p))
  zl : SegN μ (oaZL x.n x.s x.p fr) ((x.AB.map fun w => resid x.p (-1 * w)).map (· % oDig x.p))

/-- The four calls of a batch. -/
def oBatchTail : Stmt :=
  .call pWriteSrc [Light.v Light.Batch.OS.aX, Light.v aRAC, Light.v aNL, Light.v aRS, Light.v Nn, Light.v Kd, Light.v Pr, Light.v Dd, Light.v N0, Light.v J0, Light.v KKb, Light.v Light.Batch.OS.S,
    Light.v Sz] Tmp ;;
  .call pWriteSrc [Light.v Light.Batch.OS.aY, Light.v aRNBC, Light.v aHK, Light.v aRS, Light.v Nn, Light.v Kd, Light.v Pr, Light.v Dd, Light.v N0, Light.v J0, Light.v KKb, Light.v Light.Batch.OS.S,
    Light.v Sz] Tmp ;;
  .call Proc.preCore [Light.v Ll, Light.v Mm, Light.v Tt, Light.v Nv, Light.v Dd, Light.v Light.Batch.OS.aX, Light.v Light.Batch.OS.aY, Light.v B0] Tmp ;;
  .call pQueries [Light.v B0, Light.v aZL, Light.v aZH, Light.v aRV, Light.v Nn, Light.v Kd, Light.v N0, Light.v K0, Light.v KKb, Light.v Light.Batch.OS.Cnt, Light.v J0] Tmp

/-- The time of the calls of a batch. -/
def tOBTail (n s p : ℕ) : ℕ :=
  2 * tWriteSrc (oSz n s p) (oKK s p) n (oDig p) s + tPreCoreP cShared30 (oPar n s p) (oT s p)
    + tQueries (oKK s p) n (tQueryAt (oL s p) (oM s p) (oT s p)) + 100

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.ORunTables
/-!
# The batched counting oracle: the run of its tables

`oTables_ends`: the third piece `oTables` of the oracle fills the tables `RAC`, `RNBC`, `RNAB`
(residues of the matrices), `NL`, `HK` (residues of the digits), `RS`, `RV` (first rows), `DK`
(quotients by `k`) and `ZH`, `ZL` (the digits of `RNAB`), within the time `tOTables`; it changes no
cell below the free pointer.  The tables, as the specifications give them, are `OTabs`.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Batch Light.Sec4 Light.Filter2

/-- The tables that `oTables` fills, in the forms that the specifications give. -/
structure OTabs (x : Light.Filter2.OracleInst) (fr : ℕ) (μ : ℕ → ℤ) : Prop where
  rac : SegN μ fr (x.AC.map fun w => resid x.p (1 * w))
  rnbc : SegN μ (oaRNBC x.n fr) (x.BC.map fun w => resid x.p (-1 * w))
  rnab : SegN μ (oaRNAB x.n fr) (x.AB.map fun w => resid x.p (-1 * w))
  nl : SegN μ (oaNL x.n fr)
    ((List.range (oDig x.p)).map fun h : ℕ => resid x.p ((h : ℤ) * (-1) + (x.p : ℤ)))
  hk : SegN μ (oaHK x.n x.p fr)
    ((List.range (oDig x.p)).map fun h : ℕ => resid x.p ((h : ℤ) * (oDig x.p : ℤ) + 0))
  rs : SegN μ (oaRS x.n x.p fr) ((List.range (oNP x.n x.s x.p)).map fun R =>
    R / oN0 x.s x.p * (oKK x.s x.p * oN0 x.s x.p) + R % oN0 x.s x.p)
  rv : SegN μ (oaRV x.n x.s x.p fr) ((List.range (oNP x.n x.s x.p)).map fun R =>
    R / oN0 x.s x.p * (oK0 x.s x.p * oN0 x.s x.p) + R % oN0 x.s x.p)
  dk : SegN μ (oaDK x.n x.s x.p fr) ((List.range (oDig x.p * oDig x.p)).map fun z => z / oDig x.p)
  zh : SegN μ (oaZH x.n x.s x.p fr) ((x.AB.map fun w => resid x.p (-1 * w)).map (· / oDig x.p))
  zl : SegN μ (oaZL x.n x.s x.p fr) ((x.AB.map fun w => resid x.p (-1 * w)).map (· % oDig x.p))

/-- The time of `oTables`. -/
def tOTables (n U s p : ℕ) : ℕ :=
  3 * tResidTab (n * n) U p + 2 * tSmallTab (oDig p) (3 * p) p + 2 * tRowTab (oNb n s p) (oN0 s p)
    + tDivTab (oDig p) + tZTab (n * n) + 300

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.ORunParams
/-!
# The batched counting oracle: the run of its parameters

`oParams_ends`: the first piece `oParams` of the oracle computes the parameters `m, D, L, t, N0,
K0, KK, k, nb, NP, Nv, T, B` and `binom(L, m)` into the locals `oLocP`, within the time
`tOParams`; it writes `4^m` to the cell `fr` and uses the `L + 2` cells after it for the table of
binomials.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp Light.Sec4 Light.Filter2

/-- The time of binom(L, m) with the constant 30. -/
def tBinom30 (L : ℕ) : ℕ := 30 * (L + 1) ^ 2

/-- The time of `oParams`. -/
def tOParams (n s p : ℕ) : ℕ :=
  tLog4 (oM s p) + tCeilMul 56569 (oM s p) + tPowNat (oL s p - oM s p)
    + tBinom30 (oL s p) + 18 * (oK0 s p + 1) + 18 * (Nat.sqrt (p - 1) + 1)
    + Sec3.tCeilDiv (n * oDig p) (oN0 s p) + Sec3.tCeilDiv n s
    + Sec3.tCeilDiv (oProd n s) (oKK s p) + 300

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.ORunOracle
/-!
# The batched counting oracle: the run

A batch `j` (`oBatch_ends`), the loop over the batches (`oLoop_ends`), and the whole oracle.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Batch Light.Sec4 Light.Filter2 OS

/-- The number of products of the batch `j`. -/
def oKb (n s p j : ℕ) : ℕ := min (oKK s p) (oProd n s - j * oKK s p)

/-- The time of a batch. -/
def tOBatch (n s p : ℕ) : ℕ := tOBTail n s p + 30

/-- What holds after `j` batches: the tables, the counts of the first `min (j KK) T` products, and
the cells below the free pointer outside the counts. -/
def LoopInv (x : Light.Filter2.OracleInst) (fr : ℕ) (μ₀ : ℕ → ℤ) (j : ℕ) (μ : ℕ → ℤ) : Prop :=
  BTabs x fr μ ∧
    (∀ i < min (j * oKK x.s x.p) (oProd x.n x.s) * (x.n * x.n),
      μ (x.cnt + i) = (Light.Filter2.cntOf x.n x.s x.p x.AB x.BC x.AC i : ℤ)) ∧
    ∀ b < fr, (b < x.cnt ∨ x.cnt + x.len ≤ b) → μ b = μ₀ b

/-- The time of the oracle. -/
def tOracle (n U s p : ℕ) : ℕ :=
  tOParams n s p + tOTables n U s p + oBat n s p * (tOBatch n s p + 8) + 300

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Filter2.Mass
/-!
# The arithmetic of the filtering (Lemma 3.1)

For an active set `A` (a predicate on candidates `i < L`, `L = m n²`) and a modulus `M`, the mass of
the group `g` is `mass A M g = Σ_{i < L, A i, i % n² = g} cnt_M(i)`.  The count at modulus `0` is the
number of exact witnesses (`(0 : ℤ) ∣ v ↔ v = 0`).

* `sum_pool_cntOf_le` (3.2): summed over the pool, the counts at `M q` are at most `Q` times the
  exact witnesses plus `h` times the count at `M`.
* `exists_witness_of_certified` (3.3): a certified group has an exact witness.
* `contract` : with one prime `q*` minimizing the remaining mass, the remaining mass shrinks by
  `2h/Q`.
* `cntOf_pos_of_zero_pos`: an exact witness is counted at every modulus.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec

variable {n s U : ℕ} {AB BC AC : List ℤ}

/-- The mass of the group `g` over the active candidates `A` at the modulus `M`. -/
def mass (n s : ℕ) (AB BC AC : List ℤ) (A : ℕ → Bool) (M g : ℕ) : ℕ :=
  ∑ i ∈ (Finset.range ((n + s - 1) / s * (n * n))).filter (fun i => A i = true ∧ i % (n * n) = g),
    Light.Filter2.cntOf n s M AB BC AC i

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostMath
/-!
# The filtering host: the pure arithmetic

* `Arr`: an array of the memory, given by a function.
* `poolL`: the pool as a list (the primes of `[R, 2R)` in ascending order); it is never empty
  (Bertrand's postulate).
* The mass of a group as a sum over the blocks (`mass_eq_sum`); the total mass `tot` and the number
  `cardA` of active candidates.
* The step of a round (`step_contract`), the end of a round (`card_filter_le_tot`), and the count
  of the candidates (`card_lt_two_pow`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec

/-! ## Arrays -/

/-- The `L` cells from `a` hold `f 0, …, f (L - 1)`. -/
def Arr (μ : ℕ → ℤ) (a L : ℕ) (f : ℕ → ℤ) : Prop := ∀ i < L, μ (a + i) = f i

/-- A truth value as a number. -/
def bitB (b : Bool) : ℤ := if b then 1 else 0

@[simp] theorem bitB_true : Light.Filter2.bitB true = 1 := rfl
@[simp] theorem bitB_false : Light.Filter2.bitB false = 0 := rfl

/-! ## Parameters -/

/-! ## The pool as a list -/

/-- The pool: the primes of `[R, 2R)` in ascending order. -/
def poolL (n : ℕ) : List ℕ := (Sieve.primesBelow (2 * Light.Filter2.RR n)).filter fun q => decide (Light.Filter2.RR n ≤ q)

/-- The pool, read with a default. -/
def poolAt (n t : ℕ) : ℕ := (Light.Filter2.poolL n).getD t 0

/-! ## Masses as sums over the blocks -/

variable {n s U : ℕ} {AB BC AC : List ℤ}

/-- The total mass of the active candidates. -/
def tot (n s : ℕ) (AB BC AC : List ℤ) (A : ℕ → Bool) (M : ℕ) : ℕ :=
  ∑ g ∈ Finset.range (n * n), Light.Filter2.mass n s AB BC AC A M g

/-- The number of active candidates. -/
def cardA (n s : ℕ) (A : ℕ → Bool) : ℕ :=
  ∑ g ∈ Finset.range (n * n), ∑ j ∈ Finset.range ((n + s - 1) / s),
    if A (j * (n * n) + g) then 1 else 0

/-- The candidates of the groups that are not done, among those that are marked active. -/
def eff (N : ℕ) (act D : ℕ → Bool) : ℕ → Bool := fun i => act i && !D (i % N)

/-! ## Witnesses -/

/-! ## The step of a round -/

/-- The certificate of the group `g`, from the counts `act` at the moduli `M` and `M q`. -/
def certOf (n s U : ℕ) (AB BC AC : List ℤ) (act : ℕ → Bool) (M g : ℕ) : Prop :=
  ∀ t < Light.Filter2.QOf n, Light.Filter2.hOf n U * Light.Filter2.mass n s AB BC AC act M g < Light.Filter2.QOf n * Light.Filter2.mass n s AB BC AC act (M * Light.Filter2.poolAt n t) g

/-! ## Rounds -/

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostProcs
/-!
# The filtering host: small procedures

* `div(a, b) = ⌊a/b⌋` and `pow2(e) = 2^e`, by counting up.
* `aggr(cnt, act, dst, N, m)`: `dst[g] := Σ_{j<m} act[j N + g] · cnt[j N + g]` for `g < N`.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

variable {lim : Limits} {P : Program} {d : ℕ}

/-! ## Helpers for expressions -/

/-! ## Division and powers of two, by counting up -/

/-- div(a, b).  Locals: 0 a, 1 b, 2 the candidate q, 3 the multiple (q + 1) b. -/
def divBody : Stmt :=
  .set 2 (Light.k 0) ;;
  .set 3 (Light.v 1) ;;
  .while (Light.v 3 ≤' Light.v 0) (
    .set 3 (Light.v 3 +' Light.v 1) ;;
    .set 2 (Light.v 2 +' Light.k 1)) ;;
  .set 0 (Light.v 2)

/-- The time of div. -/
def tDiv (a b : ℕ) : ℕ := 14 * (a / b) + 12

/-- pow2(e).  Locals: 0 e, 1 the counter i, 2 the power 2^i. -/
def pow2Body : Stmt :=
  .set 1 (Light.k 0) ;;
  .set 2 (Light.k 1) ;;
  .while (Light.v 1 <' Light.v 0) (
    .set 2 (Light.v 2 +' Light.v 2) ;;
    .set 1 (Light.v 1 +' Light.k 1)) ;;
  .set 0 (Light.v 2)

/-- The time of pow2. -/
def tPow2 (e : ℕ) : ℕ := 12 * e + 12

/-! ## Aggregation -/

/-- The inner loop of aggr.  Locals: 0 cnt, 1 act, 2 dst, 3 N, 4 m, 5 g, 6 acc, 7 j. -/
def aggrInner : Stmt :=
  Stmt.for 7 (Light.v 4) (.set 6 (Light.v 6 +' M (Light.v 1 +' Light.v 7 *' Light.v 3 +' Light.v 5) *' M (Light.v 0 +' Light.v 7 *' Light.v 3 +' Light.v 5)))

/-- aggr(cnt, act, dst, N, m): `dst[g] := Σ_{j<m} act[j N + g] · cnt[j N + g]` for `g < N`. -/
def aggrBody : Stmt :=
  Stmt.for 5 (Light.v 3) (.set 6 (Light.k 0) ;; Light.Filter2.aggrInner ;; .store (Light.v 2 +' Light.v 5) (Light.v 6)) ;;
  .set 0 (Light.k 0)

/-- The sum over the first `j` blocks of the counts of the active candidates of the group `g`. -/
def aggrSum (N g : ℕ) (A : ℕ → Bool) (c : ℕ → ℕ) (j : ℕ) : ℕ :=
  ∑ j' ∈ Finset.range j, if A (j' * N + g) then c (j' * N + g) else 0

/-- The time of aggr. -/
def tAggr (N m : ℕ) : ℕ := N * (30 * m + 30) + 10

/-! ## Certification -/

/-- The inner loop of cert: the test of the group `g` against every prime of the pool.
Locals: 0 ua, 1 gt, 2 done, 3 N, 4 Q, 5 h, 6 g, 7 flag, 8 t. -/
def certInner : Stmt :=
  Stmt.for 8 (Light.v 4) (.ite (Light.v 4 *' M (Light.v 1 +' Light.v 8 *' Light.v 3 +' Light.v 6) <' Light.v 5 *' M (Light.v 0 +' Light.v 6) +' Light.k 1)
    (.set 7 (Light.k 0)) .skip)

/-- cert(ua, gt, done, N, Q, h): every group `g < N` that is not done is certified, `done[g] := 1`,
if `h · ua[g] < Q · gt[t N + g]` for every `t < Q`. -/
def certBody : Stmt :=
  Stmt.for 6 (Light.v 3) (.ite (M (Light.v 2 +' Light.v 6) =' Light.k 0)
    (.set 7 (Light.k 1) ;; Light.Filter2.certInner ;; .store (Light.v 2 +' Light.v 6) (Light.v 7)) .skip) ;;
  .set 0 (Light.k 0)

/-- The time of cert. -/
def tCert (N Q : ℕ) : ℕ := N * (30 * Q + 30) + 10

/-! ## The choice of the prime -/

/-- The inner loop of argmin: the remaining mass of the row `t`.
Locals: 0 gt, 1 done, 2 N, 3 Q, 4 t, 5 S, 6 g, 7 best, 8 arg. -/
def argInner : Stmt :=
  Stmt.for 6 (Light.v 2) (.set 5 (Light.v 5 +' (Light.k 1 -' M (Light.v 1 +' Light.v 6)) *' M (Light.v 0 +' Light.v 4 *' Light.v 2 +' Light.v 6)))

/-- argmin(gt, done, N, Q): the first row `t < Q` of least remaining mass
`Σ_{g < N} (1 - done[g]) gt[t N + g]`. -/
def argminBody : Stmt :=
  Stmt.for 4 (Light.v 3) (.set 5 (Light.k 0) ;; Light.Filter2.argInner ;;
    .ite (Light.v 4 =' Light.k 0) (.set 7 (Light.v 5) ;; .set 8 (Light.v 4))
      (.ite (Light.v 5 <' Light.v 7) (.set 7 (Light.v 5) ;; .set 8 (Light.v 4)) .skip)) ;;
  .set 0 (Light.v 8)

/-- The time of argmin. -/
def tArgmin (N Q : ℕ) : ℕ := Q * (30 * N + 30) + 10

/-- The remaining mass of the first `g` groups in the row `t`. -/
def remSum (D : ℕ → Bool) (Gv : ℕ → ℕ → ℕ) (t g : ℕ) : ℕ :=
  ∑ g' ∈ Finset.range g, if D g' then 0 else Gv t g'

/-! ## The end of a round -/

/-- filterAct(act, cnt, L): `act[i] := 0` where `cnt[i] = 0`.  Locals: 0 act, 1 cnt, 2 L, 3 i. -/
def filterActBody : Stmt :=
  Stmt.for 3 (Light.v 2) (.ite (M (Light.v 1 +' Light.v 3) =' Light.k 0) (.store (Light.v 0 +' Light.v 3) (Light.k 0)) .skip) ;;
  .set 0 (Light.k 0)

/-- The time of filterAct. -/
def tFilterAct (L : ℕ) : ℕ := 20 * L + 10

/-! ## The pool -/

/-- poolFilter(src, c, R, dst): the entries `≥ R` of the `c` cells from `src` are written, in their
order, to the cells from `dst`; it returns their number.
Locals: 0 src, 1 c, 2 R, 3 dst, 4 i, 5 the number so far. -/
def poolFilterBody : Stmt :=
  Stmt.for 4 (Light.v 1) (.ite (Light.v 2 ≤' M (Light.v 0 +' Light.v 4))
    (.store (Light.v 3 +' Light.v 5) (M (Light.v 0 +' Light.v 4)) ;; .set 5 (Light.v 5 +' Light.k 1)) .skip) ;;
  .set 0 (Light.v 5)

/-- The time of poolFilter. -/
def tPoolFilter (c : ℕ) : ℕ := 30 * c + 10

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostFinal
/-!
# The filtering host: the last phase

For every group `g` that is not done, the blocks of its active candidates are scanned
(`groupBody`); the flag of `g` is 1 if it is done or a scan finds a zero triangle.  The time depends
on the number of active candidates (`final_spec`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits} {P : Program} {d : ℕ}

/-! ## Counting loops whose rounds take different times -/

/-! ## The scan of a block -/

/-- The scan of the block `j` for the pair of the group `g`. -/
def blockHit (n s : ℕ) (AB BC AC : List ℤ) (j g : ℕ) : Bool :=
  scanHit n AB BC AC (g / n) (g % n) (j * s) (min s (n - j * s))

/-! ## One group -/

/-- One block of the group.  Locals: 0 n, 1 ab, 2 bc, 3 ac, 4 out, 5 done, 6 act, 7 s, 8 m, 9 N,
10 a, 11 b, 12 g, 13 hit, 14 j, 15 c0, 16 len. -/
def groupRound (pScan : ℕ) : Stmt :=
  .ite (Light.v 13 =' Light.k 0)
    (.ite (M (Light.v 6 +' Light.v 14 *' Light.v 9 +' Light.v 12) =' Light.k 0) .skip
      (.set 15 (Light.v 14 *' Light.v 7) ;;
       .set 16 (Light.v 7) ;;
       .ite (Light.v 0 -' Light.v 15 <' Light.v 7) (.set 16 (Light.v 0 -' Light.v 15)) .skip ;;
       .call pScan [Light.v 1, Light.v 2, Light.v 3, Light.v 0, Light.v 10, Light.v 11, Light.v 15, Light.v 16] 13))
    .skip

/-- group(n, ab, bc, ac, out, done, act, s, m, N, a, b, g): the flag of the group `g = a n + b`. -/
def groupBody (pScan : ℕ) : Stmt :=
  .set 13 (M (Light.v 5 +' Light.v 12)) ;;
  Stmt.for 14 (Light.v 8) (Light.Filter2.groupRound pScan) ;;
  .store (Light.v 4 +' Light.v 12) (Light.v 13) ;;
  .set 0 (Light.k 0)

/-- The flag of the group `g` after the first `j` blocks. -/
def hitB (n s : ℕ) (AB BC AC : List ℤ) (act D : ℕ → Bool) (g j : ℕ) : Bool :=
  D g || (List.range j).any fun j' => act (j' * (n * n) + g) && Light.Filter2.blockHit n s AB BC AC j' g

/-- The number of active candidates of a group that is not done. -/
def cntG (N m : ℕ) (act D : ℕ → Bool) (g : ℕ) : ℕ :=
  ∑ j ∈ Finset.range m, if (act (j * N + g) && !D g) then 1 else 0

/-- The time of group. -/
def tGroup (m s c : ℕ) : ℕ := 60 * m + 20 + c * tScan s

/-! ## All groups -/

/-- final(n, ab, bc, ac, out, done, act, s, m, N): the flags of all groups.
Locals: 0 n, 1 ab, 2 bc, 3 ac, 4 out, 5 done, 6 act, 7 s, 8 m, 9 N, 10 a, 11 b, 12 g, 13 r. -/
def finalBody (pGroup : ℕ) : Stmt :=
  Stmt.for 12 (Light.v 9) (
    .call pGroup [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.v 6, Light.v 7, Light.v 8, Light.v 9, Light.v 10, Light.v 11, Light.v 12] 13 ;;
    .set 11 (Light.v 11 +' Light.k 1) ;;
    .ite (Light.v 11 =' Light.v 0) (.set 11 (Light.k 0) ;; .set 10 (Light.v 10 +' Light.k 1)) .skip) ;;
  .set 0 (Light.k 0)

/-- The time of final, with `c` the number of active candidates of the groups that are not done. -/
def tFinal (N m s c : ℕ) : ℕ := N * (60 * m + 60) + c * tScan s + 10

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostStep
/-!
# The filtering host: the layout, and one step of a round

* The layout of the memory of the host above the free pointer `fr`: the primes below `2R`, the pool,
  the marks `act` of the `L = m n²` candidates, the flags `done` of the `n²` groups, the masses `ua`
  of the groups, the counts `cnt` of the oracle, the table `gt` of the masses at the moduli `M q`,
  and the free pointer `top` of the oracle.
* `Env`: what the host knows about the limits and the instance; `FMem`: the memory between the
  steps.
* `stepBody`: one step of a round (`step_spec`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits}

/-! ## The layout -/

/-- The number of candidates, `m n²`. -/
def LL (n : ℕ) : ℕ := Light.Filter2.mOf n * (n * n)

/-- The pool. -/
def aPL (n fr : ℕ) : ℕ := fr + 2 * Light.Filter2.RR n
/-- The marks of the candidates. -/
def aACT (n fr : ℕ) : ℕ := fr + 4 * Light.Filter2.RR n
/-- The flags of the groups. -/
def aDONE (n fr : ℕ) : ℕ := Light.Filter2.aACT n fr + Light.Filter2.LL n
/-- The masses of the groups at the current modulus. -/
def aUA (n fr : ℕ) : ℕ := Light.Filter2.aDONE n fr + n * n
/-- The counts written by the oracle. -/
def aCNT (n fr : ℕ) : ℕ := Light.Filter2.aUA n fr + n * n
/-- The table of the masses at the moduli `M q`. -/
def aGT (n fr : ℕ) : ℕ := Light.Filter2.aCNT n fr + Light.Filter2.LL n
/-- The free pointer of the oracle (and the table of the sieve). -/
def aTOP (n fr : ℕ) : ℕ := Light.Filter2.aGT n fr + 2 * Light.Filter2.RR n * (n * n)

/-- The largest time of the oracle at the moduli up to `Pmax`. -/
def supTo (To : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  (Finset.range (Light.Filter2.PmaxOf n + 1)).sup fun p => To [n, U, Light.Filter2.sOf n, p]

/-- A bound on the numbers that the host forms. -/
def filterWord (n U : ℕ) : ℕ :=
  10 ^ 9 * ((n + 1) ^ 3 * (U + 1) * (Light.Filter2.RR n + 1) * (Light.Filter2.PmaxOf n + 1)) ^ 2

/-- What the host knows: the instance (in the first memory `μ0`), the limits, and the need of the
oracle at all moduli up to `Pmax`.  The host runs at depth `d`. -/
structure Env (lim : Limits) (d : ℕ) (x : AEInst) (fr : ℕ) (μ0 : ℕ → ℤ) (needo : List ℕ → Light.Need) :
    Prop where
  pre : x.Pre μ0 fr
  hw : (lim.space : ℤ) ≤ lim.word
  space : Light.Filter2.aTOP x.n fr + 2 * Light.Filter2.RR x.n ≤ lim.space
  word : (Light.Filter2.filterWord x.n x.U : ℤ) ≤ lim.word
  oracle : ∀ p ≤ Light.Filter2.PmaxOf x.n, (needo [x.n, x.U, Light.Filter2.sOf x.n, p]).Ok lim (Light.Filter2.aTOP x.n fr) (d + 4)
  depth : d + 5 ≤ lim.depth

/-- The memory between two steps: nothing below `fr` has changed, and the pool, the marks `act` and
the flags `D` are in place. -/
structure FMem (x : AEInst) (fr : ℕ) (μ0 μ : ℕ → ℤ) (act D : ℕ → Bool) : Prop where
  kept : Light.Kept μ0 μ fr
  pool : Light.Filter2.Arr μ (Light.Filter2.aPL x.n fr) (Light.Filter2.QOf x.n) fun t => (Light.Filter2.poolAt x.n t : ℤ)
  act : Light.Filter2.Arr μ (Light.Filter2.aACT x.n fr) (Light.Filter2.LL x.n) fun i => Light.Filter2.bitB (act i)
  done : Light.Filter2.Arr μ (Light.Filter2.aDONE x.n fr) (x.n * x.n) fun g => Light.Filter2.bitB (D g)

/-- The masses of the groups that are not done, at the modulus `M`. -/
def UOk (x : AEInst) (fr : ℕ) (μ : ℕ → ℤ) (act D : ℕ → Bool) (M : ℕ) : Prop :=
  ∀ g < x.n * x.n, D g = false →
    μ (Light.Filter2.aUA x.n fr + g) = (Light.Filter2.mass x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC act M g : ℤ)

/-- The instance of the oracle at the modulus `p`. -/
def oInst (x : AEInst) (fr p : ℕ) : OracleInst :=
  { toTriInst := x.toTriInst, s := Light.Filter2.sOf x.n, p := p, cnt := Light.Filter2.aCNT x.n fr }

/-! ## Bounds on the numbers -/

/-- The product behind `filterWord`. -/
def KW (n U : ℕ) : ℕ := (n + 1) ^ 3 * (U + 1) * (Light.Filter2.RR n + 1) * (Light.Filter2.PmaxOf n + 1)

/-! ## Keeping the memory -/

section keep

variable {x : AEInst} {fr : ℕ} {μ0 μ μ' : ℕ → ℤ} {act D : ℕ → Bool} {M : ℕ}

/-- The table of the masses at the moduli `M q`, for the first `t` primes of the pool. -/
def GOk (x : AEInst) (fr : ℕ) (μ : ℕ → ℤ) (act : ℕ → Bool) (M t : ℕ) : Prop :=
  ∀ t' < t, ∀ g < x.n * x.n, μ (Light.Filter2.aGT x.n fr + (t' * (x.n * x.n) + g)) =
    (Light.Filter2.mass x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC act (M * Light.Filter2.poolAt x.n t') g : ℤ)

end keep

/-! ## One step -/

/-- step(n, U, ab, bc, ac, s, pl, act, done, ua, cnt, gt, top, N, m, Q, h, M): the oracle at every
`M q`, the certificate, the choice of `q*`; it returns `M q*`.  Locals: 18 t, 19 scratch, 20 the
row `t*`. -/
def stepBody (pO pAggr pCert pArg pCopy : ℕ) : Stmt :=
  Stmt.for 18 (Light.v 15) (
    .call pO [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.v 17 *' M (Light.v 6 +' Light.v 18), Light.v 10, Light.v 12] 19 ;;
    .call pAggr [Light.v 10, Light.v 7, Light.v 11 +' Light.v 18 *' Light.v 13, Light.v 13, Light.v 14] 19) ;;
  .call pCert [Light.v 9, Light.v 11, Light.v 8, Light.v 13, Light.v 15, Light.v 16] 19 ;;
  .call pArg [Light.v 11, Light.v 8, Light.v 13, Light.v 15] 20 ;;
  .call pCopy [Light.v 11 +' Light.v 20 *' Light.v 13, Light.v 9, Light.v 13] 19 ;;
  .set 0 (Light.v 17 *' M (Light.v 6 +' Light.v 20))

/-- The arguments of step. -/
def stepArgs (x : AEInst) (fr M : ℕ) : List ℤ :=
  [x.n, x.U, x.ab, x.bc, x.ac, Light.Filter2.sOf x.n, Light.Filter2.aPL x.n fr, Light.Filter2.aACT x.n fr, Light.Filter2.aDONE x.n fr, Light.Filter2.aUA x.n fr,
    Light.Filter2.aCNT x.n fr, Light.Filter2.aGT x.n fr, Light.Filter2.aTOP x.n fr, ((x.n * x.n : ℕ) : ℤ), Light.Filter2.mOf x.n, Light.Filter2.QOf x.n, Light.Filter2.hOf x.n x.U, M]

/-- The time of step. -/
def tStep (To : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  Light.Filter2.QOf n * (Light.Filter2.supTo To n U + Light.Filter2.tAggr (n * n) (Light.Filter2.mOf n) + 40) + Light.Filter2.tCert (n * n) (Light.Filter2.QOf n) +
    Light.Filter2.tArgmin (n * n) (Light.Filter2.QOf n) + copyTime (n * n) + 60

/-- The procedures that step calls. -/
structure StepProcs (P : Program) (pAggr pCert pArg pCopy : ℕ) : Prop where
  aggr : P[pAggr]? = some Light.Filter2.aggrBody
  cert : P[pCert]? = some Light.Filter2.certBody
  arg : P[pArg]? = some Light.Filter2.argminBody
  copy : P[pCopy]? = some copyBody

instance (n s U : ℕ) (AB BC AC : List ℤ) (act : ℕ → Bool) (M g : ℕ) :
    Decidable (Light.Filter2.certOf n s U AB BC AC act M g) := by
  unfold Light.Filter2.certOf; infer_instance

section step

variable {P₀ R : Program} {pO pAggr pCert pArg pCopy : ℕ} {To : List ℕ → ℕ}
  {needo : List ℕ → Need} {d0 : ℕ} {x : AEInst} {fr : ℕ} {μ0 : ℕ → ℤ}

/-- The flags after the certificate of a step at the modulus `Mc`. -/
def stepD (x : AEInst) (act D : ℕ → Bool) (Mc : ℕ) : ℕ → Bool := fun g =>
  D g || decide (∀ t < Light.Filter2.QOf x.n, Light.Filter2.hOf x.n x.U * Light.Filter2.mass x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC act Mc g <
    Light.Filter2.QOf x.n * Light.Filter2.mass x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC act (Mc * Light.Filter2.poolAt x.n t) g)

/-- The remaining mass at the modulus `M`, outside the groups that are done. -/
def remMass (x : AEInst) (act D : ℕ → Bool) (M : ℕ) : ℕ :=
  ∑ g ∈ Finset.range (x.n * x.n), if D g then 0 else Light.Filter2.mass x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC act M g

end step

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostRound
/-!
# The filtering host: one round

A round: the masses at the modulus `1`, then `r` steps, each of which certifies some groups and
multiplies the modulus by a prime of the pool, and at the end the marks of the candidates whose
count at the final modulus is `0` are cleared (`round_spec`).

The invariants: a group that is done has a witness (`Inv1`), and every candidate of a group that is
not done whose block has a witness is marked (`Inv2`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits}

/-- A group that is done has a witness. -/
def Inv1 (x : AEInst) (D : ℕ → Bool) : Prop :=
  ∀ g < x.n * x.n, D g = true → ∃ c < x.n, Light.Filter2.triSum x.n x.AB x.BC x.AC g c = 0

/-- A candidate of a group that is not done, whose block has a witness, is marked. -/
def Inv2 (x : AEInst) (act D : ℕ → Bool) : Prop :=
  ∀ i < Light.Filter2.LL x.n, D (i % (x.n * x.n)) = false → 0 < Light.Filter2.cntOf x.n (Light.Filter2.sOf x.n) 0 x.AB x.BC x.AC i →
    act i = true

/-- round(n, U, ab, bc, ac, s, pl, act, done, ua, cnt, gt, top, N, m, Q, h, r).  Locals: 18 the
modulus `M`, 19 the step, 20 scratch. -/
def roundBody (pO pAggr pStep pFA : ℕ) : Stmt :=
  .call pO [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.k 1, Light.v 10, Light.v 12] 20 ;;
  .call pAggr [Light.v 10, Light.v 7, Light.v 9, Light.v 13, Light.v 14] 20 ;;
  .set 18 (Light.k 1) ;;
  Stmt.for 19 (Light.v 17) (.call pStep [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.v 6, Light.v 7, Light.v 8, Light.v 9, Light.v 10, Light.v 11,
    Light.v 12, Light.v 13, Light.v 14, Light.v 15, Light.v 16, Light.v 18] 18) ;;
  .call pO [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.v 18, Light.v 10, Light.v 12] 20 ;;
  .call pFA [Light.v 7, Light.v 10, Light.v 14 *' Light.v 13] 20

/-- The time of round. -/
def tRound (To : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  2 * Light.Filter2.supTo To n U + Light.Filter2.tAggr (n * n) (Light.Filter2.mOf n) + Light.Filter2.rOf n * (Light.Filter2.tStep To n U + 58) + Light.Filter2.tFilterAct (Light.Filter2.LL n) + 200

/-- The locals of round. -/
def rFrame (x : AEInst) (fr Mc st : ℕ) (z : ℤ) : List ℤ :=
  Light.Filter2.stepArgs x fr (Light.Filter2.rOf x.n) ++ [(Mc : ℤ), (st : ℤ), z]

/-- The procedures that round calls. -/
structure RoundProcs (P : Program) (pO pAggr pStep pFA pCert pArg pCopy : ℕ) : Prop where
  step : P[pStep]? = some (Light.Filter2.stepBody pO pAggr pCert pArg pCopy)
  sp : Light.Filter2.StepProcs P pAggr pCert pArg pCopy
  fa : P[pFA]? = some Light.Filter2.filterActBody

section round

variable {P₀ R : Program} {pO pAggr pStep pFA pCert pArg pCopy : ℕ} {To : List ℕ → ℕ}
  {needo : List ℕ → Need} {d0 : ℕ} {x : AEInst} {fr : ℕ} {μ0 : ℕ → ℤ}

/-- The invariant of the loop over the steps. -/
def StepInv (x : AEInst) (fr : ℕ) (μ0 : ℕ → ℤ) (act D : ℕ → Bool) (st : ℕ) (σ : State) : Prop :=
  ∃ (Mc : ℕ) (z : ℤ) (D' : ℕ → Bool) (μ' : ℕ → ℤ), σ = ⟨Light.frame (Light.Filter2.rFrame x fr Mc st z), μ'⟩ ∧
    Light.Filter2.FMem x fr μ0 μ' act D' ∧ Light.Filter2.UOk x fr μ' act D' Mc ∧ 1 ≤ Mc ∧ Mc ≤ (2 * Light.Filter2.RR x.n) ^ st ∧
    Light.Filter2.Inv1 x D' ∧ (∀ g, D g = true → D' g = true) ∧
    Light.Filter2.QOf x.n ^ st * Light.Filter2.tot x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC (Light.Filter2.eff (x.n * x.n) act D') Mc ≤
      (2 * Light.Filter2.hOf x.n x.U) ^ st * Light.Filter2.tot x.n (Light.Filter2.sOf x.n) x.AB x.BC x.AC (Light.Filter2.eff (x.n * x.n) act D) 1

end round

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostCore
/-!
# The filtering host: the rounds and the final scans

`core` runs the rounds and then scans the blocks of the candidates that are still marked, in the
groups that are not done (`core_spec`).  The flags it writes are the flags of all-edges Exact
Triangle, whatever the guard; with the guard, no candidate is marked after the rounds, so that the
final scans cost nothing.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits}

/-! ## The flags -/

/-! ## The program -/

/-- core(n, U, ab, bc, ac, s, pl, act, done, ua, cnt, gt, top, N, m, Q, h, r, rounds, out).  Locals:
20 the round, 21 scratch. -/
def coreBody (pRound pFinal : ℕ) : Stmt :=
  Stmt.for 20 (Light.v 18) (.call pRound [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 5, Light.v 6, Light.v 7, Light.v 8, Light.v 9, Light.v 10, Light.v 11,
    Light.v 12, Light.v 13, Light.v 14, Light.v 15, Light.v 16, Light.v 17] 21) ;;
  .call pFinal [Light.v 0, Light.v 2, Light.v 3, Light.v 4, Light.v 19, Light.v 8, Light.v 7, Light.v 5, Light.v 14, Light.v 13] 21 ;;
  .set 0 (Light.k 0)

/-- The arguments of core. -/
def coreArgs (x : AEInst) (fr : ℕ) : List ℤ :=
  Light.Filter2.stepArgs x fr (Light.Filter2.rOf x.n) ++ [(Light.Filter2.roundsOf x.n : ℤ), (x.out : ℤ)]

/-- The number of candidates that the final scans may meet: none with the guard, all without. -/
def cFin (n U : ℕ) : ℕ := if Light.Filter2.Guard n U then 0 else Light.Filter2.LL n

/-- The time of core. -/
def tCore (To : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  Light.Filter2.roundsOf n * (Light.Filter2.tRound To n U + 30) + Light.Filter2.tFinal (n * n) (Light.Filter2.mOf n) (Light.Filter2.sOf n) (Light.Filter2.cFin n U) + 30

/-- The procedures of the final scans. -/
structure FinalProcs (P : Program) (pFinal pGroup pScan : ℕ) : Prop where
  final : P[pFinal]? = some (Light.Filter2.finalBody pGroup)
  group : P[pGroup]? = some (Light.Filter2.groupBody pScan)
  scan : P[pScan]? = some scanBody

section core

variable {P₀ R : Program} {pO pAggr pStep pFA pCert pArg pCopy : ℕ} {To : List ℕ → ℕ}
  {needo : List ℕ → Need} {d0 : ℕ} {x : AEInst} {fr : ℕ} {μ0 : ℕ → ℤ}

/-- The invariant of the rounds. -/
def CoreInv (x : AEInst) (fr : ℕ) (μ0 : ℕ → ℤ) (ρ : ℕ) (μ : ℕ → ℤ) : Prop :=
  ∃ act D : ℕ → Bool, Light.Filter2.FMem x fr μ0 μ act D ∧ Light.Filter2.Inv1 x D ∧ Light.Filter2.Inv2 x act D ∧
    (Light.Filter2.Guard x.n x.U → 2 ^ (Light.Filter2.dOf x.n * ρ) * Light.Filter2.cardA x.n (Light.Filter2.sOf x.n) (Light.Filter2.eff (x.n * x.n) act D) ≤ Light.Filter2.LL x.n)

end core

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostTop
/-!
# The filtering host: the parameters

The first part of the host computes the parameters of the filtering from `n` and `U` and lays out
the memory above the free pointer (`params_spec`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits}

/-! ## The computation of the parameters -/

/-- The parameters, from `n` (local 0), `U` (local 1) and the free pointer (local 6): 7 `k`, 8 `ℓ`,
9 `R`, 10 `b`, 11 `d`, 12 the number of rounds, 13 `s`, 14 `m`, 15 `c_K`, 16 `r`, 17 `B`, 18 `h`,
19 `N = n²`, and the places 20 `pl`, 21 `act`, 22 `done`, 23 `ua`, 24 `cnt`, 25 `gt`, 26 `top`. -/
def paramsStmt (pLog pSqrt pPow pDiv : ℕ) : Stmt :=
  .call pLog [Light.v 0] 7 ;; .set 7 (Light.v 7 +' Light.k 1) ;;
  .call pSqrt [Light.v 7] 8 ;; .set 8 (Light.v 8 +' Light.k 1) ;;
  .call pPow [Light.v 8] 9 ;;
  .call pSqrt [Light.v 8] 10 ;; .set 10 (Light.v 8 *' (Light.v 10 +' Light.k 1)) ;;
  .call pDiv [Light.v 10, Light.k 2] 11 ;;
  .call pDiv [Light.k 3 *' Light.v 7, Light.v 11] 12 ;; .set 12 (Light.v 12 +' Light.k 1) ;;
  .call pDiv [Light.k 336175 *' Light.v 7, Light.k 10000000] 13 ;; .call pPow [Light.v 13] 13 ;;
  .call pDiv [Light.v 0 +' Light.v 13 -' Light.k 1, Light.v 13] 14 ;;
  .call pLog [Light.v 13] 15 ;; .set 15 (Light.v 15 +' Light.k 1) ;;
  .call pDiv [Light.v 15 +' Light.v 10, Light.v 8] 16 ;; .set 16 (Light.v 16 +' Light.k 1) ;;
  .call pLog [Light.k 3 *' Light.v 1] 17 ;; .set 17 (Light.v 17 +' Light.k 1) ;;
  .call pDiv [Light.v 17 -' Light.k 1, Light.v 8] 18 ;;
  .set 19 (Light.v 0 *' Light.v 0) ;;
  .set 20 (Light.v 6 +' Light.k 2 *' Light.v 9) ;;
  .set 21 (Light.v 6 +' Light.k 4 *' Light.v 9) ;;
  .set 22 (Light.v 21 +' Light.v 14 *' Light.v 19) ;;
  .set 23 (Light.v 22 +' Light.v 19) ;;
  .set 24 (Light.v 23 +' Light.v 19) ;;
  .set 25 (Light.v 24 +' Light.v 14 *' Light.v 19) ;;
  .set 26 (Light.v 25 +' Light.k 2 *' Light.v 9 *' Light.v 19)

/-- The locals after the parameters. -/
def topFrame (x : AEInst) (fr : ℕ) : List ℤ :=
  [x.n, x.U, x.ab, x.bc, x.ac, x.out, fr, Light.Filter2.kOf x.n, Light.Filter2.ellOf x.n, Light.Filter2.RR x.n, Light.Filter2.bOf x.n, Light.Filter2.dOf x.n,
    Light.Filter2.roundsOf x.n, Light.Filter2.sOf x.n, Light.Filter2.mOf x.n, Light.Filter2.cKOf x.n, Light.Filter2.rOf x.n, Light.Filter2.BOf x.U, Light.Filter2.hOf x.n x.U, ((x.n * x.n : ℕ) : ℤ),
    Light.Filter2.aPL x.n fr, Light.Filter2.aACT x.n fr, Light.Filter2.aDONE x.n fr, Light.Filter2.aUA x.n fr, Light.Filter2.aCNT x.n fr, Light.Filter2.aGT x.n fr, Light.Filter2.aTOP x.n fr]

/-- The time of the parameters. -/
def tParams (n U : ℕ) : ℕ :=
  log2Time n + sqrtTime (Light.Filter2.kOf n) + Light.Filter2.tPow2 (Light.Filter2.ellOf n) + sqrtTime (Light.Filter2.ellOf n) + Light.Filter2.tDiv (Light.Filter2.bOf n) 2 +
    Light.Filter2.tDiv (3 * Light.Filter2.kOf n) (Light.Filter2.dOf n) + Light.Filter2.tDiv (336175 * Light.Filter2.kOf n) (10 ^ 7) + Light.Filter2.tPow2 (336175 * Light.Filter2.kOf n / 10 ^ 7) +
    Light.Filter2.tDiv (n + Light.Filter2.sOf n - 1) (Light.Filter2.sOf n) + log2Time (Light.Filter2.sOf n) + Light.Filter2.tDiv (Light.Filter2.cKOf n + Light.Filter2.bOf n) (Light.Filter2.ellOf n) +
    log2Time (3 * U) + Light.Filter2.tDiv (Light.Filter2.BOf U - 1) (Light.Filter2.ellOf n) + 400

/-- The procedures for the parameters. -/
structure ParamProcs (P : Program) (pLog pSqrt pPow pDiv : ℕ) : Prop where
  log : P[pLog]? = some log2Body
  sqrt : P[pSqrt]? = some sqrtBody
  pow : P[pPow]? = some Light.Filter2.pow2Body
  div : P[pDiv]? = some Light.Filter2.divBody

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostAssemble
/-!
# The filtering host: the top procedure

After the parameters, the host sieves the primes below `2R`, keeps those of the pool, marks all
candidates and no group, and runs `core` (`top_spec`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

variable {lim : Limits}

/-- The pool, the marks and the flags, and the core.  Locals 27 `Q`, 28 scratch. -/
def restStmt (pSieve pPF pFill pCore : ℕ) : Stmt :=
  .call pSieve [Light.k 2 *' Light.v 9 -' Light.k 1, Light.v 6, Light.v 26] 27 ;;
  .call pPF [Light.v 6, Light.v 27, Light.v 9, Light.v 20] 27 ;;
  .call pFill [Light.v 21, Light.v 14 *' Light.v 19, Light.k 1] 28 ;;
  .call pFill [Light.v 22, Light.v 19, Light.k 0] 28 ;;
  .call pCore [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 13, Light.v 20, Light.v 21, Light.v 22, Light.v 23, Light.v 24, Light.v 25, Light.v 26, Light.v 19,
    Light.v 14, Light.v 27, Light.v 18, Light.v 16, Light.v 12, Light.v 5] 28 ;;
  .set 0 (Light.k 0)

/-- The host. -/
def topBody (pLog pSqrt pPow pDiv pSieve pPF pFill pCore : ℕ) : Stmt :=
  Light.Filter2.paramsStmt pLog pSqrt pPow pDiv ;; Light.Filter2.restStmt pSieve pPF pFill pCore

/-- The time of the host. -/
def tTop (To : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  Light.Filter2.tParams n U + sieveTime (2 * Light.Filter2.RR n - 1) + Light.Filter2.tPoolFilter (2 * Light.Filter2.RR n) + fillTime (Light.Filter2.LL n) +
    fillTime (n * n) + Light.Filter2.tCore To n U + 200

/-- The procedures of the host besides core. -/
structure TopProcs (P : Program) (pSieve pPF pFill : ℕ) : Prop where
  sieve : P[pSieve]? = some sieveBody
  pf : P[pPF]? = some Light.Filter2.poolFilterBody
  fill : P[pFill]? = some fillBody

section top

variable {P₀ R : Program} {pO pAggr pStep pFA pCert pArg pCopy : ℕ} {To : List ℕ → ℕ}
  {needo : List ℕ → Need} {d : ℕ} {x : AEInst} {fr : ℕ} {μ0 : ℕ → ℤ}

end top

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostMain
/-!
# The filtering host solves all-edges Exact Triangle

`filter_host`: from a solver of the counting oracle, the host builds a solver of all-edges Exact
Triangle, within `filterTime To` steps and with the need `filterNeed needo`; its need is
polynomially bounded if that of the oracle is.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

/-! ## The program -/

/-- The procedures of the host, appended to a program of length `o` whose procedure `po` is the
counting oracle.  The host is procedure `o + 18`. -/
def hostProcs (o po : ℕ) : Program :=
  [Light.Filter2.aggrBody, Light.Filter2.certBody, Light.Filter2.argminBody, copyBody, Light.Filter2.filterActBody,
    Light.Filter2.stepBody po o (o + 1) (o + 2) (o + 3),
    Light.Filter2.roundBody po o (o + 5) (o + 4),
    scanBody, Light.Filter2.groupBody (o + 7), Light.Filter2.finalBody (o + 8), Light.Filter2.coreBody (o + 6) (o + 9),
    log2Body, sqrtBody, Light.Filter2.pow2Body, Light.Filter2.divBody, sieveBody, Light.Filter2.poolFilterBody, fillBody,
    Light.Filter2.topBody (o + 11) (o + 12) (o + 13) (o + 14) (o + 15) (o + 16) (o + 17) (o + 10)]

/-! ## The time and the need -/

/-- **The time of the host**: the parameters, the sieve, the pool, the marks, and the core, which
runs `roundsOf n` rounds of `rOf n` steps; a step calls the oracle `QOf n` times (once per prime of
the pool) and a round twice more, each call within `supTo To n U`. -/
def filterTime (To : List ℕ → ℕ) (n U : ℕ) : ℕ := Light.Filter2.tTop To n U

/-- The largest word of the oracle at the moduli up to `Pmax`. -/
def supW (needo : List ℕ → Light.Need) (n U : ℕ) : ℕ :=
  (Finset.range (Light.Filter2.PmaxOf n + 1)).sup fun p => (needo [n, U, Light.Filter2.sOf n, p]).word

/-- The most cells of the oracle at the moduli up to `Pmax`. -/
def supC (needo : List ℕ → Light.Need) (n U : ℕ) : ℕ :=
  (Finset.range (Light.Filter2.PmaxOf n + 1)).sup fun p => (needo [n, U, Light.Filter2.sOf n, p]).cells

/-- The largest depth of the oracle at the moduli up to `Pmax`. -/
def supD (needo : List ℕ → Light.Need) (n U : ℕ) : ℕ :=
  (Finset.range (Light.Filter2.PmaxOf n + 1)).sup fun p => (needo [n, U, Light.Filter2.sOf n, p]).depth

/-- **The need of the host.** -/
def filterNeed (needo : List ℕ → Light.Need) (n U : ℕ) : Light.Need where
  word := Light.Filter2.filterWord n U + Light.Filter2.supW needo n U
  cells := Light.Filter2.aTOP n 0 + 2 * Light.Filter2.RR n + Light.Filter2.supC needo n U
  depth := Light.Filter2.supD needo n U + 5

variable {lim : Limits}

/-! ## The host solves all-edges Exact Triangle -/

/-! ## The need is polynomially bounded -/

/-- A bound on the need of the oracle at the moduli up to `Pmax`. -/
def oBound (s0 k0 n U : ℕ) : ℕ :=
  2 ^ s0 * ((n + 1) * ((U + 1) * ((2 * n + 2) * ((2 ^ 8 * (2 * n + 1) ^ 11 + 1) * 1)))) ^ k0

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Filter2.HostTime
/-!
# The time of the filtering host

`filterTime_le`: the host calls the oracle `roundsOf n * (rOf n * QOf n + 2)` times, each within
`supTo To n U`; besides, it spends `1000 (n + 1)² (Q + 1) (m + Q + 1)` steps per step of a round and
once for the final pass, `cFin n U * tScan (sOf n)` steps for the final scans (`0` under the guard,
`m n² (24 s + 35)` without it), and `tPrelude n U` steps for the parameters, the sieve, the pool and
the marks.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

/-- The time of the parameters, the sieve, the pool and the marks. -/
def tPrelude (n U : ℕ) : ℕ :=
  Light.Filter2.tParams n U + sieveTime (2 * Light.Filter2.RR n - 1) + Light.Filter2.tPoolFilter (2 * Light.Filter2.RR n) + fillTime (Light.Filter2.LL n) +
    fillTime (n * n) + 300

/-- The work of one step of a round, apart from the calls of the oracle (a bound). -/
def stepWork (n : ℕ) : ℕ := 1000 * ((n + 1) ^ 2 * (Light.Filter2.QOf n + 1) * (Light.Filter2.mOf n + Light.Filter2.QOf n + 1))

end Light.Filter2
end
end