APSP in $O(n^{2.99791})$: definitions, part 3 of 5

/-
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 27a2fa4bab32.

import Definitions.Def_APSP299791_Defs2

set_option autoImplicit false
set_option relaxedAutoImplicit false

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Query
/-!
# A query, from the block that holds the data structure (proof of Theorem 30, "Query")

"Given (I, J), we find its tile from the bands of row I and column J, the subset Q of its block
product, and its output string w […], all in O(L) operations". Below, the output string is called η.
queryAt(I, J, b0) has four parts. It reads the sizes, the addresses of the tables and the bands of I
and J from the block at b0 (queryAtReads), computes the addresses of the areas of Section 4
(queryAtAreas), lets outDigits write the digits of η (queryAtOut), reads the root of the trie of the
tile from the table of roots, and lets queryCore compute the sum of Lemma 28 from the two encodings
and this trie (queryAtCore). queryCore is handed the encodings as the memory holds them; at the
leaves that it reads they are the true ones (`coreArgs_agree`). There is one lemma for each part;
for each of the two calls a record holds the arguments (outArgs, coreArgsP) and a lemma says that the
block holds what the routine assumes (outArgs_pre, coreArgs_pre). queryAtP_spec puts the four parts
together.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Sec4

open ThreeSumApsp.Spec

/-! ## The program -/

namespace QueryAt

end QueryAt

section Proof

variable {lim : Limits} {P : Program} {d : ℕ} {p : Sec2.Par} {t : ℕ} {hmL : p.m ≤ p.L}
  {aX aY b0 : ℕ} {X : Matrix (Fin p.N) (Fin (D p.m)) ℤ} {Y : Matrix (Fin (D p.m)) (Fin p.N) ℤ}
  {U : ℤ} {μ : ℕ → ℤ}

/-! ## The first two parts -/

/-! ## The call of outDigits -/

/-! ## The call of queryCore -/

/-- The arguments of the call of queryCore, and the data behind them: the encodings are what the
memory μ holds. -/
def coreArgsP (p : Sec2.Par) (t : ℕ) (hmL : p.m ≤ p.L) (b0 : ℕ)
    (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ) (Y : Matrix (Fin (ThreeSumApsp.D p.m)) (Fin p.N) ℤ) (μ : ℕ → ℤ)
    (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 := (dsTries p t hmL X Y).root (tileOf p I J)
  wd := aWD p b0
  ss := aSS p b0
  box := aBOX p b0
  T := (dsTries p t hmL X Y).cells

section Tile

variable (I J : Fin p.N)

variable (ht : t ≤ p.m)

include ht

end Tile

/-! ## The routine -/

end Proof

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Routines
/-!
# The routines of Section 4, assembled

The fourteen routines with the numbers 40 to 53, as a list (procs40), and the theorem that every
program that holds them at these numbers (Has40) meets all their specifications (specs40). Each
routine is proved under the specifications of the routines that it calls; here these assumptions are
discharged, in the order in which the routines call each other. For the programs of Theorem 30, its
offline form and Corollary 26, which hold the fourteen routines behind forty others
(has40_of_append), this gives the specifications of the two routines that know the memory map:
preCoreSpec_of and preCoreSpec_all for the preprocessing, queryAtSpec_all for a query.
-/

@[expose] public section

namespace Light.Sec4

/-- The specifications of all the routines of Section 4. -/
structure Specs40P (lim : Limits) (P : Program) : Prop where
  nineFirst : NineFirstSpec lim P
  nineNext : NineNextSpec lim P
  scatter : ScatterSpec lim P
  starFirst : StarFirstSpec lim P
  horner : HornerSpec lim P
  lookup : LookupSpec lim P
  insert : InsertSpec lim P
  newRoot : NewRootSpec lim P
  sumTen : SumTenSpec lim P
  fillList : FillListSpecP lim P
  tile : TileSpecP lim P
  allTiles : AllTilesSpecP lim P
  outDigits : OutDigitsSpec lim P
  queryCore : QueryCoreSpecP lim P

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec2.PrunedEnc.AllBands
/-!
# The pruned encodings of all bands of one matrix

encodeBandsP is encodeBands with the pruned encoder: for each band β it calls bandArray, and then
encodeP with the budget m of nines, which leaves the encoding of the band at the leaves with at most
m nines in the T = 10^L cells from enc + β T.  The arguments are those of encodeBands.

* `encodeBandsPBody`, `encodeBandsP_proc`: the loop, for callees described by what they do here.
* `EncodeBandsPLSpec`, `EncodeBandsPRSpec`: the two entries, met by `encodeBandsPL_entry` and
  `encodeBandsPR_entry`.
-/

@[expose] public section

namespace Light.Sec2

open ThreeSumApsp

variable {lim : Limits} {P : Program}

/-! ## Pruned segments -/

open Bands in
/-- encodeBandsP(L, m, N, D, K0, N0, S7, T, nB, side, aA, sI, sK, mask, dig3, dig4, arr, zs, tab,
p7, p10, enc): for each band, the input array is formed at arr and its encoding is written into the
T cells from enc + β T, at the leaves with at most m nines. -/
def encodeBandsPBody : Stmt :=
  .for Light.Sec2.Bands.BandNo (Light.v NumBands) (
    .call pBandArray [Light.v Light.Sec2.Bands.Levels, Light.v Light.Sec2.Bands.Size, Light.v Light.Sec2.Bands.Rows, Light.v Light.Sec2.Bands.Cols, Light.v Light.Sec2.Bands.NumBlocks, Light.v Light.Sec2.Bands.Width, Light.v Light.Sec2.Bands.Cells, Light.v Light.Sec2.Bands.BandNo,
      Light.v Light.Sec2.Bands.Side, Light.v Light.Sec2.Bands.Source, Light.v Light.Sec2.Bands.StepRow, Light.v Light.Sec2.Bands.StepCol, Light.v Light.Sec2.Bands.MaskTable, Light.v Light.Sec2.Bands.Digits3, Light.v Light.Sec2.Bands.Digits4, Light.v Arr] Light.Sec2.Bands.Res ;;
    .call pEncodeP [Light.v Light.Sec2.Bands.Levels, Light.v Arr, Light.v Enc +' Light.v Light.Sec2.Bands.BandNo *' Light.v Light.Sec2.Bands.Leaves, Light.v Light.Sec2.Bands.Scratch, Light.v Light.Sec2.Bands.Table, Light.v Light.Sec2.Bands.Pow7,
      Light.v Light.Sec2.Bands.Pow10, Light.v Light.Sec2.Bands.Size] Light.Sec2.Bands.Res)

variable {p : Par} {x : BandsArgs} {coefs : List ℤ} {V : ℤ} {μ μ' : ℕ → ℤ}

/-! ## The two entries -/

/-- encodeBandsP for the row bands of X: side = 0, the coefficients φ at tab. -/
def EncodeBandsPLSpec (lim : Limits) (P : Program) (c : ℕ) : Prop :=
  ∀ (p : Par) (hmL : p.m ≤ p.L) (x : BandsArgs) (aX : ℕ) (μ : ℕ → ℤ)
    (X : Matrix (Fin p.N) (Fin (ThreeSumApsp.D p.m)) ℤ) (V : ℤ),
    BandsPre lim p μ x Spec.phiFlat V → MatAt μ aX X → aX + p.N * p.D ≤ x.enc →
    (∀ i j, |X i j| ≤ V) →
    ∀ 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 (bandArrayL (Spec.stdLayout hmL) X β)))) ∧
          SameOutside μ μ' x.enc (x.zs + p.S7 - x.enc)

/-- encodeBandsP for the column bands of Y: side = 1, the coefficients ψ at tab. -/
def EncodeBandsPRSpec (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.N) ℤ) (V : ℤ),
    BandsPre lim p μ x Spec.psiFlat V → MatAt μ aY Y → aY + p.D * p.N ≤ x.enc →
    (∀ i j, |Y i j| ≤ V) →
    ∀ 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 (bandArrayR (Spec.stdLayout hmL) Y β)))) ∧
          SameOutside μ μ' x.enc (x.zs + p.S7 - x.enc)

end Light.Sec2
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec2.PrunedEnc.SharedStage
/-!
# The shared stage with pruned encodings

sharedP(L, m, N, D, aX, aY, b0) is shared, except that both calls of encodeBands go to encodeBandsP:
the encodings of the bands are computed only at the leaves with at most m nines.  The first two
parts (sharedA, sharedB) and the directory are those of shared; their lemmas are restated here for
the callees of sharedP (`SharedCalleesP`), and the third part (sharedCP) is proved as sharedC is.
sharedP_entry: if every callee meets its entry with the constant c, then sharedP meets `SharedSpecP`
with the constant 12 c + 600.
-/

@[expose] public section

namespace Light.Sec2

open ThreeSumApsp

variable {lim : Limits} {P : Program}

open Shared

/-- The entries of the routines that the shared stage with pruned encodings calls. -/
structure SharedCalleesP (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 : EncodeBandsPLSpec lim P c
  bandsR : EncodeBandsPRSpec lim P c

/-- The third part of sharedP: as sharedC, with the pruned encodings of all bands. -/
def sharedCP : Stmt :=
  sharedBands ;;
  .set AdrEncB (Light.v AdrEncA +' Light.v Light.Sec2.Shared.Bands *' Light.v Light.Sec2.Shared.Leaves) ;;
  .set AdrArr (Light.v AdrEncB +' Light.v Light.Sec2.Shared.Bands *' Light.v Light.Sec2.Shared.Leaves) ;;
  .set AdrZs (Light.v AdrArr +' Light.v Strings) ;;
  .set AdrEnd (Light.v AdrZs +' Light.v Strings) ;;
  .call pEncodeBandsP [Light.v Light.Sec2.Shared.Levels, Light.v Light.Sec2.Shared.Inner, Light.v Light.Sec2.Shared.Rows, Light.v Light.Sec2.Shared.Cols, Light.v Light.Sec2.Shared.Root, Light.v BlockRows, Light.v Strings, Light.v Light.Sec2.Shared.Leaves,
    Light.v Light.Sec2.Shared.Bands, Light.k 0, Light.v Light.Sec2.Shared.AdrX, Light.v Light.Sec2.Shared.Cols, Light.k 1, Light.v AdrMask, Light.v AdrDig3, Light.v AdrDig4, Light.v AdrArr, Light.v AdrZs, Light.v AdrPhi,
    Light.v AdrP7, Light.v AdrP10, Light.v AdrEncA] Light.Sec2.Shared.Void ;;
  .call pEncodeBandsP [Light.v Light.Sec2.Shared.Levels, Light.v Light.Sec2.Shared.Inner, Light.v Light.Sec2.Shared.Rows, Light.v Light.Sec2.Shared.Cols, Light.v Light.Sec2.Shared.Root, Light.v BlockRows, Light.v Strings, Light.v Light.Sec2.Shared.Leaves,
    Light.v Light.Sec2.Shared.Bands, Light.k 1, Light.v Light.Sec2.Shared.AdrY, Light.k 1, Light.v Light.Sec2.Shared.Rows, Light.v AdrMask, Light.v AdrDig3, Light.v AdrDig4, Light.v AdrArr, Light.v AdrZs, Light.v AdrPsi,
    Light.v AdrP7, Light.v AdrP10, Light.v AdrEncB] Light.Sec2.Shared.Void ;;
  storeLocals 0 [Light.Sec2.Shared.Levels, Light.Sec2.Shared.Inner, Light.Sec2.Shared.Rows, Light.Sec2.Shared.Cols, Light.Sec2.Shared.Outer, BlockRows, Subsets, Light.Sec2.Shared.Root, Light.Sec2.Shared.Table, Light.Sec2.Shared.Bands, Light.Sec2.Shared.Leaves,
    Strings, Light.Sec2.Shared.AdrX, Light.Sec2.Shared.AdrY, AdrP3, AdrP4, AdrP7, AdrP10, AdrPas, AdrPhi, AdrPsi, AdrMask, AdrBand,
    AdrBlock, AdrDig3, AdrDig4, AdrEncA, AdrEncB, AdrArr, AdrZs, AdrEnd]

/-- sharedP(L, m, N, D, aX, aY, b0). -/
def sharedPBody : Stmt := sharedA ;; sharedB ;; sharedCP

/-- The time of the third part of sharedP. -/
def sharedTimeCP (c : ℕ) (p : Par) : ℕ :=
  2 * (c * (p.nB * (retW p.L p.m + bandArrayShape p) + 1)) + 236

end Light.Sec2
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec2.PrunedEnc.Recursion
/-!
# The pruned encoding, by the recursion of Section 2.4.1

The pruned encoder is `encode` with one more argument, the budget r of digits 9.  At level `L + 1`
the terms 0, …, 8 are treated as in `encode`, with the same budget for the recursive call; the term
number 9 (the term P₀ of Schönhage's identity) is treated only if `0 < r`, with the budget `r - 1`.
So the leaves whose code has more than r digits 9 are not written.  Before the term 9 the budget is
cut down to the level: a budget of more than `L` nines prunes nothing at level `L`, and so the
budget `r - 1` that is passed on is at most `L` and fits in a word.

* `encPTime L r` is the number of steps, and `encPTime_le` bounds it by `O(retW L r)`.
* `encodePTerm_spec` treats the two calls for one term, `encodeP_leaf` the level 0, `encodeP_node`
  the levels `L + 1`, and `encodeP_spec` is the induction.
-/

@[expose] public section

namespace Light.Sec2

open Finset ThreeSumApsp

namespace EncodeP

/-- The local variables of encodeP: the arguments L, src, out, scr, tab, p7, p10, r; the lengths
7^{L-1} and 10^{L-1}; the number of the term; the result of a call, which is not used. -/
abbrev Level : ℕ := 0
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Src : ℕ := 1
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Out : ℕ := 2
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Scr : ℕ := 3
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Table : ℕ := 4
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Pow7 : ℕ := 5
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Pow10 : ℕ := 6
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Budget : ℕ := 7
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Len7 : ℕ := 8
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Len10 : ℕ := 9
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev TermNo : ℕ := 10
@[inherit_doc Light.Sec2.EncodeP.Level] abbrev Res : ℕ := 11

end EncodeP

open EncodeP

/-- The two calls for the term number TermNo, the recursive one with the budget b. -/
def encodePTerm (pS pE : ℕ) (b : Expr) : Stmt :=
  .call pS [Light.v Light.Sec2.EncodeP.Src, Light.v Light.Sec2.EncodeP.Scr, Light.v Len7, Light.v TermNo, Light.v Light.Sec2.EncodeP.Table] Light.Sec2.EncodeP.Res ;;
  .call pE [Light.v Light.Sec2.EncodeP.Level -' Light.k 1, Light.v Light.Sec2.EncodeP.Scr, Light.v Light.Sec2.EncodeP.Out +' Light.v TermNo *' Light.v Len10, Light.v Light.Sec2.EncodeP.Scr +' Light.v Len7, Light.v Light.Sec2.EncodeP.Table, Light.v Light.Sec2.EncodeP.Pow7,
    Light.v Light.Sec2.EncodeP.Pow10, b] Light.Sec2.EncodeP.Res

/-- encodeP(L, src, out, scr, tab, p7, p10, r).  The loop over the terms 0, …, 8 leaves 9 in
TermNo; the term 9 is treated only if `0 < r`, with the budget `min(r, L) - 1`.  pS and pE are the
numbers that the procedures encStep and encodeP have in the program. -/
def encodePBody (pS pE : ℕ) : Stmt :=
  .ite (Light.v Light.Sec2.EncodeP.Level =' Light.k 0)
    (.store (Light.v Light.Sec2.EncodeP.Out) (M (Light.v Light.Sec2.EncodeP.Src)))
    (.set Len7 (M (Light.v Light.Sec2.EncodeP.Pow7 +' Light.v Light.Sec2.EncodeP.Level -' Light.k 1)) ;;
     .set Len10 (M (Light.v Light.Sec2.EncodeP.Pow10 +' Light.v Light.Sec2.EncodeP.Level -' Light.k 1)) ;;
     .for TermNo (Light.k 9) (encodePTerm pS pE (Light.v Budget)) ;;
     .ite (Light.k 0 <' Light.v Budget)
       (.ite (Light.v Light.Sec2.EncodeP.Level <' Light.v Budget) (.set Budget (Light.v Light.Sec2.EncodeP.Level)) .skip ;;
        encodePTerm pS pE (Light.v Budget -' Light.k 1))
       .skip)

/-! ## The number of steps -/

/-- The number of steps of `encodeP`: at level `L + 1`, nine times the two calls for a term with the
same budget, and if `0 < r` once more the two calls with the budget `min (r - 1) L`. -/
def encPTime : ℕ → ℕ → ℕ
  | 0, _ => 8
  | L + 1, 0 => 100 + 9 * (27 + tEncStep (7 ^ L) + encPTime L 0)
  | L + 1, r + 1 => 106 + 9 * (27 + tEncStep (7 ^ L) + encPTime L (r + 1)) +
      (27 + tEncStep (7 ^ L) + encPTime L (min r L))

/-- The time of the two calls for one term, with the budget `r` for the recursive call. -/
def tTermP (L r : ℕ) : ℕ := 27 + tEncStep (7 ^ L) + encPTime L r

/-- The constant of the running time of `encodeP`. -/
def cEncodeP : ℕ := 1035

/-! ## Facts about the positions -/

variable {lim : Limits} {P : Program} {d : ℕ} {μ μ' μ₁ μ₂ : ℕ → ℤ}
  {src out scr tab p7 p10 L lam r : ℕ} {V : ℤ} {a c : ℕ → ℤ}

namespace EncodePre

end EncodePre

/-! ## What encodeP guarantees -/

/-- What `encodeP` guarantees: the output holds the encoding at the positions with at most `r`
nines, and only the `10^L` cells of `out` and the `scrSize L` cells of the scratch area have
changed. -/
def EncodePPost (μ : ℕ → ℤ) (out scr L r : ℕ) (a c : ℕ → ℤ) (μ' : ℕ → ℤ) : Prop :=
  (∀ i < 10 ^ L, ninesIdx L i ≤ r → μ' (out + i) = encIdx c L a i) ∧
    SameOutside2 μ μ' out (10 ^ L) scr (scrSize L)

/-- The memory before the round of the term number `lam`: the parts of the output below the earlier
terms are filled at the positions with at most `r` nines, and only the output and the scratch area
have changed. -/
def TermsDoneP (μ : ℕ → ℤ) (out scr L r : ℕ) (a c : ℕ → ℤ) (lam : ℕ) (μ' : ℕ → ℤ) : Prop :=
  (∀ i < lam * 10 ^ L, ninesIdx (L + 1) i ≤ r → μ' (out + i) = encIdx c (L + 1) a i) ∧
    SameOutside2 μ μ' out (10 ^ (L + 1)) scr (scrSize (L + 1))

/-- The list of the local variables at level L + 1, with the value `rb` of the budget. -/
abbrev EncodeP.locals (L src out scr tab p7 p10 rb lam : ℕ) (res : ℤ) : List ℤ :=
  [(L + 1 : ℕ), src, out, scr, tab, p7, p10, rb, (7 ^ L : ℕ), (10 ^ L : ℕ), lam, res]

/-- The specification of encodeP at level L, for a caller at depth d, with every budget. -/
def EncodePMeets (lim : Limits) (P : Program) (pE d tab p7 p10 L : ℕ) (c : ℕ → ℤ) : Prop :=
  ∀ (μ : ℕ → ℤ) (src out scr : ℕ) (V : ℤ) (a : ℕ → ℤ) (r : ℕ),
    EncodePre lim μ src out scr tab p7 p10 V L a c →
    Meets lim P pE (d + 1) [L, src, out, scr, tab, p7, p10, r] μ (encPTime L r) fun _ μ' =>
      EncodePPost μ out scr L r a c μ'

variable {pS pE : ℕ}

/-! ## The recursion -/

end Light.Sec2
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.Theorem30.Program
/-!
# Theorem 30 on the word RAM: the program

`program30` is the list of procedures: those of Section 2 (numbers 0 to 28, of which the
preprocessing uses the shared stage), the routines of Section 4 (40 to 53), preCore, queryAt and the
two procedures of the offline form (54 to 57), and the two main procedures of Theorem 30 (80 and
81); the unused numbers hold the empty statement.

The first 58 procedures (`base58P`) are also the beginning of the programs for Corollaries 26, 31
and 32.  So the assumptions of the earlier files are discharged for every program `base58P ++ R` that
begins with them: it meets the specifications of the preprocessing and of a query at a given place
of the memory (`preCore_base58`, `queryAt_base58`).  With `theorem_30_of` and `theorem_30_wanted_of`
this gives Theorem 30 and its offline form without hypotheses (`wordRam_theorem_30`,
`wordRam_theorem_30_wanted`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.WordRam

/-- The procedures of the shared stage with pruned encodings, numbers 29 to 31. -/
def prunedProcs : Program :=
  [Sec2.encodePBody Sec2.pEncStep Sec2.pEncodeP, Sec2.encodeBandsPBody, Sec2.sharedPBody]

/-- The procedures number 0 to 57: Section 2's program for the thin matrix product (0 to 28), the
shared stage with pruned encodings (29 to 31), eight unused numbers, the routines of Section 4, the two routines that know the memory map, and the two
routines of the offline form. -/
def base58P : Program :=
  Sec2.programThin ++ prunedProcs ++ List.replicate 8 .skip ++ procs40 ++
    [preCoreBodyP, queryAtBody, wantedCoreBody, wantedMainBody]

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec4.ChoosingParameters.Program
/-!
# Corollaries 26, 31 and 32: the concrete program for given parameters

program31P G = base58P (Theorem 5's and Theorem 30's procedures, numbers 0 to 57) followed by the
procedures 58 to 73 of Section 4.4 (`procs31`): copying, filling, m = ⌈log₄ D⌉, padding and the
inner product; the preprocessing pre31Body G at 64 and the query; their main procedures; the offline
routine and its main procedure; the solver of Corollary 26 for all instances with its test; and the
two procedures for L = ⌈am/b⌉ and t = ⌈pm/q⌉ at 72 and 73. The specifications of the preprocessing,
the query and the offline routine hold for it, for all limits, without hypotheses. At the parameters
of Corollary 26 it is `program26`.

With the two end theorems for light programs and the compiler this gives the two statements about
the word RAM from which Corollaries 26, 31 and 32 follow: on every domain of inputs, and for all
bounds that dominate the two costs of Theorem 30 at the parameters G (`CostsWithin`), the compiled
program is a data structure (`isDataStructure_of_costsWithin`) and solves the offline problem
(`solves_of_costsWithin`).
-/

@[expose] public section

namespace Light.Sec4

open ThreeSumApsp ThreeSumApsp.WordRam

/-! ## Two procedures for Corollary 26

The program carries them for all parameters G. What they do at the parameters of Corollary 26 is
proved in `regimeTest26_meets` and `allInstances26_solves`. -/

/-! ## The program -/

/-- The program for the parameters G. -/
def program31P (G : RatParams) : Program := base58P ++ procs31 G

/-! ## The program on the word RAM -/

end Light.Sec4
end
end

section
-- Source module: ThreeSumApsp.Sec4.Table2.LogBounds
/-!
# The numerical toolkit for Table 2

Table 2 and the proof of Corollary 26 evaluate the exponents `q(θ)` and `γ = θ ln(1/ρ_c)/ln 4` of
Corollary 31 and the bound `R_c(γ)` of (11) at rational points. All three are built from logarithms
of rational numbers. This file turns each such claim into inequalities between rational numbers,
which the tactic `numerics` checks by evaluating both sides.

* `log_mem` encloses `log θ`: write `θ = 2^k (1 + x)/(1 - x)` with `2^k` near `θ`, and sum `n` terms
  of `log ((1 + x)/(1 - x)) = 2 (x + x³/3 + x⁵/5 + ⋯)`.
* `qOf_le` and `le_qOf` bound `q(θ)` at a rational `θ`.
* `gammaOf_one_mem`, `lnΛ_zero_mem`, `lt_Rc_of_lnΛ_zero_mem` and `Rc_lt_of_lnΛ_zero_mem` treat the
  two numbers on which a row of the table depends: `ln(1/ρ_c)/ln 4` and the denominator of (11) at
  `γ = 0`.
* `Table2Query.intro`, `Table2Ninth.intro` and `Table2Density.intro` give an entry of the table from
  two rational numbers that enclose its `θ`; the `θ` itself comes from the intermediate value
  theorem.
-/

@[expose] public section

namespace ThreeSumApsp

/-! ## Logarithms of rational numbers -/

/-- A lower bound on `ln 10 = 2.302585092994…`. -/
noncomputable def logTenLo : ℝ := 2.3025850922

/-- An upper bound on `ln 10 = 2.302585092994…`. -/
noncomputable def logTenHi : ℝ := 2.3025850938

lemma log_two_mem : Real.log 2 ∈ Set.Icc logTwoLo logTwoHi :=
  ⟨Real.log_two_gt_d9.le, Real.log_two_lt_d9.le⟩

/-- The enclosure of `log θ`, between two rational numbers if `θ` is rational. It holds for every
`k` and `n`; they decide only how tight it is. -/
theorem log_mem_b299791 (k : ℤ) (n : ℕ) {θ : ℝ} (hθ : 0 < θ := by (norm_num [ThreeSumApsp.logApprox, ThreeSumApsp.logErr, ThreeSumApsp.seriesArg,
                                                                     ThreeSumApsp.logTwoLo, ThreeSumApsp.logTwoHi, ThreeSumApsp.logThreeLo, ThreeSumApsp.logThreeHi,
                                                                     ThreeSumApsp.logTenLo, ThreeSumApsp.logTenHi, Finset.sum_range_succ, ThreeSumApsp.rhoC])) :
    Real.log θ ∈ Set.Icc (logApprox k n θ - logErr k n θ) (logApprox k n θ + logErr k n θ) := by
  have hpow : (0 : ℝ) < 2 ^ k := by positivity
  have hx : |seriesArg k θ| < 1 := by
    rw [seriesArg, abs_div, abs_of_pos (add_pos hθ hpow), div_lt_one (add_pos hθ hpow), abs_lt]
    constructor <;> linarith
  have hsq : 0 < 1 - seriesArg k θ ^ 2 := sub_pos.2 ((sq_lt_one_iff_abs_lt_one _).2 hx)
  -- `log θ = k log 2 + log ((1 + x)/(1 - x))`
  have hlog : Real.log θ
      = k * Real.log 2 + Real.log ((1 + seriesArg k θ) / (1 - seriesArg k θ)) := by
    have hdiv : (1 + seriesArg k θ) / (1 - seriesArg k θ) = θ / 2 ^ k := by
      rw [seriesArg]
      field_simp
      ring
    rw [hdiv, Real.log_div hθ.ne' hpow.ne', Real.log_zpow]
    ring
  have htwo : |k * Real.log 2 - k * ((logTwoLo + logTwoHi) / 2)|
      ≤ |(k : ℝ)| * ((logTwoHi - logTwoLo) / 2) := by
    rw [← mul_sub, abs_mul]
    refine mul_le_mul_of_nonneg_left (abs_le.2 ⟨?_, ?_⟩) (abs_nonneg _) <;>
      linarith [log_two_mem.1, log_two_mem.2]
  -- Mathlib: the sum of `n` terms differs from `½ log ((1 + x)/(1 - x))` by at most
  -- `|x|^(2n+1)/(1 - x²)`. We use `|x|^(2n+1) ≤ x^(2n)`, which is free of absolute values.
  have hseries := (Real.sum_range_sub_log_div_le hx n).trans
    (div_le_div_of_nonneg_right (pow_le_pow_of_le_one (abs_nonneg _) hx.le (Nat.le_succ _)) hsq.le)
  rw [(even_two_mul n).pow_abs] at hseries
  rw [abs_le] at htwo hseries
  rw [hlog, logApprox, logErr]
  constructor <;> linarith [htwo.1, htwo.2, hseries.1, hseries.2]

/-! ## The exponent `q` of the query time (Corollary 31) -/

/-! ## The exponent `γ` (Corollary 31) -/

/-! ## The bound `R_c(γ)` (equation (11)) -/

/-! ## Entries of Table 2

An entry needs two facts on its row, `hg` (an enclosure of `ln(1/ρ_c)/ln 4`) and `hR` (the `ε` of
the row is below `R_c(γ)` up to the largest `γ` of the row), and two rational numbers `θlo ≤ θhi`
that enclose its `θ`. -/

end ThreeSumApsp
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec5.Corollary39.GraphH.Fill
/-!
# Corollary 39 in the light language: filling a weight matrix of H

Corollary 39: a clique on k vertices of a k-partite graph G is a triangle of a graph H.
The vertices of the first t = k - 3⌊k/3⌋ parts are fixed, and the other parts form three groups.
With gw = ⌊k/3⌋, a vertex of H is a choice of one vertex in each of the gw parts of a group, so a
part of H has N = n^gw vertices; the chosen vertices are the digits of its number in base n, lowest
first.
The weight between a vertex of H in group i and one in group i + 1 is the sum of the weights of G
over a list L of pairs of parts: the pairs between the two groups, the pairs inside group i, the
pairs of a fixed part and a part of group i, and, for one of the three matrices, the pairs of fixed
parts.  This file writes one N × N weight matrix of H, for any list L.

The weights of G stand at g: the weight between vertex u of part p and vertex v of part q in cell
g + ((p k + q) n + u) n + v.  The current choice of one vertex in each part stands in the k cells at
vt.  The weights lie below these cells and below the matrix.  `sumPairs` adds up the weights over L.
`fillMat` runs through all pairs (row, column) of vertices of H by two odometers inside vt and
writes the sums, row by row.

The proof goes from the inside to the outside: the address of a weight (`wAddr_gives`), the sum
(`sumPairs_spec`), one step of an odometer (`inc_spec`, and `incCol_spec` for the column odometer),
and then one invariant, `FillInv`, "before column c of row r all earlier entries are written", which
an entry (`fillEntry_spec`), a row (`fillRow_spec`) and the step to the next row (`nextRow_spec`)
take to its next value; `fillMat_spec` puts them together.
-/

@[expose] public section

namespace Light.Sec5

open ThreeSumApsp

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

/-! ## The local variables

The local variables 0 to 16 of the hosts for k-Clique, except local 3, which each host uses in its
own way.  `fillMat` and its parts read `Size`, `Weights`, `Side` and `Choice` and assign `Row`,
`Col`, `Place` and `Entry`. -/

namespace KC

/-- n, the number of vertices of a part. -/
abbrev Size : ℕ := 0

/-- g, the address of the weights. -/
abbrev Weights : ℕ := 2

/-- The address of the k cells of the chosen vertices. -/
abbrev Choice : ℕ := 6

end KC

open KC

/-! ## The sum of the weights over a list of pairs of parts -/

/-- The address of the weight between the chosen vertices of the parts p and q. -/
def wAddr (kk p q : ℕ) : Expr :=
  Light.v Light.Sec5.KC.Weights +' ((Light.k (p * kk + q) *' Light.v Light.Sec5.KC.Size +' M (Light.v Light.Sec5.KC.Choice +' Light.k p)) *' Light.v Light.Sec5.KC.Size +' M (Light.v Light.Sec5.KC.Choice +' Light.k q))

@[simp] theorem wAddr_cost (kk p q : ℕ) : (wAddr kk p q).cost = 17 := rfl

section sum

variable {μ μ' loc : ℕ → ℤ} {kk n g N vt U : ℕ} {vf : ℕ → ℕ}

end sum

/-! ## The odometers -/

section odometer

variable {μ loc : ℕ → ℤ} {vf vf0 : ℕ → ℕ} {kk n vt off gw ro co r c : ℕ}

end odometer

/-! ## An entry, a row, and the whole matrix -/

section fill

end fill

end Light.Sec5
end
end

section
-- Source module: ThreeSumApsp.Programs.Sec5.Corollary40.DataStructureBounds
/-!
# Time, space and limits of the relocatable data structure

The programs for Corollary 40 call the preprocessing and the query routine of a data structure for
the entries of a thin matrix product, placed at an address fr of their choice.  `KitRates` collects
what they use about it: bounds on the two running times and on the number of cells, and limits
within which the routines run for matrices of zeros and ones.  The limits are a function of the
number sp of memory cells and serve every structure that lies within the memory, so that several
structures can lie side by side.  `kit31_rates` proves this for all parameters and all bounds that
dominate the two costs of Theorem 30; `kitRates_corollary26` (Corollary 26) and
`kitRates_corollary31` (Corollary 31) are its instances.
-/

@[expose] public section

namespace Light.Sec5

open ThreeSumApsp Light.Sec4

/-! ## The bounds on the two running times -/

/-! ## The limits -/

/-- The first cell after an input with no wanted positions: three sizes and the two matrices. -/
def kitInputEnd (N D₀ : ℕ) : ℕ := 3 + 2 * (N * D₀)

/-- The limits of Section 4 for a structure that lies directly after such an input, for matrices of
zeros and ones. -/
def baseLim (G : RatParams) (N D₀ : ℕ) : Limits := lim31 G N D₀ (kitInputEnd N D₀) 0

/-- Limits for the routines with general parameters on matrices of zeros and ones, with `sp` cells
of memory.  Only the number of cells depends on the place of a structure.  So the bound on the words
and the depth are those of `baseLim`, with `sp` more for the addresses and 20 more levels of calls
for the procedures that call the routines. -/
def limKit31 (G : RatParams) (N D₀ sp : ℕ) : Limits :=
  ⟨(baseLim G N D₀).word + sp, sp, (baseLim G N D₀).depth + 20⟩

@[simp] theorem limKit31_word (G : RatParams) (N D₀ sp : ℕ) :
    (limKit31 G N D₀ sp).word = (baseLim G N D₀).word + sp :=
  rfl

@[simp] theorem limKit31_space (G : RatParams) (N D₀ sp : ℕ) : (limKit31 G N D₀ sp).space = sp :=
  rfl

@[simp] theorem limKit31_depth (G : RatParams) (N D₀ sp : ℕ) :
    (limKit31 G N D₀ sp).depth = (baseLim G N D₀).depth + 20 :=
  rfl

/-! ## What the programs use about a kit -/

section

end

/-! ## The data structure of Section 4 as a kit -/

end Light.Sec5
end
end

section
-- Source module: ThreeSumApsp.AllEdges.ScanPairs
/-!
# All-edges Exact Triangle: reading the answers of one instance and marking pairs

The all-edges form of the scans of Theorem 17.  The host keeps an array of `n²` marks, one for each
pair `(a, b)`; the mark of `(a, b)` is set once a zero triangle through `(a, b)` has been found.

scanPairsAE(out, qa, qb, w, fl, ab, bc, ac, n, c0, len): for every query pair number `i < w` whose
answer `out[i]` is not 0 and whose pair `(qa[i], qb[i])` is not marked yet, scan the piece
`{c0, …, c0 + len - 1}` for a zero triangle through the pair, and store the result as its mark.

A scan that succeeds sets a mark that was not set before, so over the whole run there are at most
`n²` successful scans; a scan that fails is an accepted pair whose scan fails, as in Theorem 17.
-/

@[expose] public section

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

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

/-! ## The pure side -/

/-- The marks after the first `i` query pairs, starting from the marks `g`: the place `pl j` of an
accepted query pair whose scan succeeds is marked. -/
def marksUpto (acc hit : ℕ → Bool) (pl : ℕ → ℕ) (g : ℕ → Bool) (i : ℕ) : ℕ → Bool :=
  fun q => g q || (List.range i).any fun j => acc j && hit j && decide (pl j = q)

/-- Whether the query pair number `i` is scanned: it is accepted and its place is not marked. -/
def aeExecAt (acc hit : ℕ → Bool) (pl : ℕ → ℕ) (g : ℕ → Bool) (i : ℕ) : Bool :=
  acc i && !marksUpto acc hit pl g i (pl i)

/-- The number of scans among the first `w` query pairs. -/
def aeExecsUpto (acc hit : ℕ → Bool) (pl : ℕ → ℕ) (g : ℕ → Bool) (w : ℕ) : ℕ :=
  ((List.range w).filter (aeExecAt acc hit pl g)).length

/-- The number of marked places among the first `N`. -/
def markCount (N : ℕ) (g : ℕ → Bool) : ℕ := ((List.range N).filter g).length

/-! ## The marks in the memory -/

/-- The memory `μ` with the `N` cells from `fl` holding the marks `g`. -/
def withMarks (μ : ℕ → ℤ) (fl N : ℕ) (g : ℕ → Bool) : ℕ → ℤ :=
  fun a => if fl ≤ a ∧ a < fl + N then Light.Sec3.bit (g (a - fl)) else μ a

/-! ## The routine -/

namespace ScanPairsAE

/-- The locals of `scanPairsAE`.  The arguments: the addresses of the answers and of the rows and
the columns of the query pairs, their number `w`, the address of the marks, the addresses of the
three arrays of weights, `n`, and the first vertex and the number of vertices of the piece.  Then
the counter, the place of the pair, and the result of a scan. -/
abbrev Out : ℕ := 0
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Rows : ℕ := 1
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Cols : ℕ := 2
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Num : ℕ := 3
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Fl : ℕ := 4
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev AdrAB : ℕ := 5
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev AdrBC : ℕ := 6
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev AdrAC : ℕ := 7
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Size : ℕ := 8
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev First : ℕ := 9
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Len : ℕ := 10
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Cnt : ℕ := 11
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Pl : ℕ := 12
@[inherit_doc Light.Sec3.ScanPairsAE.Out] abbrev Res : ℕ := 13

end ScanPairsAE

open ScanPairsAE in
/-- One query pair: if out[i] ≠ 0 then pl := qa[i] n + qb[i]; if fl[pl] = 0 then
fl[pl] := scan(ab, bc, ac, n, qa[i], qb[i], c0, len). -/
def scanPairsAEStep (pScan : ℕ) : Stmt :=
  .ite (M (Light.v Light.Sec3.ScanPairsAE.Out +' Light.v Light.Sec3.ScanPairsAE.Cnt) =' Light.k 0) .skip
    (.set Pl (M (Light.v Light.Sec3.ScanPairsAE.Rows +' Light.v Light.Sec3.ScanPairsAE.Cnt) *' Light.v Light.Sec3.ScanPairsAE.Size +' M (Light.v Light.Sec3.ScanPairsAE.Cols +' Light.v Light.Sec3.ScanPairsAE.Cnt)) ;;
      .ite (M (Light.v Light.Sec3.ScanPairsAE.Fl +' Light.v Pl) =' Light.k 0)
        (.call pScan [Light.v Light.Sec3.ScanPairsAE.AdrAB, Light.v Light.Sec3.ScanPairsAE.AdrBC, Light.v Light.Sec3.ScanPairsAE.AdrAC, Light.v Light.Sec3.ScanPairsAE.Size, M (Light.v Light.Sec3.ScanPairsAE.Rows +' Light.v Light.Sec3.ScanPairsAE.Cnt), M (Light.v Light.Sec3.ScanPairsAE.Cols +' Light.v Light.Sec3.ScanPairsAE.Cnt),
          Light.v Light.Sec3.ScanPairsAE.First, Light.v Light.Sec3.ScanPairsAE.Len] Light.Sec3.ScanPairsAE.Res ;;
          .store (Light.v Light.Sec3.ScanPairsAE.Fl +' Light.v Pl) (Light.v Light.Sec3.ScanPairsAE.Res))
        .skip)

open ScanPairsAE in
/-- scanPairsAE(out, qa, qb, w, fl, ab, bc, ac, n, c0, len); the parameter is the procedure number
of `scan`. -/
def scanPairsAEBody (pScan : ℕ) : Stmt :=
  .set Light.Sec3.ScanPairsAE.Cnt (Light.k 0) ;;
  .while (Light.v Light.Sec3.ScanPairsAE.Cnt <' Light.v Light.Sec3.ScanPairsAE.Num) (scanPairsAEStep pScan ;; .set Light.Sec3.ScanPairsAE.Cnt (Light.v Light.Sec3.ScanPairsAE.Cnt +' Light.k 1)) ;;
  .set 0 (Light.k 0)

/-- The time of one round, given whether a scan is made. -/
def tAERound (len : ℕ) (exec : Bool) : ℕ := 60 + if exec then tScanCall len + 20 else 0

/-- The time for reading w answers and marks. -/
def tAEAnswers (w : ℕ) : ℕ := 70 * w + 8

/-- The time of scanPairsAE, given the number of scans. -/
def tScanPairsAE (w len execs : ℕ) : ℕ := tAEAnswers w + (tScanCall len + 20) * execs

section

variable {pScan : ℕ} {μ : ℕ → ℤ} {out qa qb w fl ab bc ac n c0 len U : ℕ} {OUT AB BC AC : List ℤ}
  {QA QB : List ℕ} {g : ℕ → Bool}

/-- The place of the query pair number `i`. -/
abbrev plOf (n : ℕ) (QA QB : List ℕ) (i : ℕ) : ℕ := QA.getD i 0 * n + QB.getD i 0

/-- What `scanPairsAE` assumes about where the arrays lie: everything it reads lies outside the
marks, and the marks lie in the memory. -/
structure MarksApart (lim : Limits) (fl out qa qb w ab bc ac n : ℕ) : Prop where
  out : Apart out w fl (n * n)
  qa : Apart qa w fl (n * n)
  qb : Apart qb w fl (n * n)
  ab : Apart ab (n * n) fl (n * n)
  bc : Apart bc (n * n) fl (n * n)
  ac : Apart ac (n * n) fl (n * n)
  space : fl + n * n ≤ lim.space

/-- The state of `scanPairsAE` when `i` query pairs have been treated and the counter is `j`. -/
def scanPairsAEState (μ : ℕ → ℤ) (out qa qb w fl ab bc ac n c0 len : ℕ) (OUT AB BC AC : List ℤ)
    (QA QB : List ℕ) (g : ℕ → Bool) (i j : ℕ) (pl r : ℤ) : State :=
  ⟨Light.frame [out, qa, qb, w, fl, ab, bc, ac, n, c0, len, j, pl, r],
    Light.Sec3.withMarks μ fl (n * n)
      (marksUpto (accOf OUT) (hitOf n AB BC AC QA QB c0 len) (plOf n QA QB) g i)⟩

end

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.Task
/-!
# All-edges Exact Triangle as a task

All-edges Exact Triangle: for every pair `(a, b)`, whether some `c` has
`w(a,b) + w(b,c) + w(a,c) = 0`.  The answers are `n²` flags, cell `a n + b` for the pair `(a, b)`.

`aeModel` is the reading of "solved in time T" of `lightModel` with Exact Triangle replaced by its
all-edges form, so that the deductions of the paper between running times (Theorems 17 and 19)
apply to it unchanged.

The brute force (`aeBruteBody`) scans all of `C` for every pair.
-/

@[expose] public section

open ThreeSumApsp

namespace Light.Sec3

open ThreeSumApsp.Spec

/-! ## The task -/

/-- Three `n × n` matrices of weights, and the place for `n²` flags. -/
structure AEInst : Type extends TriInst where
  out : ℕ

/-- The matrices and the place for the flags lie below the free pointer; the flags do not meet the
matrices. -/
structure AEInst.Pre (x : AEInst) (μ : ℕ → ℤ) (fr : ℕ) : Prop
    extends TriInst.Pre x.toTriInst μ fr where
  belowOut : x.out + x.n * x.n ≤ fr
  apartAB : Apart x.ab (x.n * x.n) x.out (x.n * x.n)
  apartBC : Apart x.bc (x.n * x.n) x.out (x.n * x.n)
  apartAC : Apart x.ac (x.n * x.n) x.out (x.n * x.n)

/-- The flags: cell `q = a n + b` holds 1 if some `c < n` has `w(a,b) + w(b,c) + w(a,c) = 0`, and 0
if not. -/
noncomputable def aeFlags (n : ℕ) (AB BC AC : List ℤ) : List ℤ :=
  (List.range (n * n)).map fun q =>
    flag (∃ c < n, AB.getD q 0 + BC.getD (q % n * n + c) 0 + AC.getD (q / n * n + c) 0 = 0)

/-- **All-edges Exact Triangle**: ae(n, U, ab, bc, ac, out, fr) writes the `n²` flags to out. -/
noncomputable def aeTask : Task where
  Inst := AEInst
  size x := x.n
  bound x := x.U
  args x := [x.n, x.U, x.ab, x.bc, x.ac, x.out]
  Pre := AEInst.Pre
  Post x μ fr _ μ' := Seg μ' x.out (aeFlags x.n x.AB x.BC x.AC) ∧
    KeptBut μ μ' fr x.out (x.n * x.n)

end Light.Sec3

namespace Light

/-- The reading of `lightModel`, with all-edges Exact Triangle in the place of Exact Triangle. -/
noncomputable def aeModel : DetTimeModel :=
  { lightModel with exactTriangle := SolvedIn Sec3.aeTask }

end Light

namespace Light.Sec3

open ThreeSumApsp.Spec

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

/-! ## The flags as scans -/

/-- The flag of the place `q`, as a scan of all of `C`. -/
def aeHit (n : ℕ) (AB BC AC : List ℤ) (q : ℕ) : Bool := scanHit n AB BC AC (q / n) (q % n) 0 n

/-! ## The brute force -/

namespace AEBrute

/-- The locals of `aeBrute`.  The arguments: `n`, `U`, the addresses of the three arrays and of the
flags, and the free pointer.  Then the vertices `a` and `b`, and the result of a scan. -/
abbrev Size : ℕ := 0

@[inherit_doc Light.Sec3.AEBrute.Size] abbrev AdrAB : ℕ := 2
@[inherit_doc Light.Sec3.AEBrute.Size] abbrev AdrBC : ℕ := 3
@[inherit_doc Light.Sec3.AEBrute.Size] abbrev AdrAC : ℕ := 4
@[inherit_doc Light.Sec3.AEBrute.Size] abbrev Out : ℕ := 5

@[inherit_doc Light.Sec3.AEBrute.Size] abbrev VtxA : ℕ := 7
@[inherit_doc Light.Sec3.AEBrute.Size] abbrev VtxB : ℕ := 8
@[inherit_doc Light.Sec3.AEBrute.Size] abbrev Res : ℕ := 9

end AEBrute

open AEBrute in
/-- One pair `(a, b)`: out[a n + b] := scan(ab, bc, ac, n, a, b, 0, n). -/
def aeBrutePair (pScan : ℕ) : Stmt :=
  .call pScan [Light.v Light.Sec3.AEBrute.AdrAB, Light.v Light.Sec3.AEBrute.AdrBC, Light.v Light.Sec3.AEBrute.AdrAC, Light.v Light.Sec3.AEBrute.Size, Light.v VtxA, Light.v VtxB, Light.k 0, Light.v Light.Sec3.AEBrute.Size] Light.Sec3.AEBrute.Res ;;
  .store ((Light.v Light.Sec3.AEBrute.Out +' (Light.v VtxA *' Light.v Light.Sec3.AEBrute.Size)) +' Light.v VtxB) (Light.v Light.Sec3.AEBrute.Res)

open AEBrute in
/-- The inner loop: all `b` for one `a`. -/
def aeBruteInner (pScan : ℕ) : Stmt := Stmt.for VtxB (Light.v Light.Sec3.AEBrute.Size) (aeBrutePair pScan)

open AEBrute in
/-- aeBrute(n, U, ab, bc, ac, out, fr); the parameter is the procedure number of `scan`. -/
def aeBruteBody (pScan : ℕ) : Stmt :=
  Stmt.for VtxA (Light.v Light.Sec3.AEBrute.Size) (aeBruteInner pScan) ;; .set 0 (Light.k 0)

/-- The time of aeBrute. -/
def tAEBrute (n : ℕ) : ℕ := 200 * n ^ 3 + 200

/-- What aeBrute needs: sums of three weights and the places of the flags, no cells, one call. -/
def aeBruteNeed (_n U : ℕ) : Need := ⟨3 * U + 1, 0, 1⟩

/-- The flags of the first `j` places, written. -/
abbrev aeWrote (μ : ℕ → ℤ) (out n : ℕ) (AB BC AC : List ℤ) (j : ℕ) : ℕ → ℤ :=
  wrote μ out (fun q => Light.Sec3.bit (aeHit n AB BC AC q)) j

section

variable {pScan : ℕ} {x : AEInst} {μ : ℕ → ℤ} {fr : ℕ}

end

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.Data
/-!
# All-edges host: the marks, as pure data

The all-edges host runs the instances of Theorem 17 and keeps a mark for every pair `(a, b)`.
`HostData.marks t` are the marks before instance `t`; `HostData.aeExecs t` is the number of scans
of instance `t`.  A scan is a failed scan or sets a new mark, so all instances together make at most
`F(p) + n²` scans (`HostData.sum_aeExecs_le`).  After all instances the mark of `(a, b)` is set if
and only if some `c` has `S(a,b,c) = 0` (`HostData.marks_m_iff`).
-/

@[expose] public section

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

namespace HostData

variable (X : HostData)

/-- The place of the query pair number `i` of instance `t`, as scanPairsAE computes it. -/
abbrev aePl (t : ℕ) : ℕ → ℕ := plOf X.n (X.WI t) (X.WJ t)

/-- The marks before instance `t`. -/
def marks : ℕ → ℕ → Bool
  | 0 => fun _ => false
  | t + 1 => marksUpto (X.acc t) (X.hit t) (X.aePl t) (marks t) (X.w t)

/-- The number of scans of instance `t`. -/
def aeExecs (t : ℕ) : ℕ := aeExecsUpto (X.acc t) (X.hit t) (X.aePl t) (X.marks t) (X.w t)

variable {X}

end HostData

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.Loop
/-!
# All-edges host: the loop over the instances

The loop of the host of Theorem 17, with `scanPairsAE` in the place of `scanPairs`: the marks of the
`n²` pairs lie in `n²` cells from `fl`, below the matrices of the instances.  The loop writes, for
every pair, whether a zero triangle through it has been found (`aeLoop_spec`).

The memory before instance `t` agrees below `x` with the first memory, with the marks `marks t` in
the cells from `fl`; so the facts `HostSetting` of the loop of Theorem 17 hold with that memory as
the first one, and the lemmas about the calls of writeX, writeY and the solver are used unchanged.
-/

@[expose] public section

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

variable {lim : Limits} {d : ℕ}

/-! ## The procedure -/

namespace AELocal

/-- The address of the marks; the arguments before it are those of hostLoop. -/
abbrev Fl : ℕ := 20
/-- The number t of the instance. -/
abbrev Inst : ℕ := 21
/-- The number ch of its chunk. -/
abbrev ChunkNo : ℕ := 22
/-- The first vertex c0 of its piece. -/
abbrev PieceStart : ℕ := 23

end AELocal

open HostLocal AELocal

/-- The four calls for one instance: the two matrices, the solver, the scans. -/
def aeCalls (pS pWriteX pWriteY pScanPairs : ℕ) : Stmt :=
  .call pWriteX [Light.v Light.Sec3.HostLocal.MatX, Light.v Light.Sec3.HostLocal.ResAC, Light.v Light.Sec3.HostLocal.Size, Light.v Light.Sec3.HostLocal.ParD, Light.v Light.Sec3.HostLocal.ThePrime, Light.v AELocal.PieceStart, Light.v Light.Sec3.HostLocal.Len,
    Light.v Residue] Light.Sec3.HostLocal.Unused ;;
  .call pWriteY [Light.v Light.Sec3.HostLocal.MatY, Light.v Light.Sec3.HostLocal.ResBC, Light.v Light.Sec3.HostLocal.Size, Light.v Light.Sec3.HostLocal.ParD, Light.v Light.Sec3.HostLocal.ThePrime, Light.v AELocal.PieceStart, Light.v Light.Sec3.HostLocal.Len] Light.Sec3.HostLocal.Unused ;;
  .call pS [Light.v Light.Sec3.HostLocal.Size, Light.v Light.Sec3.HostLocal.ParD, Light.v NumPairs, Light.k 1, Light.v Light.Sec3.HostLocal.MatX, Light.v Light.Sec3.HostLocal.MatY, Light.v Light.Sec3.HostLocal.Rows +' Light.v Light.Sec3.HostLocal.Start, Light.v Light.Sec3.HostLocal.Cols +' Light.v Light.Sec3.HostLocal.Start,
    Light.v Light.Sec3.HostLocal.AdrOut, Light.v Light.Sec3.HostLocal.SolverFree] Light.Sec3.HostLocal.Unused ;;
  .call pScanPairs [Light.v Light.Sec3.HostLocal.AdrOut, Light.v Light.Sec3.HostLocal.Rows +' Light.v Light.Sec3.HostLocal.Start, Light.v Light.Sec3.HostLocal.Cols +' Light.v Light.Sec3.HostLocal.Start, Light.v NumPairs, Light.v Fl, Light.v Light.Sec3.HostLocal.AdrAB,
    Light.v Light.Sec3.HostLocal.AdrBC, Light.v Light.Sec3.HostLocal.AdrAC, Light.v Light.Sec3.HostLocal.Size, Light.v AELocal.PieceStart, Light.v Light.Sec3.HostLocal.Len] Light.Sec3.HostLocal.Unused

/-- The numbers of the next instance, of its chunk, and the first vertex of its piece. -/
def aeNext : Stmt :=
  .set AELocal.ChunkNo (Light.v AELocal.ChunkNo +' Light.k 1) ;;
  .ite (Light.v AELocal.ChunkNo =' Light.v Light.Sec3.HostLocal.NumChunks)
    (.set AELocal.ChunkNo (Light.k 0) ;; .set AELocal.PieceStart (Light.v AELocal.PieceStart +' Light.v Light.Sec3.HostLocal.PieceLen))
    .skip ;;
  .set AELocal.Inst (Light.v AELocal.Inst +' Light.k 1)

/-- The parameters of the instance. -/
def aeParams : Stmt :=
  .set Light.Sec3.HostLocal.Len (Light.v Light.Sec3.HostLocal.PieceLen) ;;
  .ite (Light.v Light.Sec3.HostLocal.Size -' Light.v AELocal.PieceStart <' Light.v Light.Sec3.HostLocal.PieceLen) (.set Light.Sec3.HostLocal.Len (Light.v Light.Sec3.HostLocal.Size -' Light.v AELocal.PieceStart))
    .skip ;;
  .set Residue (M (Light.v Light.Sec3.HostLocal.TabR +' Light.v AELocal.ChunkNo)) ;;
  .set Light.Sec3.HostLocal.Start (M (Light.v Light.Sec3.HostLocal.TabL +' Light.v AELocal.ChunkNo)) ;;
  .set NumPairs (M (Light.v Light.Sec3.HostLocal.TabW +' Light.v AELocal.ChunkNo))

/-- One round of aeLoop. -/
def aeRound (pS pWriteX pWriteY pScanPairs : ℕ) : Stmt :=
  aeParams ;;
  aeCalls pS pWriteX pWriteY pScanPairs ;;
  aeNext

/-- aeLoop(n, D, p, q, chunkCount, m, ab, bc, ac, rac, rbc, qi, qj, cr, cl, cw, x, y, out, fr, fl).
-/
def aeLoopBody (pS pWriteX pWriteY pScanPairs : ℕ) : Stmt :=
  .set AELocal.Inst (Light.k 0) ;; .set AELocal.ChunkNo (Light.k 0) ;; .set AELocal.PieceStart (Light.k 0) ;;
  .while (Light.v AELocal.Inst <' Light.v Light.Sec3.HostLocal.NumInst) (aeRound pS pWriteX pWriteY pScanPairs) ;;
  .set 0 (Light.k 0)

/-- The number of steps of aeLoop, if the solver takes Tn. -/
def tAELoop (Tn : List ℕ → ℕ) (X : HostData) : ℕ :=
  (∑ t ∈ Finset.range X.m, (tWrites X.n X.D (X.len t) + Tn [X.n, X.D, X.w t] +
    tScanPairsAE (X.w t) (X.len t) (X.aeExecs t))) + 14

/-! ## What the loop relies on -/

/-- The marks lie below x, apart from everything that the loop reads. -/
structure AEMarksLay (X : HostData) (A : HostAddr) (fl : ℕ) : Prop where
  below : fl + X.n * X.n ≤ A.x
  ab : Apart A.ab (X.n * X.n) fl (X.n * X.n)
  bc : Apart A.bc (X.n * X.n) fl (X.n * X.n)
  ac : Apart A.ac (X.n * X.n) fl (X.n * X.n)
  rac : Apart A.rac (X.n * X.n) fl (X.n * X.n)
  rbc : Apart A.rbc (X.n * X.n) fl (X.n * X.n)
  qi : Apart A.qi (X.n * X.n) fl (X.n * X.n)
  qj : Apart A.qj (X.n * X.n) fl (X.n * X.n)
  cr : Apart A.cr X.chunkCount fl (X.n * X.n)
  cl : Apart A.cl X.chunkCount fl (X.n * X.n)
  cw : Apart A.cw X.chunkCount fl (X.n * X.n)

/-- The context of aeLoop: a solver of Lop-AE-SparseTri, and the procedures writeX, writeY,
scanPairsAE, scan in the program. -/
structure AEHostCtx (P₀ R : Program) (pS pWriteX pWriteY pScanPairs pScan : ℕ) (Tn : List ℕ → ℕ)
    (need : List ℕ → Light.Need) : Prop where
  sol : SolvesN lopDetectTask P₀ pS Tn need
  writeX : (P₀ ++ R)[pWriteX]? = some writeXBody
  writeY : (P₀ ++ R)[pWriteY]? = some writeYBody
  scanPairs : (P₀ ++ R)[pScanPairs]? = some (scanPairsAEBody pScan)
  scan : (P₀ ++ R)[pScan]? = some scanBody

/-- The arguments of aeLoop. -/
def aeLoopArgs (X : HostData) (A : HostAddr) (fl : ℕ) : List ℤ :=
  [X.n, X.D, X.p, X.q, X.chunkCount, X.m, A.ab, A.bc, A.ac, A.rac, A.rbc, A.qi, A.qj, A.cr, A.cl,
    A.cw, A.x, A.y, A.out, A.fr, fl]

/-! ## Marks in the memory -/

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

variable {P₀ R : Program} {pS pWriteX pWriteY pScanPairs pScan : ℕ} {Tn : List ℕ → ℕ}
  {need : List ℕ → Need} {X : HostData} {U : ℕ} {A : HostAddr} {fl : ℕ} {μ μ' : ℕ → ℤ} {t : ℕ}

/-- The state of aeLoop: the arguments, the numbers t, ch and c0 of the instance, and what the
locals Len, Residue, Start, NumPairs and Unused hold. -/
abbrev aeState (X : HostData) (A : HostAddr) (fl t ch c0 : ℕ) (len rho lo w res : ℤ)
    (μ' : ℕ → ℤ) : State :=
  ⟨Light.frame [X.n, X.D, X.p, X.q, X.chunkCount, X.m, A.ab, A.bc, A.ac, A.rac, A.rbc, A.qi, A.qj, A.cr,
    A.cl, A.cw, A.x, A.y, A.out, A.fr, fl, t, ch, c0, len, rho, lo, w, res], μ'⟩

/-- The memory before instance `t`, below `x`: the first memory with the marks of `t`. -/
abbrev aeBase (X : HostData) (fl : ℕ) (μ : ℕ → ℤ) (t : ℕ) : ℕ → ℤ :=
  withMarks μ fl (X.n * X.n) (X.marks t)

/-- The time of the four calls. -/
def tAECalls (Tn : List ℕ → ℕ) (X : HostData) (t : ℕ) : ℕ :=
  tWriteX X.n X.D (X.len t) + tWriteY X.n X.D (X.len t) + Tn [X.n, X.D, X.w t] +
    tScanPairsAE (X.w t) (X.len t) (X.aeExecs t) + 52

/-! ## The loop -/

/-- The invariant of the loop, before instance t. -/
def AEInv (lim : Limits) (d : ℕ) (X : HostData) (U : ℕ) (A : HostAddr) (fl : ℕ)
    (need : List ℕ → Light.Need) (μ : ℕ → ℤ) (t : ℕ) (σ : State) : Prop :=
  ∃ (len rho lo w res : ℤ) (μ' : ℕ → ℤ),
    σ = aeState X A fl t (t % X.chunkCount) (X.c0 t) len rho lo w res μ' ∧
    HostSetting lim d X U A need (aeBase X fl μ t) μ'

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.Host
/-!
# All-edges host: the top procedures

The host of all-edges Exact Triangle has two procedures above the loop.

* aeTop(n, U, ab, bc, ac, out, fr) fills the `n²` flags with 0, puts the address `out` into the cell
  `fr`, and calls aeCore with the free pointer `fr + 1`.
* aeCore(n, U, ab, bc, ac, fr) is et17 with the loop of the all-edges form; the address of the flags
  is read from the cell `fr - 1`.  The first two parts of et17 (the parameters; the prime, the sizes
  and the addresses) are used unchanged, with their specifications.  In the small case aeCore calls
  the all-edges brute force.

The procedures of the host of Theorem 17 stay in the program, so that its context `Et17Ctx` holds and
the specifications of the parts can be used.
-/

@[expose] public section

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

/-- aeTop(n, U, ab, bc, ac, out, fr). -/
def aeTopBody (pFill pCore : ℕ) : Stmt :=
  .call pFill [Light.v 5, Light.v 0 *' Light.v 0, Light.k 0] 7 ;;
  .store (Light.v 6) (Light.v 5) ;;
  .call pCore [Light.v 0, Light.v 1, Light.v 2, Light.v 3, Light.v 4, Light.v 6 +' Light.k 1] 7 ;;
  .set 0 (Light.k 0)

/-! ## Time -/

/-- The time of the third part of aeCore in a run. -/
def aeRunTime (Tn : List ℕ → ℕ) (X : HostData) (U : ℕ) : ℕ :=
  tDblTable (bitLen U) + 3 * tResidues (X.n * X.n) (bitLen U) + tClasses X.n X.p
    + tChunks X.p X.chunkCount + tAELoop Tn X + 80

/-! ## The context -/

variable {P₀ R : Program} {ν : Et17Nums} {pAEBrute pAELoop pScanPairsAE pCore pFill : ℕ}
  {Tn : List ℕ → ℕ} {need : List ℕ → Need} {Dfun Gfun tD tG wD wG : ℕ → ℕ}

/-! ## The third part -/

/-! ## The time of a run -/

namespace AETime

variable {x : TriInst} {μ : ℕ → ℤ} {fr D g : ℕ}

end AETime

/-! ## The procedures -/

/-! ## The program, assembled -/
                            -- 33

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.MidHost
/-!
# All-edges host with the true middle dimension

The all-edges host of `AllEdges/Host.lean` declares the inner dimension `D` for the matrices of its
instances and for the calls of the solver.  The middle part `C_k × ℤ_p` of an instance has at most
`q p ≤ q s` vertices, where `q = ⌈s/g⌉` and `s = ⌊√D⌋`.  This host declares the inner dimension
`midDim D g = q s` instead (`midTables`, `midCoreBody`); everything else is as before.

The loop over the instances is generic in the field `D` of its data, so it runs with the data
`midHostData x D g`, which is `hostData x D g` with `D := q s`.  The arrays stay where they are
(they were laid out for `D ≥ q s`); the need asks in addition for the need of the solver at the
middle dimension (`midNeedAt`).
-/

@[expose] public section

namespace ThreeSumApsp

/-- The middle dimension of the instances of Theorem 17. -/
def midDim (D g : ℕ) : ℕ := Spec.pieceSizeNat D g * Nat.sqrt D

end ThreeSumApsp

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

/-! ## The data with another inner dimension -/

/-- The data `X` with the inner dimension `D'`. -/
def HostData.withD (X : HostData) (D' : ℕ) : HostData := { X with D := D' }

/-- The data of the loop of the host with the middle dimension. -/
def midHostData (x : TriInst) (D g : ℕ) : HostData := (hostData x D g).withD (midDim D g)

namespace HostData

variable (X : HostData) (D' : ℕ)

@[simp] theorem withD_n : (X.withD D').n = X.n := rfl
@[simp] theorem withD_D : (X.withD D').D = D' := rfl
@[simp] theorem withD_p : (X.withD D').p = X.p := rfl
@[simp] theorem withD_q : (X.withD D').q = X.q := rfl
@[simp] theorem withD_cap : (X.withD D').cap = X.cap := rfl
@[simp] theorem withD_AB : (X.withD D').AB = X.AB := rfl
@[simp] theorem withD_BC : (X.withD D').BC = X.BC := rfl
@[simp] theorem withD_AC : (X.withD D').AC = X.AC := rfl
@[simp] theorem withD_chunkCount : (X.withD D').chunkCount = X.chunkCount := rfl
@[simp] theorem withD_m : (X.withD D').m = X.m := rfl
@[simp] theorem withD_h : (X.withD D').h = X.h := rfl
@[simp] theorem withD_w (t : ℕ) : (X.withD D').w t = X.w t := rfl
@[simp] theorem withD_len (t : ℕ) : (X.withD D').len t = X.len t := rfl

end HostData

/-! ## The middle dimension -/

/-! ## The procedure -/

open Et17 in
/-- The third part of the core: as `aeTables`, with the inner dimension `q s` (the locals
`PieceLen` and `RootD`) for the loop over the instances. -/
def midTables (ν : Et17Nums) (pAELoop : ℕ) : Stmt :=
  .call ν.pDbl [Light.v Light.Sec3.Et17.Free, Light.v ThePrime, Light.v Bits] Light.Sec3.Et17.Small ;;
  .call ν.pResidues [Light.v AdrAB, Light.v ResAB, Light.v SizeSq, Light.v Light.Sec3.Et17.Free, Light.v Bits] Light.Sec3.Et17.Small ;;
  .call ν.pResidues [Light.v AdrBC, Light.v ResBC, Light.v SizeSq, Light.v Light.Sec3.Et17.Free, Light.v Bits] Light.Sec3.Et17.Small ;;
  .call ν.pResidues [Light.v AdrAC, Light.v ResAC, Light.v SizeSq, Light.v Light.Sec3.Et17.Free, Light.v Bits] Light.Sec3.Et17.Small ;;
  .call ν.pClasses [Light.v ResAB, Light.v Light.Sec3.Et17.Size, Light.v ThePrime, Light.v Cls, Light.v Light.Sec3.Et17.Cur, Light.v Light.Sec3.Et17.Rows, Light.v Light.Sec3.Et17.Cols] Light.Sec3.Et17.Small ;;
  .call ν.pChunks [Light.v Cls, Light.v ThePrime, Light.v Cap, Light.v TabR, Light.v TabL, Light.v TabW] NumChunks ;;
  .set NumInst (Light.v NumPieces *' Light.v NumChunks) ;;
  .call pAELoop [Light.v Light.Sec3.Et17.Size, Light.v PieceLen *' Light.v RootD, Light.v ThePrime, Light.v PieceLen, Light.v NumChunks, Light.v NumInst,
    Light.v AdrAB, Light.v AdrBC, Light.v AdrAC, Light.v ResAC, Light.v ResBC, Light.v Light.Sec3.Et17.Rows, Light.v Light.Sec3.Et17.Cols, Light.v TabR, Light.v TabL, Light.v TabW, Light.v Light.Sec3.Et17.MatX,
    Light.v Light.Sec3.Et17.MatY, Light.v AdrOut, Light.v Light.Sec3.Et17.SolverFree, M (Light.v Light.Sec3.Et17.Free -' Light.k 1)] 0

open Et17 in
/-- midCore(n, U, ab, bc, ac, fr), with the address of the flags in the cell fr - 1. -/
def midCoreBody (ν : Et17Nums) (pAEBrute pAELoop : ℕ) : Stmt :=
  et17Params ν ;;
  .ite (Light.v Light.Sec3.Et17.Small =' Light.k 1)
    (.call pAEBrute [Light.v Light.Sec3.Et17.Size, Light.v Light.Sec3.Et17.Bound, Light.v AdrAB, Light.v AdrBC, Light.v AdrAC, M (Light.v Light.Sec3.Et17.Free -' Light.k 1), Light.v Light.Sec3.Et17.Free] 0)
    (et17Sizes ν ;; midTables ν pAELoop)

/-! ## Time -/

/-- A bound on the time of the loop over the instances, with the solver at the middle
dimension. -/
noncomputable def midLoopBound (Tn : List ℕ → ℕ) (n U D g : ℕ) : ℕ :=
  4 * n * g *
      (tWrites n D (pieceSizeNat D g) + supTime Tn n (midDim D g) (queryCapNat n D)
        + tAEAnswers (queryCapNat n D))
    + (tScanCall (pieceSizeNat D g) + 20) * (falsePositiveBound n U D + n * n) + 14

/-- The time of the core after the tests, if n is not small. -/
noncomputable def midHostMain (Tn : List ℕ → ℕ) (n U D g : ℕ) : ℕ :=
  chooseTime n U D + tQueryCapNat n D + tCeilDiv (Nat.sqrt D) g + tCeilDiv n (pieceSizeNat D g)
    + tBitLen U + tDblTable (bitLen U) + 3 * tResidues (n * n) (bitLen U) + tClasses n (Nat.sqrt D)
    + tChunks (Nat.sqrt D) (4 * n * g) + midLoopBound Tn n U D g + 300

/-- A bound on the time of the core in the worst case. -/
noncomputable def midCoreTime (Dfun Gfun tD tG : ℕ → ℕ) (Tn : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  hostSetup tD tG n (Dfun n) +
    if SmallCase n (Dfun n) (Gfun (Dfun n)) then tAEBrute n + 20
    else midHostMain Tn n U (Dfun n) (Gfun (Dfun n))

/-- **A bound on the time of the all-edges host with the middle dimension.** -/
noncomputable def midHostTime (Dfun Gfun tD tG : ℕ → ℕ) (Tn : List ℕ → ℕ) (n U : ℕ) : ℕ :=
  fillTime (n * n) + midCoreTime Dfun Gfun tD tG Tn n U + 40

/-! ## Need -/

/-- The need of the core for given parameters: that of the host of Theorem 17, and the need of the
solver at the middle dimension. -/
def midNeedAt (a b : ℕ) (need : List ℕ → Light.Need) (n U D g : ℕ) : Need :=
  ⟨(hostNeedAt a b need n U D g).word + (supNeed need n (midDim D g) (queryCapNat n D)).word,
    (hostNeedAt a b need n U D g).cells + (supNeed need n (midDim D g) (queryCapNat n D)).cells,
    (hostNeedAt a b need n U D g).depth + (supNeed need n (midDim D g) (queryCapNat n D)).depth⟩

/-- The need of the core. -/
def midCoreNeed (Dfun Gfun wD wG : ℕ → ℕ) (need : List ℕ → Light.Need) (n U : ℕ) : Need :=
  midNeedAt (wD n) (wG (Dfun n)) need n U (Dfun n) (Gfun (Dfun n))

/-- **The need of the all-edges host with the middle dimension.** -/
def midHostNeed (Dfun Gfun wD wG : ℕ → ℕ) (need : List ℕ → Light.Need) (n U : ℕ) : Need :=
  ⟨(midCoreNeed Dfun Gfun wD wG need n U).word + n * n,
    (midCoreNeed Dfun Gfun wD wG need n U).cells + 1,
    (midCoreNeed Dfun Gfun wD wG need n U).depth + 1⟩

/-! ## The context -/

/-- The context of the top procedure: the context of et17, and the procedures of the host. -/
structure MidCtx (P₀ R : Program) (ν : Et17Nums) (pAEBrute pAELoop pScanPairsAE pCore pFill : ℕ)
    (Tn : List ℕ → ℕ) (need : List ℕ → Light.Need) (Dfun Gfun tD tG wD wG : ℕ → ℕ) : Prop where
  et : Et17Ctx P₀ R ν Tn need Dfun Gfun tD tG wD wG
  hScanPairsAE : (P₀ ++ R)[pScanPairsAE]? = some (scanPairsAEBody ν.pScan)
  hAELoop : (P₀ ++ R)[pAELoop]? = some (aeLoopBody ν.pS ν.pWriteX ν.pWriteY pScanPairsAE)
  hAEBrute : (P₀ ++ R)[pAEBrute]? = some (aeBruteBody ν.pScan)
  hCore : (P₀ ++ R)[pCore]? = some (midCoreBody ν pAEBrute pAELoop)
  hFill : (P₀ ++ R)[pFill]? = some fillBody

variable {P₀ R : Program} {ν : Et17Nums} {pAEBrute pAELoop pScanPairsAE pCore pFill : ℕ}
  {Tn : List ℕ → ℕ} {need : List ℕ → Need} {Dfun Gfun tD tG wD wG : ℕ → ℕ}

/-! ## The third part -/

/-! ## The time of a run -/

namespace MidTime

variable {x : TriInst} {μ : ℕ → ℤ} {fr D g : ℕ}

end MidTime

/-! ## The procedures -/

/-! ## The program, assembled -/

/-- The procedures of the host behind those of the host of Theorem 17, which start at the number
`o`: scanPairsAE, aeLoop, aeBrute, midCore, and the top procedure, which gets the number
`o + 33`. -/
def midExtra (pS pD pG o : ℕ) : Program :=
  let ν := et17NumsAt pS pD pG o
  [scanPairsAEBody ν.pScan,                                 -- 29
    aeLoopBody pS ν.pWriteX ν.pWriteY (o + 29),             -- 30
    aeBruteBody ν.pScan,                                    -- 31
    midCoreBody ν (o + 31) (o + 30),                        -- 32
    aeTopBody (o + 14) (o + 32)]                            -- 33

/-- The procedures of the host. -/
def midProcs (pS pD pG o : ℕ) : Program := et17Procs pS pD pG o ++ midExtra pS pD pG o

end Light.Sec3
end
end

section
-- Source module: ThreeSumApsp.AllEdges.RatParams
/-!
# Parameters with rational exponents

For fixed natural numbers `a, b, c, e` with `b, e ≥ 1`:

* `paramDr a b n = ⌊n^{a/b}⌋`, the largest `D` with `D^b ≤ n^a` (`paramDr_le_rpow`,
  `rpow_lt_paramDr_add_one`);
* `paramGr c e D = ⌈D^{c/e}⌉`, the least `g` with `g^e ≥ D^c` (`rpow_le_paramGr`,
  `paramGr_lt_rpow_add_one`).

Both are computed by procedures of the light language (`drBody`, `grBody`): a loop forms the power
`n^a` (or `D^c`), then `rootCeil` finds the root.  For `a ≤ b` the time of `drBody` is `O(n)`, and
for `c ≤ e` that of `grBody` is `O(D)`.
-/

@[expose] public section

namespace ThreeSumApsp

open Spec

/-- `⌊n^{a/b}⌋`: the largest `D` with `D^b ≤ n^a`. -/
def paramDr (a b n : ℕ) : ℕ := rootFloor b (n ^ a)

/-- `⌈D^{c/e}⌉`: the least `g` with `g^e ≥ D^c`. -/
def paramGr (c e D : ℕ) : ℕ := rootCeil e (D ^ c)

end ThreeSumApsp

namespace Light.Sec3

open ThreeSumApsp ThreeSumApsp.Spec

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

/-- dr(n), over the procedure pRoot (rootCeil): `⌊n^{a/b}⌋`.  Local 0 is the argument, local 1
counts the factors and then takes `⌈(n^a + 1)^{1/b}⌉`, local 2 is the product. -/
def drBody (a b pRoot : ℕ) : Stmt :=
  .set 1 (Light.k 0) ;;
  .set 2 (Light.k 1) ;;
  .while (Light.v 1 <' Light.k a) (
    .set 2 (Light.v 2 *' Light.v 0) ;;
    .set 1 (Light.v 1 +' Light.k 1)) ;;
  .call pRoot [Light.k b, Light.v 2 +' Light.k 1] 1 ;;
  .set 0 (Light.v 1 -' Light.k 1)

/-- The time of dr. -/
def tDr (a b n : ℕ) : ℕ := (paramDr a b n + 1) * (16 * b + 27) + 12 * a + 40

/-- A bound on the numbers that dr forms. -/
def wDr (a b n : ℕ) : ℕ := (n ^ a + 1) * (paramDr a b n + 1) + a + b + 1

/-- gr(D), over the procedure pRoot: `⌈D^{c/e}⌉`, the least g with g^e ≥ D^c.  Local 0 is the
argument, local 1 counts the factors, local 2 is their product. -/
def grBody (c e pRoot : ℕ) : Stmt :=
  .set 1 (Light.k 0) ;;
  .set 2 (Light.k 1) ;;
  .while (Light.v 1 <' Light.k c) (
    .set 2 (Light.v 2 *' Light.v 0) ;;
    .set 1 (Light.v 1 +' Light.k 1)) ;;
  .call pRoot [Light.k e, Light.v 2] 0

/-- The time of gr. -/
def tGr (c e D : ℕ) : ℕ := paramGr c e D * (16 * e + 27) + 12 * c + 30

/-- A bound on the numbers that gr forms. -/
def wGr (c e D : ℕ) : ℕ := D ^ c * paramGr c e D + c + e + 1

end Light.Sec3
end
end