feat(Query): query complexity framework with sorting examples

Cslib.lean

@@ -1,6 +1,19 @@
 module  -- shake: keep-all
 
 public import Cslib.Algorithms.Lean.MergeSort.MergeSort
+public import Cslib.Algorithms.Lean.Query.Arith.Defs
+public import Cslib.Algorithms.Lean.Query.Arith.Lemmas
+public import Cslib.Algorithms.Lean.Query.Bounds
+public import Cslib.Algorithms.Lean.Query.Prog
+public import Cslib.Algorithms.Lean.Query.QueryTree
+public import Cslib.Algorithms.Lean.Query.Sort.Insertion.Defs
+public import Cslib.Algorithms.Lean.Query.Sort.Insertion.Lemmas
+public import Cslib.Algorithms.Lean.Query.Sort.IsSort
+public import Cslib.Algorithms.Lean.Query.Sort.LEQuery
+public import Cslib.Algorithms.Lean.Query.Sort.LowerBound
+public import Cslib.Algorithms.Lean.Query.Sort.Merge.Defs
+public import Cslib.Algorithms.Lean.Query.Sort.Merge.Lemmas
+public import Cslib.Algorithms.Lean.Query.Sort.QueryTree
 public import Cslib.Algorithms.Lean.TimeM
 public import Cslib.Computability.Automata.Acceptors.Acceptor
 public import Cslib.Computability.Automata.Acceptors.OmegaAcceptor

Cslib/Algorithms/Lean/Query/Arith/Defs.lean

@@ -0,0 +1,82 @@
+/-
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+module
+
+public import Cslib.Algorithms.Lean.Query.Prog
+
+/-! # Arithmetic Queries and Complex Multiplication
+
+A simple example showing how to use `Prog.cost` with variable/parametrized query costs.
+
+`ArithQuery α` supports addition, subtraction, and multiplication, each with
+independently parametrized costs. Complex number multiplication provides a toy example
+where two algorithms (naive and Gauss's trick) trade multiplications for additions,
+and the optimal choice depends on the cost ratio.
+-/
+
+open Cslib.Query
+
+public section
+
+namespace Cslib.Query
+
+/-- Arithmetic queries: addition, subtraction, and multiplication. -/
+inductive ArithQuery (α : Type) : Type → Type where
+  | add (a b : α) : ArithQuery α α
+  | sub (a b : α) : ArithQuery α α
+  | mul (a b : α) : ArithQuery α α
+
+namespace ArithQuery
+
+/-- Lift `ArithQuery.add a b` into a `Prog` that returns the sum. -/
+@[expose] def doAdd (a b : α) : Prog (ArithQuery α) α := .liftBind (.add a b) .pure
+/-- Lift `ArithQuery.sub a b` into a `Prog` that returns the difference. -/
+@[expose] def doSub (a b : α) : Prog (ArithQuery α) α := .liftBind (.sub a b) .pure
+/-- Lift `ArithQuery.mul a b` into a `Prog` that returns the product. -/
+@[expose] def doMul (a b : α) : Prog (ArithQuery α) α := .liftBind (.mul a b) .pure
+
+/-- An honest oracle interprets arithmetic queries using the actual ring operations. -/
+@[expose] def honest [Add α] [Sub α] [Mul α] {ι : Type} : ArithQuery α ι → ι
+  | .add a b => a + b
+  | .sub a b => a - b
+  | .mul a b => a * b
+
+/-- Weighted cost model for arithmetic queries. Subtraction costs the same as addition
+    (both are linear-time on bignums). -/
+@[expose] def weight (c_add c_mul : Nat) {ι : Type} : ArithQuery α ι → Nat
+  | .add _ _ => c_add
+  | .sub _ _ => c_add
+  | .mul _ _ => c_mul
+
+end ArithQuery
+
+/-- Naive complex multiplication: `(a + bi)(c + di) = (ac - bd) + (ad + bc)i`.
+    Uses 4 multiplications, 1 subtraction, 1 addition. -/
+@[expose] def complexMulNaive (a b c d : α) : Prog (ArithQuery α) (α × α) := do
+  let ac ← ArithQuery.doMul a c
+  let bd ← ArithQuery.doMul b d
+  let ad ← ArithQuery.doMul a d
+  let bc ← ArithQuery.doMul b c
+  let real ← ArithQuery.doSub ac bd
+  let imag ← ArithQuery.doAdd ad bc
+  pure (real, imag)
+
+/-- Gauss's trick for complex multiplication: computes `(a+b)(c+d)` to save one
+    multiplication, at the cost of extra additions and subtractions.
+    Uses 3 multiplications, 2 subtractions, 2 additions. -/
+@[expose] def complexMulGauss (a b c d : α) : Prog (ArithQuery α) (α × α) := do
+  let ac ← ArithQuery.doMul a c
+  let bd ← ArithQuery.doMul b d
+  let apb ← ArithQuery.doAdd a b
+  let cpd ← ArithQuery.doAdd c d
+  let abcd ← ArithQuery.doMul apb cpd
+  let real ← ArithQuery.doSub ac bd
+  let imag ← ArithQuery.doSub abcd (← ArithQuery.doAdd ac bd)
+  pure (real, imag)
+
+end Cslib.Query
+
+end -- public section

Cslib/Algorithms/Lean/Query/Arith/Lemmas.lean

@@ -0,0 +1,68 @@
+/-
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Kim Morrison
+-/
+module
+
+public import Cslib.Algorithms.Lean.Query.Arith.Defs
+import Mathlib.Tactic.Ring
+public import Mathlib.Algebra.Ring.Defs
+
+/-! # Complex Multiplication: Correctness and Cost Analysis
+
+A simple example showing how to use `Prog.cost` with variable/parametrized query costs.
+
+We prove that both `complexMulNaive` and `complexMulGauss` correctly compute
+complex multiplication when given an honest oracle, and compute their exact
+costs under a parametric weight function. The cost theorems hold for *any* oracle
+(not just honest ones), because both algorithms are straight-line (no branching
+on query results).
+-/
+
+open Cslib.Query
+
+public section
+
+namespace Cslib.Query
+
+variable {α : Type}
+
+-- ## Correctness
+
+theorem complexMulNaive_eval_honest [Ring α] (a b c d : α) :
+    (complexMulNaive a b c d).eval ArithQuery.honest = (a * c - b * d, a * d + b * c) := by
+  simp [complexMulNaive, ArithQuery.doMul, ArithQuery.doSub, ArithQuery.doAdd, ArithQuery.honest]
+
+theorem complexMulGauss_eval_honest [CommRing α] (a b c d : α) :
+    (complexMulGauss a b c d).eval ArithQuery.honest = (a * c - b * d, a * d + b * c) := by
+  simp [complexMulGauss, ArithQuery.doMul, ArithQuery.doSub, ArithQuery.doAdd, ArithQuery.honest]
+  ring
+
+-- ## Exact cost counts
+
+theorem complexMulNaive_cost (oracle : {ι : Type} → ArithQuery α ι → ι)
+    (c_add c_mul : Nat) (a b c d : α) :
+    (complexMulNaive a b c d).cost oracle (ArithQuery.weight c_add c_mul) =
+      4 * c_mul + 2 * c_add := by
+  simp [complexMulNaive, ArithQuery.doMul, ArithQuery.doSub, ArithQuery.doAdd, ArithQuery.weight]
+  omega
+
+theorem complexMulGauss_cost (oracle : {ι : Type} → ArithQuery α ι → ι)
+    (c_add c_mul : Nat) (a b c d : α) :
+    (complexMulGauss a b c d).cost oracle (ArithQuery.weight c_add c_mul) =
+      3 * c_mul + 5 * c_add := by
+  simp [complexMulGauss, ArithQuery.doMul, ArithQuery.doSub, ArithQuery.doAdd, ArithQuery.weight]
+  omega
+
+-- ## Crossover: Gauss beats naive when multiplication costs more than 3× addition
+
+theorem gauss_le_naive (c_add c_mul : Nat) (h : 3 * c_add ≤ c_mul) :
+    3 * c_mul + 5 * c_add ≤ 4 * c_mul + 2 * c_add := by omega
+
+theorem gauss_le_naive_iff (c_add c_mul : Nat) :
+    3 * c_mul + 5 * c_add ≤ 4 * c_mul + 2 * c_add ↔ 3 * c_add ≤ c_mul := by omega
+
+end Cslib.Query
+
+end -- public section

Cslib/Algorithms/Lean/Query/Bounds.lean

@@ -0,0 +1,37 @@
+/-
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf, Kim Morrison, Shreyas Srinivas
+-/
+module
+
+public import Cslib.Algorithms.Lean.Query.Prog
+
+/-! # Upper and Lower Bounds for Query Complexity
+
+Definitions of upper and lower bounds on the number of queries a program makes,
+quantified over oracles.
+-/
+
+open Cslib.Query
+
+public section
+
+namespace Cslib.Query
+
+/-- Upper bound: for all oracles, inputs of size ≤ n make at most `bound n` queries. -/
+@[expose] def UpperBound (prog : α → Prog Q β)
+    (size : α → Nat) (bound : Nat → Nat) : Prop :=
+  ∀ (oracle : {ι : Type} → Q ι → ι) (n : Nat) (x : α),
+    size x ≤ n → (prog x).queriesOn oracle ≤ bound n
+
+/-- Lower bound: for every size n, there exists an input and oracle
+    making the program perform ≥ `bound n` queries. -/
+@[expose] def LowerBound (prog : α → Prog Q β)
+    (size : α → Nat) (bound : Nat → Nat) : Prop :=
+  ∀ (n : Nat), ∃ (x : α), size x ≤ n ∧
+    ∃ (oracle : {ι : Type} → Q ι → ι), bound n ≤ (prog x).queriesOn oracle
+
+end Cslib.Query
+
+end -- public section

Cslib/Algorithms/Lean/Query/Prog.lean

@@ -0,0 +1,125 @@
+/-
+Copyright (c) 2026 Lean FRO, LLC. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sebastian Graf, Kim Morrison, Shreyas Srinivas
+-/
+module
+
+public import Cslib.Foundations.Control.Monad.Free
+
+/-! # Prog: Programs as Free Monads over Query Types
+
+`Prog Q α` is an alias for `FreeM Q α`, representing a program that makes queries of type `Q`
+and returns a result of type `α`. A query type `Q : Type → Type` maps each query to its
+response type.
+
+The key operations are:
+- `Prog.eval oracle p`: evaluate `p` by answering each query using `oracle`
+- `Prog.queriesOn oracle p`: count the queries along the oracle-determined path
+
+Because the oracle is supplied *after* the program produces its query plan (the `Prog` tree),
+a sound implementation of `prog` has no way to "guess" what the oracle would respond.
+This is the foundation of the anti-cheating guarantee for both upper and lower bounds.
+
+This provides an alternative to the `TimeM`-based cost analysis in
+`Cslib.Algorithms.Lean.MergeSort`: here query counting is structural (derived from the
+`Prog` tree) rather than annotation-based.
+-/
+
+open Cslib
+
+public section
+
+namespace Cslib.Query
+
+/-- A program that makes queries of type `Q` and returns a result of type `α`.
+    This is `FreeM Q α`, the free monad over the query type. -/
+abbrev Prog (Q : Type → Type) (α : Type) := FreeM Q α
+
+namespace Prog
+
+variable {Q : Type → Type} {α β : Type}
+
+/-- Evaluate a program by answering each query using `oracle`. -/
+@[expose] def eval (oracle : {ι : Type} → Q ι → ι) : Prog Q α → α
+  | .pure a => a
+  | .liftBind op cont => eval oracle (cont (oracle op))
+
+/-- Count the number of queries along the path determined by `oracle`. -/
+@[expose] def queriesOn (oracle : {ι : Type} → Q ι → ι) : Prog Q α → Nat
+  | .pure _ => 0
+  | .liftBind op cont => 1 + queriesOn oracle (cont (oracle op))
+
+-- Simp lemmas for eval
+
+@[simp] theorem eval_pure (oracle : {ι : Type} → Q ι → ι) (a : α) :
+    eval oracle (.pure a : Prog Q α) = a := rfl
+
+@[simp] theorem eval_liftBind (oracle : {ι : Type} → Q ι → ι)
+    {ι : Type} (op : Q ι) (cont : ι → Prog Q α) :
+    eval oracle (.liftBind op cont) = eval oracle (cont (oracle op)) := rfl
+
+@[simp] theorem eval_bind (oracle : {ι : Type} → Q ι → ι)
+    (t : Prog Q α) (f : α → Prog Q β) :
+    eval oracle (t.bind f) = eval oracle (f (eval oracle t)) := by
+  induction t with
+  | pure a => rfl
+  | liftBind op cont ih => exact ih (oracle op)
+
+-- Simp lemmas for queriesOn
+
+@[simp] theorem queriesOn_pure (oracle : {ι : Type} → Q ι → ι) (a : α) :
+    queriesOn oracle (.pure a : Prog Q α) = 0 := rfl
+
+@[simp] theorem queriesOn_liftBind (oracle : {ι : Type} → Q ι → ι)
+    {ι : Type} (op : Q ι) (cont : ι → Prog Q α) :
+    queriesOn oracle (.liftBind op cont) = 1 + queriesOn oracle (cont (oracle op)) := rfl
+
+@[simp] theorem queriesOn_bind (oracle : {ι : Type} → Q ι → ι)
+    (t : Prog Q α) (f : α → Prog Q β) :
+    queriesOn oracle (t.bind f) =
+      queriesOn oracle t + queriesOn oracle (f (eval oracle t)) := by
+  induction t with
+  | pure a => simp [FreeM.bind]
+  | liftBind op cont ih =>
+    simp only [FreeM.bind, queriesOn_liftBind, eval_liftBind, ih (oracle op)]
+    omega
+
+/-- Weighted query cost: each query has a cost given by `weight`. -/
+@[expose] def cost (oracle : {ι : Type} → Q ι → ι)
+    (weight : {ι : Type} → Q ι → Nat) : Prog Q α → Nat
+  | .pure _ => 0
+  | .liftBind op cont => weight op + cost oracle weight (cont (oracle op))
+
+-- Simp lemmas for cost
+
+@[simp] theorem cost_pure (oracle : {ι : Type} → Q ι → ι)
+    (weight : {ι : Type} → Q ι → Nat) (a : α) :
+    cost oracle weight (.pure a : Prog Q α) = 0 := rfl
+
+@[simp] theorem cost_liftBind (oracle : {ι : Type} → Q ι → ι)
+    (weight : {ι : Type} → Q ι → Nat) {ι : Type} (op : Q ι) (cont : ι → Prog Q α) :
+    cost oracle weight (.liftBind op cont) =
+      weight op + cost oracle weight (cont (oracle op)) := rfl
+
+@[simp] theorem cost_bind (oracle : {ι : Type} → Q ι → ι)
+    (weight : {ι : Type} → Q ι → Nat) (t : Prog Q α) (f : α → Prog Q β) :
+    cost oracle weight (t.bind f) =
+      cost oracle weight t + cost oracle weight (f (eval oracle t)) := by
+  induction t with
+  | pure a => simp [FreeM.bind]
+  | liftBind op cont ih =>
+    simp only [FreeM.bind, cost_liftBind, eval_liftBind, ih (oracle op)]
+    omega
+
+theorem queriesOn_eq_cost_one (oracle : {ι : Type} → Q ι → ι) (p : Prog Q α) :
+    queriesOn oracle p = cost oracle (fun _ => 1) p := by
+  induction p with
+  | pure a => rfl
+  | liftBind op cont ih => simp [ih (oracle op)]
+
+end Prog
+
+end Cslib.Query
+
+end -- public section