APSP in $O(n^{2.995561})$: definitions, part 1 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_APSP2996001_Defs1
import Definitions.Def_APSP299791_Defs1
import Definitions.Def_APSP299791_Defs2
import Definitions.Def_APSP299791_Defs3
import Definitions.Def_APSP299791_Defs4
import Definitions.Def_APSPSource_PaperStatements
import Definitions.Def_APSPSource_RemainingDefinitions
import Definitions.Def_APSPSource_ThreeSumApsp_Lang_Logic
import Definitions.Def_APSPSource_ThreeSumApsp_Lang_Regions
import Definitions.Def_APSPSource_ThreeSumApsp_Lang_Syntax
import Definitions.Def_APSPSource_ThreeSumApsp_Lang_WordSize
import Definitions.Def_APSPSource_ThreeSumApsp_Machine_Realized
import Definitions.Def_APSPSource_ThreeSumApsp_Machine_Solving
import Definitions.Def_APSPSource_ThreeSumApsp_Programs_Sec2_Theorem5_Places
import Definitions.Def_APSPSource_ThreeSumApsp_Sec2_Lemma10
import Definitions.Def_APSPSource_ThreeSumApsp_Sec2_Lemma6
import Definitions.Def_APSPSource_ThreeSumApsp_Sec2_Lemma7_8
import Definitions.Def_APSPSource_ThreeSumApsp_Sec2_Levels
import Definitions.Def_APSPSource_ThreeSumApsp_Sec2_Tiling_Definitions
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_Boxes
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_BoxesFromLeaves
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_Lemma27_28
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_ParameterSteps
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_Table2_LogBounds
import Definitions.Def_APSPSource_ThreeSumApsp_Sec4_Theorem30
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_Alphabets
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_Arrays
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_Counters
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_Layout
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_SubsetTable
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec2_Theorem5_Subsets
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec3_Theorem17_Instances
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec4_Theorem30_Cubes
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec4_Theorem30_NineScan
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec4_Theorem30_NineStrings
import Definitions.Def_APSPSource_ThreeSumApsp_Spec_Sec4_Theorem30_Trie
import Definitions.Def_APSPSource_ThreeSumApsp_TimeClaims_Sec3_Definitions
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Asymptotics_Dominated
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Asymptotics_Scale
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Asymptotics_UpperBounds
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Digits
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Flag
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Index
import Definitions.Def_APSPSource_ThreeSumApsp_Util_List
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Log
import Definitions.Def_APSPSource_ThreeSumApsp_Util_PrimesInWindow
import Definitions.Def_APSPSource_ThreeSumApsp_Util_Weave
import Definitions.Def_TrulySubcubicAPSP_SourceSpecification
import Mathlib.Algebra.BigOperators.Associated
import Mathlib.Algebra.BigOperators.Fin
import Mathlib.Algebra.BigOperators.Group.Finset.Basic
import Mathlib.Algebra.BigOperators.GroupWithZero.Finset
import Mathlib.Algebra.BigOperators.Intervals
import Mathlib.Algebra.BigOperators.Ring.Finset
import Mathlib.Algebra.Field.GeomSum
import Mathlib.Algebra.Group.Action.Defs
import Mathlib.Algebra.Group.Int.Defs
import Mathlib.Algebra.GroupWithZero.Nat
import Mathlib.Algebra.MvPolynomial.Basic
import Mathlib.Algebra.MvPolynomial.CommRing
import Mathlib.Algebra.Order.BigOperators.Group.Finset
import Mathlib.Algebra.Order.BigOperators.Group.List
import Mathlib.Algebra.Order.BigOperators.Group.Multiset
import Mathlib.Algebra.Order.BigOperators.GroupWithZero.List
import Mathlib.Algebra.Order.BigOperators.Ring.Finset
import Mathlib.Algebra.Order.Field.Basic
import Mathlib.Algebra.Order.Field.GeomSum
import Mathlib.Algebra.Order.Floor.Div
import Mathlib.Algebra.Order.Floor.Semifield
import Mathlib.Algebra.Order.Floor.Semiring
import Mathlib.Algebra.Order.Group.Abs
import Mathlib.Algebra.Order.Group.Int
import Mathlib.Algebra.Order.Group.Nat
import Mathlib.Algebra.Order.Group.Unbundled.Abs
import Mathlib.Algebra.Order.Ring.Abs
import Mathlib.Algebra.Order.Ring.Int
import Mathlib.Algebra.Order.Ring.Nat
import Mathlib.Algebra.Order.Ring.Rat
import Mathlib.Analysis.Asymptotics.Lemmas
import Mathlib.Analysis.Complex.Exponential
import Mathlib.Analysis.Complex.ExponentialBounds
import Mathlib.Analysis.Convex.SpecificFunctions.Basic
import Mathlib.Analysis.Normed.Field.Basic
import Mathlib.Analysis.SpecialFunctions.BinaryEntropy
import Mathlib.Analysis.SpecialFunctions.Log.Base
import Mathlib.Analysis.SpecialFunctions.Log.Basic
import Mathlib.Analysis.SpecialFunctions.Pow.Asymptotics
import Mathlib.Analysis.SpecialFunctions.Pow.Real
import Mathlib.Analysis.SpecificLimits.Normed
import Mathlib.Combinatorics.Enumerative.DoubleCounting
import Mathlib.Combinatorics.SimpleGraph.Basic
import Mathlib.Data.Bool.Count
import Mathlib.Data.Fin.VecNotation
import Mathlib.Data.Finset.Powerset
import Mathlib.Data.Finset.Sort
import Mathlib.Data.Fintype.BigOperators
import Mathlib.Data.Fintype.Fin
import Mathlib.Data.Int.Cast.Lemmas
import Mathlib.Data.Int.Notation
import Mathlib.Data.Int.SuccPred
import Mathlib.Data.List.GetD
import Mathlib.Data.List.Induction
import Mathlib.Data.List.Iterate
import Mathlib.Data.List.MinMax
import Mathlib.Data.List.OfFn
import Mathlib.Data.List.Range
import Mathlib.Data.List.TakeWhile
import Mathlib.Data.Nat.Bitwise
import Mathlib.Data.Nat.Cast.Order.Ring
import Mathlib.Data.Nat.Choose.Basic
import Mathlib.Data.Nat.Choose.Bounds
import Mathlib.Data.Nat.Choose.Sum
import Mathlib.Data.Nat.Count
import Mathlib.Data.Nat.Digits.Defs
import Mathlib.Data.Nat.Digits.Lemmas
import Mathlib.Data.Nat.Log
import Mathlib.Data.Nat.Notation
import Mathlib.Data.Nat.Prime.Factorial
import Mathlib.Data.Nat.Size
import Mathlib.Data.Nat.SuccPred
import Mathlib.Data.Real.Basic
import Mathlib.Data.Set.Function
import Mathlib.Data.ZMod.Basic
import Mathlib.LinearAlgebra.Finsupp.LinearCombination
import Mathlib.LinearAlgebra.Matrix.Defs
import Mathlib.LinearAlgebra.Matrix.Notation
import Mathlib.Logic.Equiv.Basic
import Mathlib.Logic.Function.Basic
import Mathlib.Logic.Function.Iterate
import Mathlib.MeasureTheory.Integral.Bochner.Basic
import Mathlib.NumberTheory.Bertrand
import Mathlib.NumberTheory.Chebyshev
import Mathlib.NumberTheory.PrimeCounting
import Mathlib.Order.ConditionallyCompleteLattice.Basic
import Mathlib.Order.Filter.AtTopBot.Basic
import Mathlib.Order.Interval.Finset.Fin
import Mathlib.Order.Interval.Finset.Nat
import Mathlib.Probability.Independence.Basic
import Mathlib.RingTheory.AdjoinRoot
import Mathlib.Tactic.Common
import Mathlib.Tactic.DeriveFintype
import Mathlib.Tactic.GCongr
import Mathlib.Tactic.Linarith
import Mathlib.Tactic.LinearCombination
import Mathlib.Tactic.NoncommRing
import Mathlib.Tactic.NormNum
import Mathlib.Tactic.Positivity
import Mathlib.Tactic.Ring
import Mathlib.Tactic.SplitIfs

set_option autoImplicit false
set_option relaxedAutoImplicit false

section
-- Source module: ThreeSumApsp.Lang.Lib.Emod
/-!
# The remainder of a division

emod(x, M, fr) returns the remainder of the integer x modulo M ≥ 1, a number in 0, …, M - 1.  The
language has no division: the routine writes M, 2M, 4M, … into a table at the free pointer fr as
long as they are at most |x|, subtracts them from |x| from the largest down where possible, and
corrects the sign (`emod_meets`).  With k = `emodRounds |x| M` doublings it takes `emodTime |x| M`
steps, a constant times k + 1, and k cells.
-/

@[expose] public section

namespace Light

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

/-- The local variables of emod: the arguments x, M and fr (the result replaces x), what is left
of |x|, the number of entries of the table, and the next multiple of M. -/
abbrev Arg : ℕ := 0
@[inherit_doc Light.Emod.Arg] abbrev Modulus : ℕ := 1
@[inherit_doc Light.Emod.Arg] abbrev Free : ℕ := 2
@[inherit_doc Light.Emod.Arg] abbrev Rest : ℕ := 3
@[inherit_doc Light.Emod.Arg] abbrev Count : ℕ := 4
@[inherit_doc Light.Emod.Arg] abbrev Mult : ℕ := 5

end Emod

open Emod in
/-- One entry of the table: the multiple is written to the table and doubled. -/
def emodTableRound : Stmt :=
  .store (Light.v Light.Emod.Free +' Light.v Light.Emod.Count) (Light.v Light.Emod.Mult) ;;
  .set Light.Emod.Count (Light.v Light.Emod.Count +' Light.k 1) ;;
  .set Light.Emod.Mult (Light.v Light.Emod.Mult +' Light.v Light.Emod.Mult)

open Emod in
/-- The table: M, 2M, 4M, … are written to the free pointer as long as they are at most |x|. -/
def emodTable : Stmt := .while (Light.v Light.Emod.Mult ≤' Light.v Rest) emodTableRound

open Emod in
/-- The last entry of the table that has not been tried is subtracted if possible. -/
def emodReduceRound : Stmt :=
  .set Light.Emod.Count (Light.v Light.Emod.Count -' Light.k 1) ;;
  .ite (Light.M (Light.v Light.Emod.Free +' Light.v Light.Emod.Count) ≤' Light.v Rest) (.set Rest (Light.v Rest -' Light.M (Light.v Light.Emod.Free +' Light.v Light.Emod.Count))) .skip

open Emod in
/-- The entries of the table are subtracted, from the largest down, where possible. -/
def emodReduce : Stmt := .while (Light.k 0 <' Light.v Light.Emod.Count) emodReduceRound

open Emod in
/-- The sign: the remainder r of |x| gives the remainder of x. -/
def emodSign : Stmt :=
  .ite (Light.v Light.Emod.Arg <' Light.k 0)
    (.ite (Light.v Rest =' Light.k 0) (.set Light.Emod.Arg (Light.k 0)) (.set Light.Emod.Arg (Light.v Modulus -' Light.v Rest)))
    (.set Light.Emod.Arg (Light.v Rest))

open Emod in
/-- emod(x, M, fr). -/
def emodBody : Stmt :=
  .ite (Light.v Light.Emod.Arg <' Light.k 0) (.set Rest (Light.k 0 -' Light.v Light.Emod.Arg)) (.set Rest (Light.v Light.Emod.Arg)) ;;
  .set Light.Emod.Count (Light.k 0) ;;
  .set Light.Emod.Mult (Light.v Modulus) ;;
  emodTable ;;
  emodReduce ;;
  emodSign

/-- The number of doublings: the least k with a < M 2^k. -/
def emodRounds (a M : ℕ) : ℕ := Nat.size (a / M)

/-- The time of emod, for a = |x|. -/
@[simp] def emodTime (a M : ℕ) : ℕ := 43 * emodRounds a M + 34

namespace Emod

/-- What emod asks of the limits, for a modulus M ≥ 1, numbers up to V and the free pointer fr. -/
structure Pre (lim : Limits) (Mo V fr : ℕ) : Prop where
  mod_pos : 1 ≤ Mo
  space : (lim.space : ℤ) ≤ lim.word
  word : (4 * V + 4 * Mo + 16 : ℤ) ≤ lim.word
  free : fr + (Nat.log 2 (V + 1) + 2) ≤ lim.space

/-- The first i entries of the table are in place, and no cell below the free pointer has changed.
-/
def Table (μ μ' : ℕ → ℤ) (Mo fr i : ℕ) : Prop :=
  (∀ j < i, μ' (fr + j) = ((Mo * 2 ^ j : ℕ) : ℤ)) ∧ Light.Kept μ μ' fr

variable {μ : ℕ → ℤ} {x : ℤ} {Mo V fr a : ℕ}

end Emod

end Light
end
end

section
-- Source module: ThreeSumApsp.Batch.Tiling
/-!
# Batched tiles: independent products in the slots of one run

Lemma 9 holds for every family `Q ↦ (X_Q, Y_Q)`, so the `K₀²` slots of one run of `Full` may hold
block products of `K₀²` different matrix products. Here the slot `(g, h)` of the band `β` holds the
block number `srcBlock β g h = (β K₀ + g) K₀ + h` of a stacked source matrix, on both sides. If the
source stacks the products interleaved (the row block `i` of product `u` is the block `i K₀² + u`),
then the slot `(g, h)` of the tile `(β, β')` computes the block `(β, β')` of the product
`u = g K₀ + h`.

* `bFamilyL`, `bFamilyR`, `bArrayL`, `bArrayR`: the families and the input arrays of a band.
* `Full_bArray_eq_block_mul`: one run computes the product of the two source blocks of each slot.
* `queryValue_bArray`: a query for a position `(I, J)` of the tile, as in Theorem 30, returns the
  inner product of the source row `srcRow I J` and the source column `srcCol I J`.
-/

@[expose] public section

open Finset

namespace ThreeSumApsp.Batch

/-- The source block of the slot `(g, h)` of the band `β`. -/
def srcBlock (L m β g h : ℕ) : ℕ := (β * ThreeSumApsp.K0 L m + g) * ThreeSumApsp.K0 L m + h

/-- The left family of a band: the slot `(g, h)` holds the row block `srcBlock β g h`. -/
def bFamilyL {L m R : ℕ} (lay : Layout L m) (X : Matrix (Fin R) (Fin (ThreeSumApsp.D m)) ℤ) (β : ℕ) :
    Finset (Fin L) → LeftMat L m :=
  fun Q => ∑ gh, if lay.table gh = Q then rowBlock lay X (srcBlock L m β gh.1 gh.2) else 0

/-- The right family of a band: the slot `(g, h)` holds the column block `srcBlock β g h`. -/
def bFamilyR {L m R : ℕ} (lay : Layout L m) (Y : Matrix (Fin (ThreeSumApsp.D m)) (Fin R) ℤ) (β : ℕ) :
    Finset (Fin L) → RightMat L m :=
  fun Q => ∑ gh, if lay.table gh = Q then colBlock lay Y (srcBlock L m β gh.1 gh.2) else 0

/-- The left input array of a band. -/
def bArrayL {L m R : ℕ} (lay : Layout L m) (X : Matrix (Fin R) (Fin (ThreeSumApsp.D m)) ℤ) (β : ℕ) :
    LeftStr L → ℤ :=
  arrayL m (bFamilyL lay X β)

/-- The right input array of a band. -/
def bArrayR {L m R : ℕ} (lay : Layout L m) (Y : Matrix (Fin (ThreeSumApsp.D m)) (Fin R) ℤ) (β : ℕ) :
    RightStr L → ℤ :=
  arrayR m (bFamilyR lay Y β)

/-- The source row of the position `(I, J)` of a tile. -/
def srcRow {L m : ℕ} (lay : Layout L m) (I J : ℕ) : ℕ :=
  srcBlock L m (bandOf L m I) (blockOf lay I) (blockOf lay J) * ThreeSumApsp.N0 L m + I % ThreeSumApsp.N0 L m

/-- The source column of the position `(I, J)` of a tile. -/
def srcCol {L m : ℕ} (lay : Layout L m) (I J : ℕ) : ℕ :=
  srcBlock L m (bandOf L m J) (blockOf lay I) (blockOf lay J) * ThreeSumApsp.N0 L m + J % ThreeSumApsp.N0 L m

/-! ## Bounds -/

end ThreeSumApsp.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.BandArrayB
/-!
# The input array of a batched band

bandArrayB has the arguments of bandArray (Section 2.3.4) and fills the batched array of
`Batch/Tiling.lean`: the slot `(g, h)` of the band `β` holds the row block `srcBlock β g h =
(β K₀ + g) K₀ + h` of the source, on both sides.  The source has `K₀ N` rows of `D` entries, row by
row (for the right side it is the transpose of the `D × K₀ N` matrix); the strides among the
arguments are not read.  The number `N` of rows of a band is a multiple of the band size
(`padN = N`), so every source row exists.  The digits of the row of a block are read from the table
of digits at the row `r` itself (`r < N₀ ≤ N`).

The proof is that of bandArray, with these changes.  `blockNoB`, `rowNoB`: the number of the block
and of the row; `BandCtxB`: the surroundings; the loops over the levels and the coverage lemma
`coverLoop` are those of bandArray.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Batch

open Light.Sec2 Light.Sec2.Band

/-! ## The program -/

/-- One entry: the code of the string, and the entry copied from the source, row by row. -/
def bandCellB : Stmt :=
  .set CodeAcc (k 0) ;;
  .set Light.Sec2.Band.Outer (Light.v Light.Sec2.Band.Digits3 +' Light.v Light.Sec2.Band.Offset *' (Light.v Light.Sec2.Band.Levels -' Light.v Light.Sec2.Band.Size)) ;;
  .set Light.Sec2.Band.Inner (Light.v Light.Sec2.Band.Digits4 +' Light.v Light.Sec2.Band.Col *' Light.v Light.Sec2.Band.Size) ;;
  bandCodeLoop ;;
  .store (Light.v Light.Sec2.Band.Dest +' Light.v CodeAcc) (M (Light.v Light.Sec2.Band.Source +' Light.v Light.Sec2.Band.Row *' Light.v Light.Sec2.Band.Cols +' Light.v Light.Sec2.Band.Col))

/-- The loop over the columns k. -/
def bandColLoopB : Stmt := .for Light.Sec2.Band.Col (Light.v Light.Sec2.Band.Cols) bandCellB

/-- One row r of a block. -/
def bandRowB : Stmt :=
  .set Light.Sec2.Band.Row (Light.v BlockIdx *' Light.v Light.Sec2.Band.Width +' Light.v Light.Sec2.Band.Offset) ;;
  bandColLoopB

/-- The loop over the rows r of a block. -/
def bandRowLoopB : Stmt := .for Light.Sec2.Band.Offset (Light.v Light.Sec2.Band.Width) bandRowB

/-- The address of the mask of the subset of the slot (g, h), and the rows of the block. -/
def bandRowsB : Stmt :=
  .set MaskRow (Light.v Light.Sec2.Band.MaskTable +' (Light.v PosG *' Light.v Light.Sec2.Band.NumBlocks +' Light.v PosH) *' Light.v Light.Sec2.Band.Levels) ;; bandRowLoopB

/-- One slot (g, h): the number of the source block, the same on both sides. -/
def bandBlockB : Stmt :=
  .set BlockIdx ((Light.v Light.Sec2.Band.BandNo *' Light.v Light.Sec2.Band.NumBlocks +' Light.v PosG) *' Light.v Light.Sec2.Band.NumBlocks +' Light.v PosH) ;;
  bandRowsB

/-- The loop over h. -/
def bandHLoopB : Stmt := .for PosH (Light.v Light.Sec2.Band.NumBlocks) bandBlockB

/-- The loop over g. -/
def bandGLoopB : Stmt := .for PosG (Light.v Light.Sec2.Band.NumBlocks) bandHLoopB

/-- bandArrayB(L, m, N, D, K0, N0, S7, β, side, aA, sI, sK, mask, dig3, dig4, arr). -/
def bandArrayBBody : Stmt := bandZeroLoop ;; bandGLoopB

/-! ## The numbers -/

namespace BandArgsB
variable (A : BandArgs)

/-- The number of the source block of the slot `(g, h)`. -/
def blockNo (g h : ℕ) : ℕ := (A.β * A.p.K0 + g) * A.p.K0 + h

/-- The number of the source row for the slot `(g, h)` and the row `r` of the block. -/
def rowNo (g h r : ℕ) : ℕ := blockNo A g h * A.p.N0 + r

variable {A}

end BandArgsB

/-- The local variables of the loop over `h` for the slot `(g, h)`. -/
abbrev blockTmpB (A : BandArgs) (g h : ℕ) : BlockTmp :=
  ⟨h, (BandArgsB.blockNo A g h : ℕ), (A.maskRow g h : ℕ)⟩

/-- The local variables of the loop over the rows for the row `r` of the slot `(g, h)`. -/
abbrev rowTmpB (A : BandArgs) (g h r : ℕ) : RowTmp := ⟨r, (BandArgsB.rowNo A g h r : ℕ)⟩

variable {lim : Limits} {P : Program} {d : ℕ} {A : BandArgs} {μ₀ μ : ℕ → ℤ} {T : List ℤ}

/-! ## The surroundings -/

/-- Everything that the loops assume. -/
structure BandCtxB (lim : Limits) (A : BandArgs) (μ₀ : ℕ → ℤ) (T : List ℤ) : Prop where
  std : Std lim
  hmL : A.p.m ≤ A.p.L
  space : A.arr + A.p.S7 ≤ lim.space
  aA_le : A.aA + A.p.K0 * A.p.N * A.p.D ≤ A.arr
  tabs : BandTables A.p μ₀ A.mask A.dig3 A.dig4 A.arr
  β_lt : A.β < A.p.nB
  pad : padN A.p.L A.p.m A.p.N = A.p.N
  T_len : T.length = A.p.S7
  entry : ∀ g < A.p.K0, ∀ h < A.p.K0, ∀ r < A.p.N0, ∀ k < A.p.D, T.getD (A.code g h r k) 0
    = μ₀ (A.aA + (BandArgsB.rowNo A g h r * A.p.D + k))
  zero : ∀ c : ℕ, (∀ g < A.p.K0, ∀ h < A.p.K0, ∀ r < A.p.N0, ∀ k < A.p.D,
    c ≠ A.code g h r k) → T.getD c 0 = 0

namespace BandCtxB

end BandCtxB

/-! ## The loop over the columns -/

section cell
variable {g h r k : ℕ}

end cell

/-! ## The loop over the rows of a block -/

/-- The cells of the row r of the slot (g, h) are right. -/
def CoveredRowB (A : BandArgs) (T : List ℤ) (μ : ℕ → ℤ) (g h r : ℕ) : Prop :=
  ∀ k < A.p.D, Covered A T μ g h r k

section rows
variable {g h r : ℕ}

end rows

/-! ## The loops over h and g -/

section blocks
variable {g h : ℕ}

end blocks

/-! ## The whole procedure -/

/-! ## What the callers assume -/

/-- bandArrayB for the left side, at the arguments with which encodeBandsP calls it: the source `X`
(`K₀ N × D`, row by row) at aX.  Writes the batched array of the band β. -/
def BandArrayBLSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (β aX mask dig3 dig4 arr : ℕ) (μ : ℕ → ℤ)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ),
    arr + p.S7 ≤ lim.space → MatAt μ aX X → aX + p.K0 * p.N * p.D ≤ arr →
    BandTables p μ mask dig3 dig4 arr → β < p.nB → padN p.L p.m p.N = p.N →
    ∀ d, d ≤ lim.depth → Meets lim P pBandArray d
      [p.L, p.m, p.N, p.D, p.K0, p.N0, p.S7, β, 0, aX, p.D, 1, mask, dig3, dig4, arr] μ
      (c * bandArrayShape p) fun _ μ' =>
      Seg μ' arr (Spec.arrL (Batch.bArrayL (Spec.stdLayout hmL) X β)) ∧ SameOutside μ μ' arr p.S7

/-- The same for the right side: the source is the transpose of `Y` (`D × K₀ N`), that is, `Yᵀ`
row by row at aY. -/
def BandArrayBRSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (β aY mask dig3 dig4 arr : ℕ) (μ : ℕ → ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ),
    arr + p.S7 ≤ lim.space → MatAt μ aY Y.transpose → aY + p.K0 * p.N * p.D ≤ arr →
    BandTables p μ mask dig3 dig4 arr → β < p.nB → padN p.L p.m p.N = p.N →
    ∀ d, d ≤ lim.depth → Meets lim P pBandArray d
      [p.L, p.m, p.N, p.D, p.K0, p.N0, p.S7, β, 1, aY, 1, p.N, mask, dig3, dig4, arr] μ
      (c * bandArrayShape p) fun _ μ' =>
      Seg μ' arr (Spec.arrR (Batch.bArrayR (Spec.stdLayout hmL) Y β)) ∧ SameOutside μ μ' arr p.S7

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.Ready
/-!
# The invariants of Theorem 30 for arbitrary input arrays

`Sec2.SharedReadyP` and `Sec4.DSReady` name the input arrays of the bands by the matrices X and Y
(`bandArrayL (stdLayout hmL) X β`). Nothing in the routines of Section 4 depends on how the arrays
arise, so here the same invariants are stated for two families of arrays `AL β`, `AR β`, one for
each band, with the lemmas that the routines use. The batched data structure takes for them the
arrays `bArrayL`, `bArrayR` of `Batch.Tiling`.

* `SharedReadyA`: the shared block, with the pruned encodings of `AL β` and `AR β`.
* `DSReadyA`: the whole block of the data structure.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Spec Light.Sec4

/-- The input arrays of the row bands. -/
abbrev ArrL (p : Sec2.Par) : Type := ℕ → LeftStr p.L → ℤ
/-- The input arrays of the column bands. -/
abbrev ArrR (p : Sec2.Par) : Type := ℕ → RightStr p.L → ℤ

/-- What the shared stage leaves behind, for the input arrays `AL β` and `AR β`. -/
structure SharedReadyA (p : Sec2.Par) (aX aY b0 : ℕ) (AL : ArrL p) (AR : Light.Batch.ArrR p) (μ : ℕ → ℤ) :
    Prop extends Sec2.SharedTables p b0 μ where
  dir : SegN μ (p.aDIR b0) (Sec2.dirList p aX aY b0)
  encA : ∀ β < p.nB, Sec2.RetSeg p.L p.m μ (p.aENCA b0 + β * p.T) (arrT (encodingL (AL β)))
  encB : ∀ β < p.nB, Sec2.RetSeg p.L p.m μ (p.aENCB b0 + β * p.T) (arrT (encodingR (AR β)))

section

variable {p : Sec2.Par} {aX aY b0 : ℕ} {AL : ArrL p} {AR : ArrR p} {μ μ' : ℕ → ℤ}

end

/-! ## The data structure -/

/-- The encodings of the row bands. -/
def encRowA {p : Sec2.Par} (AL : ArrL p) (β : ℕ) : Leaf p.L → ℤ := encodingL (AL β)

/-- The encodings of the column bands. -/
def encColA {p : Sec2.Par} (AR : Light.Batch.ArrR p) (β : ℕ) : Leaf p.L → ℤ := encodingR (AR β)

/-- The tries of all tiles, and their roots. -/
def dsTriesA (p : Sec2.Par) (t : ℕ) (AL : ArrL p) (AR : Light.Batch.ArrR p) : TrieStore :=
  allTries p.L p.m t (tileList p.nB (encRowA AL) (encColA AR))

/-- The block from b0 on holds the data structure for the arrays `AL β`, `AR β`. -/
structure DSReadyA (p : Sec2.Par) (t aX aY b0 : ℕ) (AL : ArrL p) (AR : Light.Batch.ArrR p) (μ : ℕ → ℤ) :
    Prop where
  shared : SharedReadyA p aX aY b0 AL AR μ
  cellT : μ (b0 + 31) = t
  roots : SegN μ (aROOTS p b0) (dsTriesA p t AL AR).roots
  trie : Seg μ (aTR p b0) (dsTriesA p t AL AR).cells

section

variable {p : Sec2.Par} {t aX aY b0 : ℕ} {AL : ArrL p} {AR : ArrR p} {μ μ' : ℕ → ℤ}

end

/-- The arrays are bounded by U. -/
structure BoundA {p : Sec2.Par} (AL : ArrL p) (AR : Light.Batch.ArrR p) (U : ℤ) : Prop where
  nonneg : 0 ≤ U
  left : ∀ β u, |AL β u| ≤ U
  right : ∀ β v, |AR β v| ≤ U

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.SharedB
/-!
# The shared stage with batched band arrays

The procedure sharedP and its callees are unchanged, but for the body of bandArray, which is
`bandArrayBBody`. Then the shared stage leaves the pruned encodings of the batched arrays
`bArrayL (stdLayout hmL) X β` and `bArrayR (stdLayout hmL) Y β` behind, where the source X is
`K₀ N × D`, row by row at aX, and the source Y is `D × K₀ N`, stored as its transpose at aY.

* `EncodeBandsPLBSpec`, `EncodeBandsPRBSpec`: encodeBandsP for the batched arrays.
* `SharedCalleesB`, `SharedPreB`, `SharedSpecB`, `sharedPB_entry`: the shared stage.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp Light.Sec2 Light.Sec2.Shared

variable {lim : Limits} {P : Program}

/-! ## The encodings of all bands -/

/-- encodeBandsP for the batched row bands of the source X. -/
def EncodeBandsPLBSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (x : BandsArgs) (aX : ℕ) (μ : ℕ → ℤ)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ) (V : ℤ),
    BandsPre lim p μ x Spec.phiFlat V → MatAt μ aX X → aX + p.K0 * p.N * p.D ≤ x.enc →
    (∀ i j, |X i j| ≤ V) → padN p.L p.m p.N = p.N →
    ∀ d, d + (p.L + 2) ≤ lim.depth → Meets lim P pEncodeBandsP d (x.vals p 0 aX p.D 1) μ
      (c * (p.nB * (retW p.L p.m + bandArrayShape p) + 1)) fun _ μ' =>
        (∀ β < p.nB, RetSeg p.L p.m μ' (x.enc + β * p.T)
          (Spec.arrT (encodingL (ThreeSumApsp.Batch.bArrayL (Spec.stdLayout hmL) X β)))) ∧
          SameOutside μ μ' x.enc (x.zs + p.S7 - x.enc)

/-- encodeBandsP for the batched column bands of the source Y, stored as its transpose. -/
def EncodeBandsPRBSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (x : BandsArgs) (aY : ℕ) (μ : ℕ → ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ) (V : ℤ),
    BandsPre lim p μ x Spec.psiFlat V → MatAt μ aY Y.transpose → aY + p.K0 * p.N * p.D ≤ x.enc →
    (∀ i j, |Y i j| ≤ V) → padN p.L p.m p.N = p.N →
    ∀ d, d + (p.L + 2) ≤ lim.depth → Meets lim P pEncodeBandsP d (x.vals p 1 aY 1 p.N) μ
      (c * (p.nB * (retW p.L p.m + bandArrayShape p) + 1)) fun _ μ' =>
        (∀ β < p.nB, RetSeg p.L p.m μ' (x.enc + β * p.T)
          (Spec.arrT (encodingR (ThreeSumApsp.Batch.bArrayR (Spec.stdLayout hmL) Y β)))) ∧
          SameOutside μ μ' x.enc (x.zs + p.S7 - x.enc)

/-! ## The shared stage -/

/-- The entries of the routines that the shared stage calls, with the batched encodeBandsP. -/
structure SharedCalleesB (lim : Limits) (P : Program) (c : ℕ) : Prop where
  pow : PowSpec lim P c
  binom : BinomSpec lim P c
  sqrt : SqrtSpec lim P c
  coef : CoefSpec lim P c
  subsets : SubsetsSpec lim P c
  counters : CountersSpec lim P c
  digits : DigitsSpec lim P c
  bandsL : EncodeBandsPLBSpec lim P c
  bandsR : EncodeBandsPRBSpec lim P c

/-- What the shared stage assumes of the two sources. -/
structure SharedPreB (lim : Limits) (p : Par) (aX aY b0 : ℕ)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ) (V : ℤ) (μ : ℕ → ℤ) : Prop where
  space : p.sharedEnd b0 ≤ lim.space
  matX : MatAt μ aX X
  matY : MatAt μ aY Y.transpose
  belowX : aX + p.K0 * p.N * p.D ≤ b0
  belowY : aY + p.K0 * p.N * p.D ≤ b0
  absX : ∀ i j, |X i j| ≤ V
  absY : ∀ i j, |Y i j| ≤ V
  V_nonneg : 0 ≤ V
  seven : 7 ^ (p.L + 1) * V ≤ lim.word
  ten : ((10 ^ (p.L + 1) : ℕ) : ℤ) ≤ lim.word
  pad : padN p.L p.m p.N = p.N

/-- The batched input arrays of the row bands. -/
abbrev batchL (p : Par) (hmL : p.m ≤ p.L)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ) : ArrL p :=
  fun β => ThreeSumApsp.Batch.bArrayL (Spec.stdLayout hmL) X β

/-- The batched input arrays of the column bands. -/
abbrev batchR (p : Par) (hmL : p.m ≤ p.L)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ) : Light.Batch.ArrR p :=
  fun β => ThreeSumApsp.Batch.bArrayR (Spec.stdLayout hmL) Y β

/-- sharedP(L, m, N, D, aX, aY, b0), with the batched band arrays. -/
def SharedSpecB (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ) (μ : ℕ → ℤ)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ)
    (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ) (V : ℤ),
    SharedPreB lim p aX aY b0 X Y V μ →
    ∀ d, d + (p.L + 3) ≤ lim.depth →
    Meets lim P pSharedP d [p.L, p.m, p.N, p.D, aX, aY, b0] μ (c * sharedShapeP p) fun _ μ' =>
      SharedReadyA p aX aY b0 (batchL p hmL X) (batchR p hmL Y) μ' ∧
        SameOutside μ μ' b0 (p.sharedEnd b0 - b0)

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.QueryA
/-!
# A query of Theorem 30, for arbitrary input arrays

`queryAt` for the invariant `DSReadyA`: the proof of `Sec4.queryAt_spec` goes through as it is,
up to its last step. The value returned is the value of the query of Lemma 28 for the encodings of
the band of I and of the band of J (`queryValue`); Theorem 30 is not used.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Spec Light.Sec4 Light.Sec4.QueryAt

/-- queryAt(I, J, b0), for the data structure of the arrays `AL`, `AR`: it returns the value of the
query of the position (I, J) on the encodings of the band of I and the band of J. -/
def QueryAtASpec (lim : Limits) (P : Program) : Prop :=
  ∀ (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ) (AL : ArrL p) (AR : Light.Batch.ArrR p) (U : ℤ)
    (μ : ℕ → ℤ) (I J : Fin p.N),
    t ≤ p.m → Lim30 lim p t b0 U → BoundA AL AR U → DSReadyA p t aX aY b0 AL AR μ →
    ∀ d, d + 2 ≤ lim.depth →
    Meets lim P Proc.queryAt d [(I : ℕ), (J : ℕ), b0] μ (tQueryAt p.L p.m t) fun r μ' =>
      r = queryValue p.m t (encRowA AL ((I : ℕ) / (p.K0 * p.N0)))
          (encColA AR ((J : ℕ) / (p.K0 * p.N0))) (outStrOfPos (stdLayout hmL) I J) ∧
        DSReadyA p t aX aY b0 AL AR μ' ∧ SameOutside μ μ' (aWD p b0) (3 * p.L + p.m)

section Proof

variable {lim : Limits} {P : Program} {d : ℕ} {p : Sec2.Par} {t : ℕ} {hmL : p.m ≤ p.L}
  {aX aY b0 : ℕ} {AL : ArrL p} {AR : ArrR p} {U : ℤ} {μ : ℕ → ℤ}

/-- The arguments of the call of queryCore. -/
def coreArgsA (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (b0 : ℕ) (AL : ArrL p) (AR : Light.Batch.ArrR p)
    (μ : ℕ → ℤ) (I J : ℕ) : QueryCoreArgs where
  L := p.L
  m := p.m
  t := t
  encA := memRow p b0 μ (I / (p.K0 * p.N0))
  encB := memCol p b0 μ (J / (p.K0 * p.N0))
  η := outStrOfPos (stdLayout hmL) I J
  aA := p.aENCA b0 + I / (p.K0 * p.N0) * p.T
  aB := p.aENCB b0 + J / (p.K0 * p.N0) * p.T
  tr := aTR p b0
  root := (dsTriesA p t AL AR).root (tileOf p I J)
  wd := aWD p b0
  ss := aSS p b0
  box := aBOX p b0
  T := (dsTriesA p t AL AR).cells

section Tile

variable (I J : Fin p.N)

variable (ht : t ≤ p.m)

include ht

end Tile

end Proof

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.Program
/-!
# The batched data structure on the word RAM

`base58B` is `base58P` with the body `bandArrayBBody` at the number of bandArray. Every program that
begins with it meets the specifications of the preprocessing and of a query for the batched arrays:

* `preCoreB_base58B`: preCore(L, m, t, N, D, aX, aY, b0), for the source X (`K₀ N × D`) at aX and
  the transpose of the source Y at aY, builds the data structure `DSReadyA` of the arrays
  `bArrayL (stdLayout hmL) X β`, `bArrayR (stdLayout hmL) Y β`;
* `queryAtA_base58B`: queryAt(I, J, b0) returns the value of the query of (I, J) from it.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Spec Light.Sec4

/-- The procedures number 0 to 57, with the batched band array. -/
def base58B : Program := base58P.set Sec2.pBandArray bandArrayBBody

/-! ## The preprocessing -/

/-- preCore(L, m, t, N, D, aX, aY, b0) for the two sources: X (`K₀ N × D`, row by row at aX) and Y
(`D × K₀ N`, its transpose row by row at aY). -/
def PreCoreBSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (aX aY b0 : ℕ)
    (X : Matrix (Fin (p.K0 * p.N)) (Fin (ThreeSumApsp.D p.m)) ℤ) (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin (p.K0 * p.N)) ℤ)
    (U : ℤ) (μ : ℕ → ℤ),
    t ≤ p.m → Lim30 lim p t b0 U → 0 ≤ U → (∀ i j, |X i j| ≤ U) → (∀ i j, |Y i j| ≤ U) →
    MatAt μ aX X → MatAt μ aY Y.transpose → aX + p.K0 * p.N * p.D ≤ b0 →
    aY + p.K0 * p.N * p.D ≤ b0 → padN p.L p.m p.N = p.N →
    ∀ d, d + (p.L + 6) ≤ lim.depth →
    Meets lim P Proc.preCore d [p.L, p.m, t, p.N, p.D, aX, aY, b0] μ
      (tPreCoreP c p t) fun _ μ' =>
      DSReadyA p t aX aY b0 (batchL p hmL X) (batchR p hmL Y) μ' ∧
        SameOutside μ μ' b0 (Light.Sec4.top p t b0 - b0)

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.OParams
/-!
# The batched counting oracle: its parameters and its time

For `n` vertices, the block size `s` and the modulus `p ≥ 1`:

* `m = ⌈log₄ (4 s p)⌉ ≥ 1`, `D = 4^m ≥ 4 s p`, `L = 20 m`, `t = ⌈θ m⌉` with `θ = 56569/500000`;
* `N₀ = 3^{L-m}`, `K₀ = ⌊√binom(L, m)⌋`, `KK = K₀²` (the products of a batch);
* `k = ⌊√(p-1)⌋ + 1` (so `k² ≥ p`, `k ≤ p`), the `n k` digit rows of a product;
* `nb = ⌈n k / N₀⌉ + 1` blocks of `N₀` rows per product, `NP = nb N₀` rows, and the virtual size
  `Nv = K₀ NP` of Theorem 30;
* `T = ⌈n/s⌉` products, in `B = ⌈T / KK⌉` batches.

`oTimeF A n U s p` is the time of the oracle, with the constant `A` for all the bookkeeping.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp Light.Sec4

/-- `m = ⌈log₄ (4 s p)⌉`. -/
def oM (s p : ℕ) : ℕ := Nat.clog 4 (4 * s * p)
/-- `L = 20 m`. -/
def oL (s p : ℕ) : ℕ := 20 * oM s p
/-- `t = ⌈θ m⌉`, `θ = 56569/500000`. -/
def oT (s p : ℕ) : ℕ := (56569 * oM s p + 500000 - 1) / 500000
/-- `D = 4^m`. -/
def oD (s p : ℕ) : ℕ := 4 ^ oM s p
/-- `N₀ = 3^{L-m}`. -/
def oN0 (s p : ℕ) : ℕ := 3 ^ (oL s p - oM s p)
/-- `K₀ = ⌊√binom(L, m)⌋`. -/
def oK0 (s p : ℕ) : ℕ := Nat.sqrt ((oL s p).choose (oM s p))
/-- `KK = K₀²`. -/
def oKK (s p : ℕ) : ℕ := oK0 s p * oK0 s p
/-- `k = ⌊√(p-1)⌋ + 1`. -/
def oDig (p : ℕ) : ℕ := Nat.sqrt (p - 1) + 1
/-- `nb = ⌈n k / N₀⌉ + 1`. -/
def oNb (n s p : ℕ) : ℕ := (n * oDig p) ⌈/⌉ oN0 s p + 1
/-- `NP = nb N₀`. -/
def oNP (n s p : ℕ) : ℕ := oNb n s p * oN0 s p
/-- `Nv = K₀ NP`. -/
def oNv (n s p : ℕ) : ℕ := oK0 s p * oNP n s p
/-- The number of products, `T = ⌈n/s⌉`. -/
def oProd (n s : ℕ) : ℕ := n ⌈/⌉ s
/-- The number of batches, `B = ⌈T / KK⌉`. -/
def oBat (n s p : ℕ) : ℕ := oProd n s ⌈/⌉ oKK s p
/-- The parameters of Theorem 30 for a batch. -/
def oPar (n s p : ℕ) : Sec2.Par := ⟨oL s p, oM s p, oNv n s p⟩

/-- The time of the oracle, with the constant `A` for the bookkeeping: the tables, the parameters,
writing the sources of each batch, the loops of the queries; the preprocessing of each batch, and
one query for each pair and each product. -/
def oTimeF (A n U s p : ℕ) : ℕ :=
  if p = 0 then 0 else
    A * (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)
      + oBat n s p * tPreCoreP cShared30 (oPar n s p) (oT s p)
      + oBat n s p * (oKK s p * (n * n)) * tQueryAt (oL s p) (oM s p) (oT s p)

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.OProg
/-!
# The batched counting oracle: the program text

The program `oProg` is `base58B` (Theorem 30 with the batched band arrays, numbers 0 to 57), the
procedures 58 to 73 of Section 4.4 for the rational parameters `c = 20`, `θ = 56569/500000`
(`procs31 G20`: copy, fill, log4, …, and `⌈θ m⌉` at 73), and the procedures of the oracle from 74 on:

* 74 emod, 75 ceilDiv, 76 powNat(b, e) = b^e;
* 77 residTab(dst, src, len, sgn, p, fr): `dst[q] := (sgn · src[q]) mod p` for `q < len`;
* 78 smallTab(dst, len, mul, add, p, fr): `dst[h] := (h · mul + add) mod p` for `h < len`;
* 79 rowTab(dst, nb, N0, Mul): `dst[i N0 + o] := i · Mul + o` for `i < nb`, `o < N0`;
* 80 divTab(dst, k): `dst[h k + l] := h` for `h, l < k`;
* 81 zTab(zh, zl, rn, dk, len, k): `zh[q] := dk[rn[q]]`, then `zl[q] := rn[q] - k · zh[q]`;
* 82 writeSrc: one source of a batch;
* 83 queries: the queries of a batch;
* 84 oracle: the top procedure.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp Light.Sec4

/-! ## The procedure numbers -/

abbrev pEmod : ℕ := 74
abbrev pCeilDiv : ℕ := 75
abbrev pPowNat : ℕ := 76
abbrev pResidTab : ℕ := 77
abbrev pSmallTab : ℕ := 78
abbrev pRowTab : ℕ := 79
abbrev pDivTab : ℕ := 80
abbrev pZTab : ℕ := 81
abbrev pWriteSrc : ℕ := 82
abbrev pQueries : ℕ := 83
abbrev pOracle : ℕ := 84

/-! ## Small procedures -/

/-- powNat(b, e): locals 2 (counter) and 3 (the power). -/
def powNatBody : Stmt :=
  .set 3 (k 1) ;;
  .for 2 (Light.v 1) (.set 3 (Light.v 3 *' Light.v 0)) ;;
  .set 0 (Light.v 3)

/-- residTab(dst, src, len, sgn, p, fr): local 6 is the counter, 7 the residue. -/
def residTabBody : Stmt :=
  .for 6 (Light.v 2) (
    .call pEmod [Light.v 3 *' M (Light.v 1 +' Light.v 6), Light.v 4, Light.v 5] 7 ;;
    .store (Light.v 0 +' Light.v 6) (Light.v 7))

/-- smallTab(dst, len, mul, add, p, fr): local 6 is the counter, 7 the residue. -/
def smallTabBody : Stmt :=
  .for 6 (Light.v 1) (
    .call pEmod [Light.v 6 *' Light.v 2 +' Light.v 3, Light.v 4, Light.v 5] 7 ;;
    .store (Light.v 0 +' Light.v 6) (Light.v 7))

/-- rowTab(dst, nb, N0, Mul): locals 4 and 5 are the counters i and o. -/
def rowTabBody : Stmt :=
  .for 4 (Light.v 1) (.for 5 (Light.v 2) (.store (Light.v 0 +' Light.v 4 *' Light.v 2 +' Light.v 5) (Light.v 4 *' Light.v 3 +' Light.v 5)))

/-- divTab(dst, k): locals 2 and 3 are the counters h and l. -/
def divTabBody : Stmt :=
  .for 2 (Light.v 1) (.for 3 (Light.v 1) (.store (Light.v 0 +' Light.v 2 *' Light.v 1 +' Light.v 3) (Light.v 2)))

/-- zTab(zh, zl, rn, dk, len, k): local 6 is the counter. -/
def zTabBody : Stmt :=
  .for 6 (Light.v 4) (.store (Light.v 0 +' Light.v 6) (M (Light.v 3 +' M (Light.v 2 +' Light.v 6)))) ;;
  .for 6 (Light.v 4) (.store (Light.v 1 +' Light.v 6) (M (Light.v 2 +' Light.v 6) -' Light.v 5 *' M (Light.v 0 +' Light.v 6)))

/-! ## One source of a batch -/

namespace WS

/-- The locals of writeSrc: the arguments base, T1, T2, RS, n, k, p, D, N0, j0, KKb, s, sz; then
the counters u, a, l, cl, the first vertex c0 = (j0 + u) s of the middle block, its length len, a
temporary, and the address of the row. -/
abbrev Base : ℕ := 0
@[inherit_doc Base] abbrev T1 : ℕ := 1
@[inherit_doc Base] abbrev T2 : ℕ := 2
@[inherit_doc Base] abbrev RS : ℕ := 3
@[inherit_doc Base] abbrev N : ℕ := 4
@[inherit_doc Base] abbrev Kd : ℕ := 5
@[inherit_doc Base] abbrev Pr : ℕ := 6
@[inherit_doc Base] abbrev Dd : ℕ := 7
@[inherit_doc Base] abbrev N0 : ℕ := 8
@[inherit_doc Base] abbrev J0 : ℕ := 9
@[inherit_doc Base] abbrev KKb : ℕ := 10
@[inherit_doc Base] abbrev S : ℕ := 11
@[inherit_doc Base] abbrev Sz : ℕ := 12
@[inherit_doc Base] abbrev U : ℕ := 13
@[inherit_doc Base] abbrev A : ℕ := 14
@[inherit_doc Base] abbrev Lc : ℕ := 15
@[inherit_doc Base] abbrev Cl : ℕ := 16
@[inherit_doc Base] abbrev C0 : ℕ := 17
@[inherit_doc Base] abbrev Len : ℕ := 18
@[inherit_doc Base] abbrev Tmp : ℕ := 19
@[inherit_doc Base] abbrev Row : ℕ := 20

end WS

open WS in
/-- One cell: `tmp := T1[a n + c0 + cl] + T2[l]`, reduced mod p; `row[cl p + tmp] := 1`. -/
def writeSrcCell : Stmt :=
  .set Tmp (M (Light.v T1 +' Light.v Light.Batch.WS.A *' Light.v Light.Batch.WS.N +' Light.v Light.Batch.WS.C0 +' Light.v Cl) +' M (Light.v T2 +' Light.v Lc)) ;;
  .ite (Light.v Tmp <' Light.v Pr) .skip (.set Tmp (Light.v Tmp -' Light.v Pr)) ;;
  .store (Light.v Light.Batch.WS.Row +' Light.v Cl *' Light.v Pr +' Light.v Tmp) (k 1)

open WS in
/-- The row `R = a k + l` of the product `u`: its address `base + (RS[R] + u N0) D`, then its cells. -/
def writeSrcRow : Stmt :=
  .set Light.Batch.WS.Row (Light.v Base +' (M (Light.v RS +' Light.v Light.Batch.WS.A *' Light.v Kd +' Light.v Lc) +' Light.v Light.Batch.WS.U *' Light.v N0) *' Light.v Dd) ;;
  .for Cl (Light.v Light.Batch.WS.Len) writeSrcCell

open WS in
/-- The product `u` of the batch: its middle block, then its rows. -/
def writeSrcProd : Stmt :=
  .set Light.Batch.WS.C0 ((Light.v J0 +' Light.v Light.Batch.WS.U) *' Light.v Light.Batch.WS.S) ;;
  .ite (Light.v Light.Batch.WS.C0 +' Light.v Light.Batch.WS.S ≤' Light.v Light.Batch.WS.N) (.set Light.Batch.WS.Len (Light.v Light.Batch.WS.S)) (.set Light.Batch.WS.Len (Light.v Light.Batch.WS.N -' Light.v Light.Batch.WS.C0)) ;;
  .for Light.Batch.WS.A (Light.v Light.Batch.WS.N) (.for Lc (Light.v Kd) writeSrcRow)

open WS in
/-- writeSrc(base, T1, T2, RS, n, k, p, D, N0, j0, KKb, s, sz): clear the `sz` cells, then mark the
cells of the products `u < KKb`. -/
def writeSrcBody : Stmt :=
  .call Proc.fill [Light.v Base, Light.v Sz, k 0] Tmp ;;
  .for Light.Batch.WS.U (Light.v KKb) writeSrcProd

/-! ## The queries of a batch -/

namespace QS

/-- The locals of queries: the arguments b0, ZL, ZH, RV, n, k, N0, K0, KKb, cnt, j0; then the
slot g, h, the counters a, b, the two positions I, J, the answer, the place of the counts of the
product, and the counter u of the products. -/
abbrev B0 : ℕ := 0
@[inherit_doc B0] abbrev ZL : ℕ := 1
@[inherit_doc B0] abbrev ZH : ℕ := 2
@[inherit_doc B0] abbrev RV : ℕ := 3
@[inherit_doc B0] abbrev N : ℕ := 4
@[inherit_doc B0] abbrev Kd : ℕ := 5
@[inherit_doc B0] abbrev N0 : ℕ := 6
@[inherit_doc B0] abbrev K0 : ℕ := 7
@[inherit_doc B0] abbrev KKb : ℕ := 8
@[inherit_doc B0] abbrev Cnt : ℕ := 9
@[inherit_doc B0] abbrev J0 : ℕ := 10
@[inherit_doc B0] abbrev G : ℕ := 11
@[inherit_doc B0] abbrev H : ℕ := 12
@[inherit_doc B0] abbrev A : ℕ := 13
@[inherit_doc B0] abbrev B : ℕ := 14
@[inherit_doc B0] abbrev I : ℕ := 15
@[inherit_doc B0] abbrev J : ℕ := 16
@[inherit_doc B0] abbrev Res : ℕ := 17
@[inherit_doc B0] abbrev Dst : ℕ := 18
@[inherit_doc B0] abbrev Uq : ℕ := 19

end QS

open QS in
/-- The query of the pair `(a, b)` in the slot `(g, h)`. -/
def queryCell : Stmt :=
  .set I (M (Light.v RV +' Light.v Light.Batch.QS.A *' Light.v Kd +' M (Light.v ZL +' Light.v Light.Batch.QS.A *' Light.v Light.Batch.QS.N +' Light.v Light.Batch.QS.B)) +' Light.v G *' Light.v N0) ;;
  .set Light.Batch.QS.J (M (Light.v RV +' Light.v Light.Batch.QS.B *' Light.v Kd +' M (Light.v ZH +' Light.v Light.Batch.QS.A *' Light.v Light.Batch.QS.N +' Light.v Light.Batch.QS.B)) +' Light.v H *' Light.v N0) ;;
  .call Proc.queryAt [Light.v I, Light.v Light.Batch.QS.J, Light.v B0] Light.Batch.QS.Res ;;
  .store (Light.v Light.Batch.QS.Dst +' Light.v Light.Batch.QS.A *' Light.v Light.Batch.QS.N +' Light.v Light.Batch.QS.B) (Light.v Light.Batch.QS.Res)

open QS in
/-- The product `u = g K₀ + h` of the batch: the queries of all pairs, then the next slot. -/
def queryProd : Stmt :=
  .set Light.Batch.QS.Dst (Light.v Light.Batch.QS.Cnt +' (Light.v J0 +' Light.v Uq) *' (Light.v Light.Batch.QS.N *' Light.v Light.Batch.QS.N)) ;;
  .for Light.Batch.QS.A (Light.v Light.Batch.QS.N) (.for Light.Batch.QS.B (Light.v Light.Batch.QS.N) queryCell) ;;
  .set H (Light.v H +' k 1) ;;
  .ite (Light.v H =' Light.v K0) (.set H (k 0) ;; .set G (Light.v G +' k 1)) .skip

open QS in
/-- queries(b0, ZL, ZH, RV, n, k, N0, K0, KKb, cnt, j0): the products `u < KKb` of the batch, in the
slots `(g, h)`, `u = g K₀ + h`. -/
def queriesBody : Stmt :=
  .set G (k 0) ;; .set H (k 0) ;;
  .for Uq (Light.v KKb) queryProd

/-! ## The top procedure -/

namespace OS

/-- The locals of the oracle: the arguments n, U, ab, bc, ac, s, p, cnt, fr; the parameters m, D,
L, t, N0, K0, KK, k, nb, NP, Nv, T, B; the addresses of the tables RAC, RNBC, RNAB, NL, HK, RS,
RV, DK, ZH, ZL, of the two sources and of the data structure; a temporary; the batch, its first
product j0 and its number KKb of products; the rows and the cells of a source; binom(L, m). -/
abbrev Nn : ℕ := 0

@[inherit_doc Nn] abbrev Ab : ℕ := 2
@[inherit_doc Nn] abbrev Bc : ℕ := 3
@[inherit_doc Nn] abbrev Ac : ℕ := 4
@[inherit_doc Nn] abbrev S : ℕ := 5
@[inherit_doc Nn] abbrev Pr : ℕ := 6
@[inherit_doc Nn] abbrev Cnt : ℕ := 7
@[inherit_doc Nn] abbrev Fr : ℕ := 8
@[inherit_doc Nn] abbrev Mm : ℕ := 9
@[inherit_doc Nn] abbrev Dd : ℕ := 10
@[inherit_doc Nn] abbrev Ll : ℕ := 11
@[inherit_doc Nn] abbrev Tt : ℕ := 12
@[inherit_doc Nn] abbrev N0 : ℕ := 13
@[inherit_doc Nn] abbrev K0 : ℕ := 14
@[inherit_doc Nn] abbrev KK : ℕ := 15
@[inherit_doc Nn] abbrev Kd : ℕ := 16
@[inherit_doc Nn] abbrev Nb : ℕ := 17
@[inherit_doc Nn] abbrev NP : ℕ := 18
@[inherit_doc Nn] abbrev Nv : ℕ := 19
@[inherit_doc Nn] abbrev Tp : ℕ := 20
@[inherit_doc Nn] abbrev Bt : ℕ := 21
@[inherit_doc Nn] abbrev aRAC : ℕ := 22
@[inherit_doc Nn] abbrev aRNBC : ℕ := 23
@[inherit_doc Nn] abbrev aRNAB : ℕ := 24
@[inherit_doc Nn] abbrev aNL : ℕ := 25
@[inherit_doc Nn] abbrev aHK : ℕ := 26
@[inherit_doc Nn] abbrev aRS : ℕ := 27
@[inherit_doc Nn] abbrev aRV : ℕ := 28
@[inherit_doc Nn] abbrev aDK : ℕ := 29
@[inherit_doc Nn] abbrev aZH : ℕ := 30
@[inherit_doc Nn] abbrev aZL : ℕ := 31
@[inherit_doc Nn] abbrev aX : ℕ := 32
@[inherit_doc Nn] abbrev aY : ℕ := 33
@[inherit_doc Nn] abbrev B0 : ℕ := 34
@[inherit_doc Nn] abbrev Tmp : ℕ := 35
@[inherit_doc Nn] abbrev Bn : ℕ := 36
@[inherit_doc Nn] abbrev J0 : ℕ := 37
@[inherit_doc Nn] abbrev KKb : ℕ := 38
@[inherit_doc Nn] abbrev Rows : ℕ := 39
@[inherit_doc Nn] abbrev Sz : ℕ := 40
@[inherit_doc Nn] abbrev Kb : ℕ := 41

end OS

open OS in
/-- The parameters. -/
def oParams : Stmt :=
  .call Proc.log4 [(k 4 *' Light.v Light.Batch.OS.S) *' Light.v Pr, Light.v Fr] Mm ;;
  .set Dd (M (Light.v Fr)) ;;
  .set Ll (k 20 *' Light.v Mm) ;;
  .call Proc.switch31 [Light.v Mm] Tt ;;
  .call pPowNat [k 3, Light.v Ll -' Light.v Mm] N0 ;;
  .call Sec2.pBinom [Light.v Ll, Light.v Mm, Light.v Fr +' k 1] Kb ;;
  .call Sec2.pSqrt [Light.v Kb] K0 ;;
  .set KK (Light.v K0 *' Light.v K0) ;;
  .call Sec2.pSqrt [Light.v Pr -' k 1] Kd ;;
  .set Kd (Light.v Kd +' k 1) ;;
  .call pCeilDiv [Light.v Nn *' Light.v Kd, Light.v N0] Nb ;;
  .set Nb (Light.v Nb +' k 1) ;;
  .set NP (Light.v Nb *' Light.v N0) ;;
  .set Nv (Light.v K0 *' Light.v NP) ;;
  .call pCeilDiv [Light.v Nn, Light.v Light.Batch.OS.S] Tp ;;
  .call pCeilDiv [Light.v Tp, Light.v KK] Bt

open OS in
/-- The addresses: the tables, the two sources, the data structure. -/
def oAddrs : Stmt :=
  .set aRAC (Light.v Fr) ;;
  .set aRNBC (Light.v aRAC +' Light.v Nn *' Light.v Nn) ;;
  .set aRNAB (Light.v aRNBC +' Light.v Nn *' Light.v Nn) ;;
  .set aNL (Light.v aRNAB +' Light.v Nn *' Light.v Nn) ;;
  .set aHK (Light.v aNL +' Light.v Kd) ;;
  .set aRS (Light.v aHK +' Light.v Kd) ;;
  .set aRV (Light.v aRS +' Light.v NP) ;;
  .set aDK (Light.v aRV +' Light.v NP) ;;
  .set aZH (Light.v aDK +' Light.v Kd *' Light.v Kd) ;;
  .set aZL (Light.v aZH +' Light.v Nn *' Light.v Nn) ;;
  .set Light.Batch.OS.Rows (Light.v KK *' Light.v NP) ;;
  .set Sz (Light.v Light.Batch.OS.Rows *' Light.v Dd) ;;
  .set Light.Batch.OS.aX (Light.v aZL +' Light.v Nn *' Light.v Nn) ;;
  .set Light.Batch.OS.aY (Light.v Light.Batch.OS.aX +' Light.v Sz) ;;
  .set B0 (Light.v Light.Batch.OS.aY +' Light.v Sz)

open OS in
/-- The tables. -/
def oTables : Stmt :=
  .call pResidTab [Light.v aRAC, Light.v Ac, Light.v Nn *' Light.v Nn, k 1, Light.v Pr, Light.v B0] Tmp ;;
  .call pResidTab [Light.v aRNBC, Light.v Bc, Light.v Nn *' Light.v Nn, k 0 -' k 1, Light.v Pr, Light.v B0] Tmp ;;
  .call pResidTab [Light.v aRNAB, Light.v Ab, Light.v Nn *' Light.v Nn, k 0 -' k 1, Light.v Pr, Light.v B0] Tmp ;;
  .call pSmallTab [Light.v aNL, Light.v Kd, k 0 -' k 1, Light.v Pr, Light.v Pr, Light.v B0] Tmp ;;
  .call pSmallTab [Light.v aHK, Light.v Kd, Light.v Kd, k 0, Light.v Pr, Light.v B0] Tmp ;;
  .call pRowTab [Light.v aRS, Light.v Nb, Light.v N0, Light.v KK *' Light.v N0] Tmp ;;
  .call pRowTab [Light.v aRV, Light.v Nb, Light.v N0, Light.v K0 *' Light.v N0] Tmp ;;
  .call pDivTab [Light.v aDK, Light.v Kd] Tmp ;;
  .call pZTab [Light.v aZH, Light.v aZL, Light.v aRNAB, Light.v aDK, Light.v Nn *' Light.v Nn, Light.v Kd] Tmp

open OS in
/-- One batch: its products, the two sources, the preprocessing, the queries. -/
def oBatch : Stmt :=
  .set J0 (Light.v Bn *' Light.v KK) ;;
  .ite (Light.v J0 +' Light.v KK ≤' Light.v Tp) (.set KKb (Light.v KK)) (.set KKb (Light.v Tp -' Light.v J0)) ;;
  .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

open OS in
/-- oracle(n, U, ab, bc, ac, s, p, cnt, fr). -/
def oracleBody : Stmt :=
  oParams ;; oAddrs ;; oTables ;; .for Bn (Light.v Bt) oBatch

/-! ## The program -/

/-- The rational parameters `c = 20`, `θ = 56569/500000`. -/
def G20 : RatParams where
  a := 20
  b := 1
  p := 56569
  q := 500000
  m₀ := 1
  hb := by norm_num
  hq := by norm_num
  hc := by norm_num
  hp := by norm_num
  hθ := by norm_num
  hm₀ := by norm_num

/-- The procedures of the oracle, numbers 74 to 84. -/
def oProcs : Program :=
  [emodBody, Sec3.ceilDivBody, powNatBody, residTabBody, smallTabBody, rowTabBody, divTabBody,
    zTabBody, writeSrcBody, queriesBody, Light.Batch.oracleBody]

/-- **The program of the batched counting oracle.** -/
def oProg : Program := base58B ++ procs31 G20 ++ oProcs

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.Lift
/-!
# The digit-lifted sources and the count identity

The sources of a batch are 0/1 matrices with at most one entry 1 in each inner block of `p` columns.
A source row `ρ = (i K₀² + u) N₀ + o` belongs to the product `u` (`rowU`) and is its row
`R = i N₀ + o` (`rowR`).  The inner index `κ = cl p + r` is the vertex `cl` of the middle block and
the residue `r`.  The entry is 1 iff the triple `(u, R, cl)` is valid (`ok`) and `r` is the residue
`f u R cl`.

* `sum_srcVal`: the inner product of a row of the left source and a column of the right source of the
  same product counts the `cl` where both are valid and the two residues agree.
* `lift_count`: with the residues of the digit construction (Lemma 4.1 of the note), the residues
  agree iff `p` divides the triangle sum.
-/

@[expose] public section

open Finset

namespace ThreeSumApsp.Batch

/-- The product of a source row. -/
def rowU (KK N0 ρ : ℕ) : ℕ := ρ / N0 % KK

/-- The row of its product that a source row is. -/
def rowR (KK N0 ρ : ℕ) : ℕ := ρ / N0 / KK * N0 + ρ % N0

/-- The source row of the row `R` of the product `u`. -/
def srcOf (KK N0 u R : ℕ) : ℕ := (R / N0 * KK + u) * N0 + R % N0

/-- The entry `(ρ, κ)` of a source. -/
def srcVal (s p KK N0 : ℕ) (ok : ℕ → ℕ → ℕ → Bool) (f : ℕ → ℕ → ℕ → ℕ) (ρ κ : ℕ) : ℤ :=
  if κ / p < s ∧ ok (rowU KK N0 ρ) (rowR KK N0 ρ) (κ / p) = true ∧
      κ % p = f (rowU KK N0 ρ) (rowR KK N0 ρ) (κ / p) then 1 else 0

/-! ## The residues of the digit construction -/

/-- The residue of an integer, as a natural number. -/
def resid (p : ℕ) (w : ℤ) : ℕ := (w % (p : ℤ)).toNat

/-! ## The sources of the digit construction -/

section Digits

variable (n s p k T KK bt : ℕ) (AB BC AC : List ℤ)

/-- The product `u` of the batch `bt` is the product of the middle block `bt K₀² + u`. -/
abbrev jOf (u : ℕ) : ℕ := bt * KK + u

/-- The left source: row `R = a k + l` of the product of the middle block `j`, vertex `cl` of the
block; the residue is `(w(a, c) - l) mod p`. -/
def okX (u R cl : ℕ) : Bool := decide (ThreeSumApsp.Batch.jOf KK bt u < T ∧ R / k < n ∧ ThreeSumApsp.Batch.jOf KK bt u * s + cl < n)

def fX (u R cl : ℕ) : ℕ :=
  (ThreeSumApsp.Batch.resid p (AC.getD (R / k * n + (ThreeSumApsp.Batch.jOf KK bt u * s + cl)) 0) + (p - R % k) % p) % p

/-- The right source: column `C = b k + h`; the residue is `(h k - w(b, c)) mod p`. -/
def okY (u C cl : ℕ) : Bool := decide (ThreeSumApsp.Batch.jOf KK bt u < T ∧ C / k < n ∧ ThreeSumApsp.Batch.jOf KK bt u * s + cl < n)

def fY (u C cl : ℕ) : ℕ :=
  (C % k * k % p + ThreeSumApsp.Batch.resid p (-BC.getD (C / k * n + (ThreeSumApsp.Batch.jOf KK bt u * s + cl)) 0)) % p

end Digits

/-! ## The virtual positions of a query -/

/-- The virtual position of the row `R` of the slot row `g`. -/
def virt (K0 N0 g R : ℕ) : ℕ := R / N0 * (K0 * N0) + g * N0 + R % N0

end ThreeSumApsp.Batch
end
end

section
-- Source module: ThreeSumApsp.Filter2.Defs
/-!
# Deterministic filtering: parameters and the counting oracle

The filtering reduction of the 2.996001 note (Lemma 3.1, Lemma 4.1).  A candidate is a triple
`(a, b, j)`: the pair `(a, b)` and the block `C_j = [j s, min((j+1) s, n))` of third vertices.  It
lies at index `i = j n² + a n + b`; its group is the pair, `i % (n n) = a n + b`.  For a modulus
`p ≥ 1` its count is the number of `c ∈ C_j` with `p ∣ w(a,b) + w(b,c) + w(a,c)`.

* The parameters, as functions of `n` and `U` (§3, (3.5)): `kOf`, `ellOf`, `RR`, `bOf`, `dOf`,
  `roundsOf`, `sOf`, `mOf`, `cKOf`, `rOf`, `BOf`, `hOf`, `pool`, `QOf`, `Guard`, `PmaxOf`.
* The counts (`cntList`) and the task of the counting oracle (`oracleTask`).
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Filter2

open ThreeSumApsp.Spec Light.Sec3

/-! ## Parameters -/

/-- `k`, the bit length of `n`: `n < 2^k`. -/
def kOf (n : ℕ) : ℕ := Nat.size n

/-- `ℓ = ⌊√k⌋ + 1`. -/
def ellOf (n : ℕ) : ℕ := Nat.sqrt (Light.Filter2.kOf n) + 1

/-- `R = 2^ℓ`; the pool is the primes in `[R, 2R)`. -/
def RR (n : ℕ) : ℕ := 2 ^ Light.Filter2.ellOf n

/-- `b = ℓ (⌊√ℓ⌋ + 1)`, about `k^{3/4}`. -/
def bOf (n : ℕ) : ℕ := Light.Filter2.ellOf n * (Nat.sqrt (Light.Filter2.ellOf n) + 1)

/-- `d = ⌊b/2⌋`: a round with the guard shrinks the active set by `2^d`. -/
def dOf (n : ℕ) : ℕ := Light.Filter2.bOf n / 2

/-- The number of rounds, `⌊3k/d⌋ + 1`, so that `d · rounds > 3k`. -/
def roundsOf (n : ℕ) : ℕ := 3 * Light.Filter2.kOf n / Light.Filter2.dOf n + 1

/-- The block size `s = 2^{⌊σ k⌋}`, `σ = 0.0336175`. -/
def sOf (n : ℕ) : ℕ := 2 ^ (336175 * Light.Filter2.kOf n / 10 ^ 7)

/-- The number of blocks, `⌈n/s⌉`. -/
def mOf (n : ℕ) : ℕ := (n + Light.Filter2.sOf n - 1) / Light.Filter2.sOf n

/-- `c_K`, the bit length of the largest witness count `s`. -/
def cKOf (n : ℕ) : ℕ := Nat.size (Light.Filter2.sOf n)

/-- The number of prime factors per round, `⌊(c_K + b)/ℓ⌋ + 1`. -/
def rOf (n : ℕ) : ℕ := (Light.Filter2.cKOf n + Light.Filter2.bOf n) / Light.Filter2.ellOf n + 1

/-- `B`, the bit length of `3U`: every triangle sum has absolute value below `2^B`. -/
def BOf (U : ℕ) : ℕ := Nat.size (3 * U)

/-- `h = ⌊(B − 1)/ℓ⌋`: a nonzero integer below `2^B` in absolute value has at most `h` prime
divisors in the pool. -/
def hOf (n U : ℕ) : ℕ := (Light.Filter2.BOf U - 1) / Light.Filter2.ellOf n

/-- The pool: the primes in `[R, 2R)`. -/
def pool (n : ℕ) : Finset ℕ := (Finset.range (2 * Light.Filter2.RR n)).filter fun q => Light.Filter2.RR n ≤ q ∧ q.Prime

/-- `Q`, the size of the pool. -/
def QOf (n : ℕ) : ℕ := (Light.Filter2.pool n).card

/-- The guard (3.6), `K (2h)^r 2^d ≤ Q^r` with `K = s`.  Not tested by the program on the main
path; it holds for all large `n` along polynomially bounded `U`. -/
def Guard (n U : ℕ) : Prop := Light.Filter2.sOf n * (2 * Light.Filter2.hOf n U) ^ Light.Filter2.rOf n * 2 ^ Light.Filter2.dOf n ≤ Light.Filter2.QOf n ^ Light.Filter2.rOf n

instance (n U : ℕ) : Decidable (Light.Filter2.Guard n U) := by unfold Light.Filter2.Guard; infer_instance

/-- The largest modulus, `(2R)^r`. -/
def PmaxOf (n : ℕ) : ℕ := (2 * Light.Filter2.RR n) ^ Light.Filter2.rOf n

/-! ## The counts -/

/-- The block `j` of third vertices: `c` with `j s ≤ c < (j + 1) s` and `c < n`. -/
def InBlock (n s j c : ℕ) : Prop := j * s ≤ c ∧ c < (j + 1) * s ∧ c < n

instance (n s j c : ℕ) : Decidable (Light.Filter2.InBlock n s j c) := by unfold Light.Filter2.InBlock; infer_instance

/-- The triangle sum `w(a,b) + w(b,c) + w(a,c)` of the pair `q = a n + b` and the vertex `c`, as in
`aeFlags`. -/
def triSum (n : ℕ) (AB BC AC : List ℤ) (q c : ℕ) : ℤ :=
  AB.getD q 0 + BC.getD (q % n * n + c) 0 + AC.getD (q / n * n + c) 0

/-- The count of the candidate `i = j n² + q` at the modulus `p`: the `c` in block `j` with
`p ∣ triSum q c`. -/
def cntOf (n s p : ℕ) (AB BC AC : List ℤ) (i : ℕ) : ℕ :=
  ((Finset.range n).filter fun c =>
    Light.Filter2.InBlock n s (i / (n * n)) c ∧ (p : ℤ) ∣ Light.Filter2.triSum n AB BC AC (i % (n * n)) c).card

/-- All `m n²` counts at the modulus `p`, candidate `i` in cell `i`. -/
def cntList (n s p : ℕ) (AB BC AC : List ℤ) : List ℤ :=
  (List.range ((n + s - 1) / s * (n * n))).map fun i => (Light.Filter2.cntOf n s p AB BC AC i : ℤ)

/-! ## The counting oracle as a task -/

/-- An instance of the oracle: an all-edges instance (without the place of the flags), a block size
`s`, a modulus `p`, and the place `cnt` of the `m n²` counts. -/
structure OracleInst : Type extends TriInst where
  s : ℕ
  p : ℕ
  cnt : ℕ

/-- The number of counts, `⌈n/s⌉ n²`. -/
def OracleInst.len (x : Light.Filter2.OracleInst) : ℕ := (x.n + x.s - 1) / x.s * (x.n * x.n)

/-- The matrices and the place of the counts lie below the free pointer and do not meet;
`1 ≤ s ≤ n`, `1 ≤ p`. -/
structure OracleInst.Pre (x : Light.Filter2.OracleInst) (μ : ℕ → ℤ) (fr : ℕ) : Prop
    extends TriInst.Pre x.toTriInst μ fr where
  s_pos : 1 ≤ x.s
  s_le : x.s ≤ x.n
  p_pos : 1 ≤ x.p
  belowCnt : x.cnt + x.len ≤ fr
  apartAB : Apart x.ab (x.n * x.n) x.cnt x.len
  apartBC : Apart x.bc (x.n * x.n) x.cnt x.len
  apartAC : Apart x.ac (x.n * x.n) x.cnt x.len

/-- **The counting oracle** (Lemma 4.1): `oracle(n, U, ab, bc, ac, s, p, cnt, fr)` writes the counts
of all candidates at the modulus `p` to `cnt`, and changes no other cell below `fr`. -/
def oracleTask : TaskN where
  Inst := Light.Filter2.OracleInst
  pars x := [x.n, x.U, x.s, x.p]
  args x := [x.n, x.U, x.ab, x.bc, x.ac, x.s, x.p, x.cnt]
  Pre := Light.Filter2.OracleInst.Pre
  Post x μ fr _ μ' := Seg μ' x.cnt (Light.Filter2.cntList x.n x.s x.p x.AB x.BC x.AC) ∧
    KeptBut μ μ' fr x.cnt x.len

end Light.Filter2
end
end

section
-- Source module: ThreeSumApsp.Batch.OMath
/-!
# The batched counting oracle: the sources and the answer of a query

The two sources of a batch have the same shape (`srcMat`): the entry `(ρ, κ)` is 1 iff the product
`u` and the row `R` of the row `ρ`, and the vertex `cl = κ / p` of the middle block, are valid, and
`κ % p` is the residue `fGen T1 T2 u R cl = (T1[a n + c] + T2[l]) mod p` (`R = a k + l`,
`c = j s + cl`).  The left source has `T1 = w(a, c) mod p`, `T2[l] = -l mod p`; the right source,
transposed, has `T1 = -w(b, c) mod p`, `T2[h] = h k mod p`.

`query_eq_cntOf`: the query of the batched data structure at the virtual positions of the pair
`(a, b)` in the slot `(g, h)` returns the count of the candidate `j n² + a n + b`, `j` the middle
block of the product `u = g K₀ + h`.
-/

@[expose] public section

open Finset ThreeSumApsp ThreeSumApsp.Batch ThreeSumApsp.Spec

namespace Light.Batch

/-- The residue of the cell of the row `R` and the vertex `cl`. -/
def fGen (n s p k KK bt : ℕ) (T1 T2 : ℕ → ℕ) (u R cl : ℕ) : ℕ :=
  (T1 (R / k * n + (ThreeSumApsp.Batch.jOf KK bt u * s + cl)) + T2 (R % k)) % p

/-- A source of a batch, `rows × D`. -/
def srcMat (rows D' n s p k T KK N0 bt : ℕ) (T1 T2 : ℕ → ℕ) : Matrix (Fin rows) (Fin D') ℤ :=
  fun ρ κ => srcVal s p KK N0 (okX n s k T KK bt) (fGen n s p k KK bt T1 T2) ρ κ

/-- The tables of the left source. -/
def t1X (p : ℕ) (AC : List ℤ) (q : ℕ) : ℕ := ThreeSumApsp.Batch.resid p (AC.getD q 0)
def t2X (p : ℕ) (l : ℕ) : ℕ := (p - l) % p
/-- The tables of the right source. -/
def t1Y (p : ℕ) (BC : List ℤ) (q : ℕ) : ℕ := ThreeSumApsp.Batch.resid p (-BC.getD q 0)
def t2Y (p k : ℕ) (h : ℕ) : ℕ := h * k % p
end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.OSpecs
/-!
# The batched counting oracle: what its small procedures do

The specifications of the procedures 76 to 82 (`OProg`), as propositions about a program `P`, with
their times.  They are proved in `OTables` (the tables) and `OWrite` (writeSrc).
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Batch Light.Sec4

/-! ## Times -/

def tPowNat (e : ℕ) : ℕ := 12 * e + 12
def tResidTab (len V p : ℕ) : ℕ := len * (emodTime V p + 30) + 10
def tSmallTab (len V p : ℕ) : ℕ := len * (emodTime V p + 30) + 10
def tRowTab (nb N0 : ℕ) : ℕ := nb * (N0 * 30 + 30) + 10
def tDivTab (k : ℕ) : ℕ := k * (k * 30 + 30) + 10
def tZTab (len : ℕ) : ℕ := len * 60 + 20
/-- writeSrc: the clearing of `sz` cells, then the products, rows and cells. -/
def tWriteSrc (sz KKb n k s : ℕ) : ℕ :=
  tFill sz + KKb * (n * (k * (s * 60 + 60) + 30) + 60) + 30

/-! ## Specifications -/

variable (lim : Limits) (P : Program)

/-- powNat(b, e) returns `b^e`, for `b ≥ 1`, and changes no cell. -/
def PowNatSpec : Prop :=
  ∀ (b e : ℕ) (μ : ℕ → ℤ) (d : ℕ), 1 ≤ b → ((b ^ e + e + 1 : ℕ) : ℤ) ≤ lim.word →
    Meets lim P pPowNat d [b, e] μ (tPowNat e) fun r μ' => r = (b ^ e : ℕ) ∧ μ' = μ

/-- residTab(dst, src, len, sgn, p, fr), `sgn = ±1`: `dst[q] := resid p (sgn · src[q])`.  It changes
the `len` cells from dst and cells from fr on (the table of emod). -/
def ResidTabSpec : Prop :=
  ∀ (dst src : ℕ) (W : List ℤ) (sgn : ℤ) (p V fr : ℕ) (μ : ℕ → ℤ) (d : ℕ),
    (sgn = 1 ∨ sgn = -1) → Seg μ src W → (∀ w ∈ W, |w| ≤ V) → Emod.Pre lim p V fr →
    dst + W.length ≤ fr → src + W.length ≤ fr → Apart dst W.length src W.length →
    d + 1 ≤ lim.depth →
    Meets lim P pResidTab d [dst, src, W.length, sgn, p, fr] μ (tResidTab W.length V p)
      fun _ μ' => SegN μ' dst (W.map fun w => resid p (sgn * w)) ∧
        SameOn (fun b => b < fr ∧ Outside dst W.length b) μ μ'

/-- smallTab(dst, len, mul, add, p, fr): `dst[h] := resid p (h · mul + add)`. -/
def SmallTabSpec : Prop :=
  ∀ (dst len : ℕ) (mul add : ℤ) (p V fr : ℕ) (μ : ℕ → ℤ) (d : ℕ),
    (∀ h < len, |(h : ℤ) * mul + add| ≤ V) → |mul| ≤ V → |add| ≤ V → (len : ℤ) ≤ V →
    Emod.Pre lim p V fr → dst + len ≤ fr → d + 1 ≤ lim.depth →
    Meets lim P pSmallTab d [dst, len, mul, add, p, fr] μ (tSmallTab len V p)
      fun _ μ' => SegN μ' dst ((List.range len).map fun h => resid p ((h : ℤ) * mul + add)) ∧
        SameOn (fun b => b < fr ∧ Outside dst len b) μ μ'

/-- rowTab(dst, nb, N0, Mul): `dst[R] := R / N0 · Mul + R % N0` for `R < nb N0`. -/
def RowTabSpec : Prop :=
  ∀ (dst nb N0 Mul : ℕ) (μ : ℕ → ℤ) (d : ℕ), Std lim → 1 ≤ N0 → dst + nb * N0 < lim.space →
    (((nb + 1) * (Mul + N0 + 1) : ℕ) : ℤ) ≤ lim.word →
    Meets lim P pRowTab d [dst, nb, N0, Mul] μ (tRowTab nb N0)
      fun _ μ' => SegN μ' dst ((List.range (nb * N0)).map fun R => R / N0 * Mul + R % N0) ∧
        SameOutside μ μ' dst (nb * N0)

/-- divTab(dst, k): `dst[z] := z / k` for `z < k²`. -/
def DivTabSpec : Prop :=
  ∀ (dst k : ℕ) (μ : ℕ → ℤ) (d : ℕ), Std lim → dst + k * k < lim.space →
    Meets lim P pDivTab d [dst, k] μ (tDivTab k)
      fun _ μ' => SegN μ' dst ((List.range (k * k)).map fun z => z / k) ∧
        SameOutside μ μ' dst (k * k)

/-- zTab(zh, zl, rn, dk, len, k): `zh[q] := rn[q] / k`, `zl[q] := rn[q] % k`, from the table
`dk` of divTab. -/
def ZTabSpec : Prop :=
  ∀ (zh zl rn dk k : ℕ) (RN : List ℕ) (μ : ℕ → ℤ) (d : ℕ), Std lim → 1 ≤ k → SegN μ rn RN →
    SegN μ dk ((List.range (k * k)).map fun z => z / k) → (∀ z ∈ RN, z < k * k) →
    Apart zh RN.length zl RN.length → Apart zh RN.length rn RN.length →
    Apart zl RN.length rn RN.length → Apart zh RN.length dk (k * k) →
    zh + RN.length < lim.space → zl + RN.length < lim.space → rn + RN.length < lim.space →
    dk + k * k < lim.space →
    Meets lim P pZTab d [zh, zl, rn, dk, RN.length, k] μ (tZTab RN.length)
      fun _ μ' => SegN μ' zh (RN.map (· / k)) ∧ SegN μ' zl (RN.map (· % k)) ∧
        SameOn (fun b => Outside zh RN.length b ∧ Outside zl RN.length b) μ μ'

/-- What writeSrc assumes.  The source has `KK NP` rows of `D` cells from base; the tables T1
(`n²` residues) and T2 (`k` residues) and the table RS of the first rows `R / N0 · KK N0 + R % N0`
lie apart from it; the products `j0 + u`, `u < KKb`, have middle blocks that begin below `n`. -/
structure WriteSrcPre (μ : ℕ → ℤ) (base t1 t2 rs n k p D N0 j0 KKb s KK NP : ℕ) (T1 T2 : List ℕ) :
    Prop where
  std : Std lim
  p_pos : 1 ≤ p
  N0_pos : 1 ≤ N0
  k_pos : 1 ≤ k
  D_pos : 1 ≤ D
  seg1 : SegN μ t1 T1
  len1 : T1.length = n * n
  lt1 : ∀ x ∈ T1, x < p
  seg2 : SegN μ t2 T2
  len2 : T2.length = k
  lt2 : ∀ x ∈ T2, x < p
  segRS : SegN μ rs ((List.range NP).map fun R => R / N0 * (KK * N0) + R % N0)
  NP_mul : N0 ∣ NP
  nk_le : n * k ≤ NP
  KKb_le : KKb ≤ KK
  blocks : ∀ u < KKb, (j0 + u) * s < n
  sp_le : s * p ≤ D
  space : base + KK * NP * D < lim.space
  apart1 : Apart base (KK * NP * D) t1 (n * n)
  apart2 : Apart base (KK * NP * D) t2 k
  apartRS : Apart base (KK * NP * D) rs NP
  below1 : t1 + n * n < lim.space
  below2 : t2 + k < lim.space
  belowRS : rs + NP < lim.space
  word : (((j0 + KKb + 1) * (s + 1) + n * n + n * k + 2 * p + 4 : ℕ) : ℤ) ≤ lim.word

/-- writeSrc(base, t1, t2, rs, n, k, p, D, N0, j0, KKb, s, KK NP D): afterwards the `KK NP × D`
source at base holds `srcMat` for the batch `bt` (`j0 = bt KK`, the valid products `u < KKb`), and
no other cell has changed. -/
def WriteSrcSpec : Prop :=
  ∀ (base t1 t2 rs n k p D N0 bt KKb s KK NP : ℕ) (T1 T2 : List ℕ) (μ : ℕ → ℤ) (d : ℕ),
    WriteSrcPre lim μ base t1 t2 rs n k p D N0 (bt * KK) KKb s KK NP T1 T2 →
    d + 1 ≤ lim.depth →
    Meets lim P pWriteSrc d [base, t1, t2, rs, n, k, p, D, N0, bt * KK, KKb, s, KK * NP * D] μ
      (tWriteSrc (KK * NP * D) KKb n k s) fun _ μ' =>
      (∀ ρ < KK * NP, ∀ κ < D, μ' (base + ρ * D + κ) =
        srcVal s p KK N0 (okX n s k (bt * KK + KKb) KK bt)
          (fGen n s p k KK bt (fun q => T1.getD q 0) (fun l => T2.getD l 0)) ρ κ) ∧
      SameOutside μ μ' base (KK * NP * D)

end Light.Batch
end
end

section
-- Source module: ThreeSumApsp.Batch.OQuery
/-!
# The batched counting oracle: the queries of a batch

`queries(b0, ZL, ZH, RV, n, k, N0, K0, KKb, cnt, j0)` asks, for each product `u < KKb` of the batch
(slot `g = u / K₀`, `h = u % K₀`) and each pair `(a, b)`, the query at the virtual positions
`I = virt g (a k + ZL[a n + b])`, `J = virt h (b k + ZH[a n + b])`, and stores the answer at
`cnt + (j0 + u) n² + a n + b` (`queries_ok`).  The invariant `QInv E`: the data structure is ready,
the first `E` cells from `cnt + j0 n²` hold their answers, and no other cell below b0 has changed.
-/

@[expose] public section

namespace Light.Batch

open ThreeSumApsp ThreeSumApsp.Spec ThreeSumApsp.Batch Light.Sec4 QS

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

/-- The answer of the query of the pair `(a, b)` in the slot `(g, h)`. -/
def qVal (pv : Sec2.Par) (t : ℕ) (hmL : pv.m ≤ pv.L) (AL : ArrL pv) (AR : Light.Batch.ArrR pv) (n k : ℕ)
    (ZL ZH : List ℕ) (g h a b : ℕ) : ℤ :=
  queryValue pv.m t
    (encRowA AL (virt pv.K0 pv.N0 g (a * k + ZL.getD (a * n + b) 0) / (pv.K0 * pv.N0)))
    (encColA AR (virt pv.K0 pv.N0 h (b * k + ZH.getD (a * n + b) 0) / (pv.K0 * pv.N0)))
    (outStrOfPos (stdLayout hmL) (virt pv.K0 pv.N0 g (a * k + ZL.getD (a * n + b) 0))
      (virt pv.K0 pv.N0 h (b * k + ZH.getD (a * n + b) 0)))

/-- The answer of the cell `e` from `cnt + j0 n²`. -/
def qAns (pv : Sec2.Par) (t : ℕ) (hmL : pv.m ≤ pv.L) (AL : ArrL pv) (AR : Light.Batch.ArrR pv) (n k : ℕ)
    (ZL ZH : List ℕ) (e : ℕ) : ℤ :=
  qVal pv t hmL AL AR n k ZL ZH (e / (n * n) / pv.K0) (e / (n * n) % pv.K0) (e % (n * n) / n)
    (e % (n * n) % n)

/-- What the queries of a batch assume, in the memory μ₀ at the start. -/
structure QSet (lim : Limits) (P : Program) (pv : Sec2.Par) (t : ℕ) (aX aY b0 : ℕ) (AL : ArrL pv)
    (AR : Light.Batch.ArrR pv) (U : ℤ) (μ₀ : ℕ → ℤ) (zl zh rv n k NP KKb cnt j0 : ℕ) (ZL ZH : List ℕ) :
    Prop where
  hq : QueryAtASpec lim P
  ht : t ≤ pv.m
  hlim : Lim30 lim pv t b0 U
  hb : BoundA AL AR U
  segZL : SegN μ₀ zl ZL
  segZH : SegN μ₀ zh ZH
  lenZL : ZL.length = n * n
  lenZH : ZH.length = n * n
  ltZL : ∀ x ∈ ZL, x < k
  ltZH : ∀ x ∈ ZH, x < k
  segRV : SegN μ₀ rv ((List.range NP).map fun R => R / pv.N0 * (pv.K0 * pv.N0) + R % pv.N0)
  hN : pv.N = pv.K0 * NP
  n_pos : 1 ≤ n
  NP_mul : pv.N0 ∣ NP
  nk_le : n * k ≤ NP
  KKb_le : KKb ≤ pv.K0 * pv.K0
  belowZL : zl + n * n ≤ b0
  belowZH : zh + n * n ≤ b0
  belowRV : rv + NP ≤ b0
  belowCnt : cnt + (j0 + KKb) * (n * n) ≤ b0
  apZL : Apart (cnt + j0 * (n * n)) (KKb * (n * n)) zl (n * n)
  apZH : Apart (cnt + j0 * (n * n)) (KKb * (n * n)) zh (n * n)
  apRV : Apart (cnt + j0 * (n * n)) (KKb * (n * n)) rv NP

/-- The invariant: the data structure is ready, the first `E` cells from `cnt + j0 n²` hold their
answers, and no other cell below b0 has changed. -/
def QInv (pv : Sec2.Par) (t : ℕ) (hmL : pv.m ≤ pv.L) (aX aY b0 : ℕ) (AL : ArrL pv) (AR : Light.Batch.ArrR pv)
    (μ₀ : ℕ → ℤ) (n k : ℕ) (ZL ZH : List ℕ) (base E : ℕ) (μ : ℕ → ℤ) : Prop :=
  DSReadyA pv t aX aY b0 AL AR μ ∧ (∀ e < E, μ (base + e) = qAns pv t hmL AL AR n k ZL ZH e) ∧
    ∀ x < b0, (x < base ∨ base + E ≤ x) → μ x = μ₀ x

/-- The time of the queries of a batch. -/
def tQueries (KKb n tq : ℕ) : ℕ := KKb * (n * (n * (tq + 80) + 20) + 80) + 30

section cell

variable {pv : Sec2.Par} {t : ℕ} {aX aY b0 : ℕ} {AL : ArrL pv} {AR : ArrR pv} {U : ℤ}
  {μ₀ : ℕ → ℤ} {zl zh rv n k NP KKb cnt j0 : ℕ} {ZL ZH : List ℕ}

/-- The locals of queries. -/
abbrev qFr (b0 zl zh rv n k N0 K0 KKb cnt j0 g h a b : ℕ) (I J r : ℤ) (dst u : ℕ) : List ℤ :=
  [b0, zl, zh, rv, n, k, N0, K0, KKb, cnt, j0, g, h, a, b, I, J, r, dst, u]

/-- The time of the queries of a row `a`. -/
def tQRow (n tq : ℕ) : ℕ := n * (tq + 80) + 6

/-- The time of the queries of a product. -/
def tQRows (n tq : ℕ) : ℕ := n * (tQRow n tq + 10) + 6

/-- The time of a product. -/
def tQProd (n tq : ℕ) : ℕ := tQRows n tq + 60

end cell

end Light.Batch
end
end