feat: query complexity via free monads

Cslib.lean

@@ -1,3 +1,5 @@
+import Cslib.Algorithms.MergeSort.MergeSort
+import Cslib.Algorithms.QueryModel
 import Cslib.Computability.Automata.Acceptors.Acceptor
 import Cslib.Computability.Automata.Acceptors.OmegaAcceptor
 import Cslib.Computability.Automata.DA.Basic
@@ -21,6 +23,7 @@ import Cslib.Computability.Languages.RegularLanguage
 import Cslib.Foundations.Control.Monad.Free
 import Cslib.Foundations.Control.Monad.Free.Effects
 import Cslib.Foundations.Control.Monad.Free.Fold
+import Cslib.Foundations.Control.Monad.Time
 import Cslib.Foundations.Data.FinFun
 import Cslib.Foundations.Data.HasFresh
 import Cslib.Foundations.Data.Nat.Segment

Cslib/Algorithms/MergeSort/MergeSort.lean

@@ -0,0 +1,95 @@
+/-
+Copyright (c) 2025 Tanner Duve. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Tanner Duve
+-/
+
+import Cslib.Algorithms.QueryModel
+
+/-!
+# Merge sort in the query model
+
+This file implements merge sort as a program in the query model defined in
+`Cslib.Algorithms.QueryModel`. The algorithm uses only comparison queries.
+
+## Main definitions
+
+- `merge`     : merge step using `Prog` comparisons
+- `split`     : split a list in two by alternating elements
+- `mergeSort` : merge sort expressed in the query model
+
+We also provide simple example evaluations of `mergeSort` and its time cost.
+-/
+
+open Cslib
+
+namespace Cslib.Algorithms.MergeSort.QueryBased
+
+open Cslib.Algorithms
+
+
+/-- The Model for comparison sorting natural-number registers.-/
+inductive QueryF : Type → Type where
+  /-- Read the value stored at index `i`. -/
+  | read  : Nat → QueryF Nat
+  /-- Write value `v` at index `i`. -/
+  | write : Nat → Nat → QueryF PUnit
+  /-- Compare the values at indices `i` and `j`. -/
+  | cmp   : Nat → Nat → QueryF Bool
+
+/-- Lift a comparison into the query language at the top level. -/
+def cmpVal (x y : Nat) : Prog Bool :=
+  FreeM.lift (QueryF.cmp x y)
+
+/-- Concrete semantics for primitive queries, used to run programs. -/
+def evalQuery : {ι : Type} → QueryF ι → ι
+  | _, .read _      => 0
+  | _, .write _ _   => PUnit.unit
+  | _, .cmp x y     => x ≤ y
+
+
+
+
+abbrev SortProg := Prog QueryF
+
+/-- Lift a comparison on values into the free monad. -/
+def cmpVal (x y : Nat) : SortProg Bool :=
+  FreeM.lift (QueryF.cmp x y)
+
+/-- Merge two sorted lists using comparisons in the query monad. -/
+def merge : List Nat → List Nat → SortProg (List Nat)
+  | [], ys => pure ys
+  | xs, [] => pure xs
+  | x :: xs', y :: ys' => do
+      let b ← cmpVal x y
+      if b then
+        let rest ← merge xs' (y :: ys')
+        pure (x :: rest)
+      else
+        let rest ← merge (x :: xs') ys'
+        pure (y :: rest)
+
+/-- Split a list into two lists by alternating elements. -/
+def split (xs : List Nat) : List Nat × List Nat :=
+  let rec go : List Nat → List Nat → List Nat → List Nat × List Nat
+    | [], accL, accR => (accL.reverse, accR.reverse)
+    | [x], accL, accR => ((x :: accL).reverse, accR.reverse)
+    | x :: y :: xs, accL, accR => go xs (x :: accL) (y :: accR)
+  go xs [] []
+
+/-- Merge sort expressed as a program in the query model.
+TODO: Working version without partial -/
+partial def mergeSort : List Nat → SortProg (List Nat)
+  | []      => pure []
+  | [x]     => pure [x]
+  | xs      =>
+    let (left, right) := split xs
+    do
+      let sortedLeft  ← mergeSort left
+      let sortedRight ← mergeSort right
+      merge sortedLeft sortedRight
+
+#eval evalProg (mergeSort [5,3,8,6,2,7,4,1])
+#eval timeProg (mergeSort [5,3,8,6,2,7,4,1])
+
+end Cslib.Algorithms.MergeSort.QueryBased

Cslib/Algorithms/QueryModel.lean

@@ -0,0 +1,110 @@
+/-
+Copyright (c) 2025 Tanner Duve. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Tanner Duve
+-/
+
+import Cslib.Foundations.Control.Monad.Free.Effects
+import Cslib.Foundations.Control.Monad.Free.Fold
+import Cslib.Foundations.Control.Monad.Time
+
+/-
+# Query model
+
+This file defines a simple query language modeled as a free monad over a
+parametric type of  query operations.
+
+We equip this language with a cost model (`TimeM`) that counts how many primitive queries
+are performed. An example algorithm (merge sort) is implemented in
+`Cslib.Algorithms.MergeSort.QueryBased`.
+
+## Main definitions
+
+- `QueryF`, `Prog` : query language and programs
+- `timeOfQuery`, `timeInterp`, `timeProg` : cost model for programs
+- `evalQuery`, `evalProg` : concrete execution semantics
+
+## Tags
+
+query model, free monad, time complexity, merge sort
+-/
+
+namespace Cslib
+
+namespace Algorithms
+
+
+
+/-- Programs built as the free monad over `QueryF`. -/
+abbrev Prog (QType : Type u → Type u) (α : Type v) := FreeM QType α
+
+
+instance {QType : Type u → Type u} : Monad (Prog QType) := inferInstance
+
+
+
+
+namespace Prog
+
+
+
+/-- Conditional branching on a boolean program. -/
+def cond {QType} {α} (b : Prog QType Bool) (t e : Prog QType α) : Prog QType α :=
+  b.bind (fun b' => if b' then t else e)
+
+/-- A counting loop from `0` to `n - 1`, sequencing the body. -/
+def forLoop {QType} (n : Nat) (body : Nat → Prog QType PUnit) : Prog QType PUnit :=
+  go n
+where
+  /-- Auxiliary recursive worker for `forLoop`. -/
+  go : Nat → Prog QType PUnit
+    | 0       => pure ()
+    | i + 1   =>
+      body i >>= fun _ => go i
+
+end Prog
+
+class Query (Q : Type u → Type u) where
+  timeOfQuery : {ι : Type u} → Q ι → Nat
+  evalQuery : {ι : Type u} → Q ι → ι
+
+open Query
+-- /-- Constant time cost assigned to each primitive query. -/
+-- def timeOfQuery : {ι : Type} → QueryF ι → Nat
+--   | _, .read _       => 1
+--   | _, .write _ _    => 1
+--   | _, .cmp _ _      => 1
+
+/--
+Interpret primitive queries into the time-counting monad `TimeM`.
+-/
+def timeInterp [Query QF] [Inhabited ι] (q : QF ι) : TimeM ι :=
+  TimeM.tick default (timeOfQuery q)
+
+-- /-- Interpret primitive queries into the time-counting monad `TimeM`. -/
+-- def timeInterp : {ι : Type} → QueryF ι → TimeM ι
+--   | _, .read i      => TimeM.tick 0 (timeOfQuery (.read i))
+--   | _, .write i v   => TimeM.tick PUnit.unit (timeOfQuery (.write i v))
+--   | _, .cmp i j     => TimeM.tick false (timeOfQuery (.cmp i j))
+
+
+instance {α : Type u} {Q : Type u → Type u} [Query Q] : MonadLiftT (Prog Q) TimeM where
+  monadLift (p : Prog Q α) : TimeM α :=
+    timeInterp
+
+/-- Total time cost of running a program under the interpreter `timeInterp`. -/
+def timeProg [Query QF] {α : Type u}  (p : Prog QF α) : Nat :=
+  (p.liftM timeInterp).time
+
+
+
+/-- Evaluate a query program to a pure value using `evalQuery`. -/
+def evalProg [Query QF] {α : Type} (p : Prog QF α) : α :=
+  FreeM.foldFreeM id
+    (fun {ι} (op : QF ι) (k : ι → α) =>
+      k (evalQuery op))
+    p
+
+end Algorithms
+
+end Cslib

Cslib/Foundations/Control/Monad/Free/Effects.lean

@@ -4,7 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Tanner Duve
 -/
 import Cslib.Foundations.Control.Monad.Free
+import Mathlib.Control.Monad.Writer
+import Cslib.Foundations.Control.Monad.Time
 import Mathlib.Control.Monad.Cont
+import Mathlib.Data.Nat.Basic
 
 /-!
 # Free Monad
@@ -13,7 +16,7 @@ This file defines several canonical instances on the free monad.
 
 ## Main definitions
 
-- `FreeState`, `FreeWriter`, `FreeCont`: Specific effect monads
+- `FreeState`, `FreeWriter`, `FreeTime`, `FreeCont`: Specific effect monads
 
 ## Implementation
 
@@ -29,7 +32,7 @@ the universal property.
 
 ## Tags
 
-Free monad, state monad, writer monad, continuation monad
+Free monad, state monad, writer monad, time monad, continuation monad
 -/
 
 namespace Cslib
@@ -262,6 +265,44 @@ instance [Monoid ω] : MonadWriter ω (FreeWriter ω) where
 
 end FreeWriter
 
+/-! ### Time Monad via `FreeM` -/
+
+/-- Time monad implemented as the free writer monad over `Nat`. This models computations that
+emit natural-number costs while producing a result. -/
+abbrev FreeTime := FreeWriter Nat
+
+namespace FreeTime
+
+variable {α : Type}
+
+/-- Emit a time cost of `c` units (default `1`). -/
+def tick (c : Nat := 1) : FreeTime PUnit :=
+  FreeWriter.tell c
+
+/-- Run a `FreeTime` computation, returning the result and total time cost.
+
+The cost is accumulated additively starting from `0`. -/
+def run : FreeTime α → α × Nat
+  | .pure a => (a, 0)
+  | .liftBind (.tell c) k =>
+      let (a, t) := run (k .unit)
+      (a, c + t)
+
+/-- Interpret a `FreeTime` computation into the concrete time monad `TimeM`. -/
+def toTimeM (comp : FreeTime α) : TimeM α :=
+  let (a, t) := run comp
+  ⟨a, t⟩
+
+@[simp]
+lemma run_pure (a : α) :
+    run (.pure a : FreeTime α) = (a, 0) := rfl
+
+@[simp]
+lemma tick_def (c : Nat) :
+    tick c = (FreeWriter.tell c : FreeTime PUnit) := rfl
+
+end FreeTime
+
 /-! ### Continuation Monad via `FreeM` -/
 
 /-- Type constructor for continuation operations. -/

Cslib/Foundations/Control/Monad/Time.lean

@@ -0,0 +1,70 @@
+/-
+Copyright (c) 2025 Sorrachai Yingchareonthawornhcai. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Sorrachai Yingchareonthawornhcai, Tanner Duve
+-/
+
+import Mathlib.Control.Monad.Writer
+
+/-!
+# Time Monad
+
+`TimeM` is a monad that tracks execution time alongside computations, using natural numbers
+as a simple cost model. As plain types it is isomorphic to `WriterT Nat Id`.
+-/
+
+structure TimeM (α : Type u) where
+  /-- The result of the computation. -/
+  ret : α
+  /-- The accumulated time cost. -/
+  time : Nat
+
+namespace TimeM
+
+def pure {α} (a : α) : TimeM α :=
+  ⟨a, 0⟩
+
+def bind {α β} (m : TimeM α) (f : α → TimeM β) : TimeM β :=
+  let r := f m.ret
+  ⟨r.ret, m.time + r.time⟩
+
+instance : Monad TimeM where
+  pure := pure
+  bind := bind
+
+@[simp] def tick {α : Type} (a : α) (c : ℕ := 1) : TimeM α :=
+  ⟨a, c⟩
+
+scoped notation "✓" a:arg ", " c:arg => tick a c
+scoped notation "✓" a:arg => tick a  -- Default case with only one argument
+
+def tickUnit : TimeM Unit :=
+  ✓ () -- This uses the default time increment of 1
+
+@[simp] theorem time_of_pure {α} (a : α) : (pure a).time = 0 := rfl
+@[simp] theorem time_of_bind {α β} (m : TimeM α) (f : α → TimeM β) :
+ (TimeM.bind m f).time = m.time + (f m.ret).time := rfl
+@[simp] theorem time_of_tick {α} (a : α) (c : ℕ) : (tick a c).time = c := rfl
+@[simp] theorem ret_bind {α β} (m : TimeM α) (f : α → TimeM β) :
+  (TimeM.bind m f).ret = (f m.ret).ret := rfl
+
+-- this allow us to simplify the chain of compositions
+attribute [simp] Bind.bind Pure.pure TimeM.pure
+
+/-- `TimeM` is (definitionally) the same as the writer monad `WriterT Nat Id`. -/
+abbrev WriterNat (α : Type) := WriterT Nat Id α
+
+/-- Equivalence between `TimeM α` and `WriterT Nat Id α` as plain types. -/
+def equivWriter (α : Type) : TimeM α ≃ WriterNat α where
+  toFun m := (m.ret, m.time)
+  invFun w := ⟨w.1, w.2⟩
+  left_inv m := by cases m; rfl
+  right_inv w := by cases w; rfl
+
+@[simp] lemma equivWriter_toFun {α : Type} (m : TimeM α) :
+    (equivWriter α m : WriterNat α) = (m.ret, m.time) := rfl
+
+@[simp] lemma equivWriter_invFun {α : Type} (w : WriterNat α) :
+    (equivWriter α).invFun w = TimeM.mk w.1 w.2 := rfl
+
+end TimeM