From 7c47d30bab2ea45952a1e58ddf65f2e0ae422c64 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:11:56 -0800
Subject: [PATCH 01/13] initial commit

---
 Cslib/Algorithms/MergeSort/MergeSort.lean     | 66 +++++++++++++
 Cslib/Algorithms/QueryModel.lean              | 98 +++++++++++++++++++
 .../Control/Monad/Free/Effects.lean           | 38 ++++++-
 Cslib/Foundations/Control/Monad/Time.lean     | 82 ++++++++++++++++
 4 files changed, 282 insertions(+), 2 deletions(-)
 create mode 100644 Cslib/Algorithms/MergeSort/MergeSort.lean
 create mode 100644 Cslib/Algorithms/QueryModel.lean
 create mode 100644 Cslib/Foundations/Control/Monad/Time.lean

diff --git a/Cslib/Algorithms/MergeSort/MergeSort.lean b/Cslib/Algorithms/MergeSort/MergeSort.lean
new file mode 100644
index 00000000..cee93ffa
--- /dev/null
+++ b/Cslib/Algorithms/MergeSort/MergeSort.lean
@@ -0,0 +1,66 @@
+/-
+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
+
+/-- Merge two sorted lists using comparisons in the query monad. -/
+def merge : List Nat → List Nat → Prog (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. -/
+partial def mergeSort : List Nat → Prog (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
+
diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
new file mode 100644
index 00000000..d89c55d5
--- /dev/null
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -0,0 +1,98 @@
+/-
+Copyright (c) 2025 Tanner Duve. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Tanner Duve, Sorrachai Yingchareonthawornhcai
+-/
+
+import Cslib.Foundations.Control.Monad.Free.Effects
+import Cslib.Foundations.Control.Monad.Free.Fold
+import Cslib.Foundations.Control.Monad.Time
+
+/-
+# Query model for comparison-based algorithms
+
+This file defines a simple query language for reading, writing, and comparing natural numbers,
+modeled as a free monad over primitive 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`   : functor of primitive query operations
+- `Prog`     : free monad of query programs
+- `timeOfQuery`, `timeInterp`, `timeProg` : cost model for programs
+- `evalProg` : concrete execution semantics for programs
+
+## Tags
+
+query model, free monad, time complexity, merge sort
+-/
+
+open Cslib
+
+/-- Primitive queries on 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
+
+/-- Programs built as the free monad over `QueryF`. -/
+abbrev Prog (α : Type) := FreeM QueryF α
+
+namespace Prog
+
+/-- Lift a comparison on values into the free monad. -/
+def cmpVal (x y : Nat) : Prog Bool :=
+  FreeM.lift (QueryF.cmp x y)
+
+/-- Conditional branching on a boolean program. -/
+def cond {α} (b : Prog Bool) (t e : Prog α) : Prog α :=
+  b.bind (fun b' => if b' then t else e)
+
+/-- A counting loop from `0` to `n - 1`, sequencing the body. -/
+def forLoop (n : Nat) (body : Nat → Prog PUnit) : Prog PUnit :=
+  let rec go : Nat → Prog PUnit
+    | 0       => pure ()
+    | i + 1   =>
+      body i >>= fun _ => go i
+  go n
+
+end Prog
+
+/-- 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 : {ι : 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))
+
+/-- Total time cost of running a program under the interpreter `timeInterp`. -/
+def timeProg {α : Type} (p : Prog α) : Nat :=
+  (p.liftM timeInterp).time
+
+/-- 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
+
+/-- Evaluate a query program to a pure value using `evalQuery`. -/
+def evalProg {α : Type} (p : Prog α) : α :=
+  FreeM.foldFreeM id
+    (fun {ι} (op : QueryF ι) (k : ι → α) =>
+      k (evalQuery op))
+    p
diff --git a/Cslib/Foundations/Control/Monad/Free/Effects.lean b/Cslib/Foundations/Control/Monad/Free/Effects.lean
index 6e81e050..39ae3449 100644
--- a/Cslib/Foundations/Control/Monad/Free/Effects.lean
+++ b/Cslib/Foundations/Control/Monad/Free/Effects.lean
@@ -6,6 +6,7 @@ Authors: Tanner Duve
 import Mathlib.Control.Monad.Cont
 import Cslib.Foundations.Control.Monad.Free
 import Mathlib.Control.Monad.Writer
+import Cslib.Foundations.Control.Monad.Time
 
 /-!
 # Free Monad
@@ -14,7 +15,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
 
@@ -30,7 +31,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
@@ -263,6 +264,39 @@ 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 u}
+
+/-- 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. -/
+def run (comp : FreeTime α) : α × Nat :=
+  FreeWriter.run (ω := Nat) comp
+
+/-- 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, 1) := 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. -/
diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
new file mode 100644
index 00000000..ca2edb1b
--- /dev/null
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -0,0 +1,82 @@
+/-
+Copyright (c) 2025 Tanner Duve. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Tanner Duve
+-/
+
+import Mathlib.Control.Monad.Writer
+
+/-!
+# Time monad
+
+This file defines a simple monad `TimeM` that pairs a value with a natural number
+representing an accumulated time/cost. As plain types it is isomorphic to `Writer Nat`.
+
+## Main definitions
+
+- `TimeM`          : computations with a time component
+- `TimeM.equivWriter` : equivalence with `Writer Nat`
+-/
+
+namespace Cslib
+
+universe u
+
+/-- A computation that returns a value of type `α` together with an accumulated
+time cost (a natural number). -/
+structure TimeM (α : Type u) where
+  /-- The result of the computation. -/
+  val : α
+  /-- The accumulated time cost. -/
+  time : Nat
+
+namespace TimeM
+
+variable {α β : Type u}
+
+/-- Return a value with zero time cost. -/
+def pure (a : α) : TimeM α :=
+  ⟨a, 0⟩
+
+/-- Sequence two computations, adding their time components. -/
+def bind (m : TimeM α) (f : α → TimeM β) : TimeM β :=
+  let r := f m.val
+  ⟨r.val, m.time + r.time⟩
+
+instance : Monad TimeM where
+  pure := pure
+  bind := bind
+
+/-- Construct a value that costs `c` units of time. -/
+def tick (a : α) (c : Nat := 1) : TimeM α :=
+  ⟨a, c⟩
+
+@[simp] theorem time_of_pure (a : α) : (pure a).time = 0 := rfl
+
+@[simp] theorem time_of_bind (m : TimeM α) (f : α → TimeM β) :
+    (bind m f).time = m.time + (f m.val).time := rfl
+
+@[simp] theorem time_of_tick (a : α) (c : Nat) : (tick a c).time = c := rfl
+
+@[simp] theorem val_bind (m : TimeM α) (f : α → TimeM β) :
+    (bind m f).val = (f m.val).val := rfl
+
+/-- `TimeM` is (definitionally) the same as the writer monad `Writer Nat`. -/
+abbrev WriterNat (α : Type) := WriterT Nat Id α
+
+/-- Equivalence between `TimeM α` and `Writer Nat α` as plain types. -/
+def equivWriter (α : Type) : TimeM α ≃ WriterNat α where
+  toFun m := (m.val, 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.val, m.time) := rfl
+
+@[simp] lemma equivWriter_invFun {α : Type} (w : WriterNat α) :
+    (equivWriter α).invFun w = TimeM.mk w.1 w.2 := rfl
+
+end TimeM
+
+end Cslib

From c0c342660586acd7f66be65b01a416e2130303f1 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:15:31 -0800
Subject: [PATCH 02/13] fix up docs

---
 Cslib/Algorithms/MergeSort/MergeSort.lean |  4 ++--
 Cslib/Algorithms/QueryModel.lean          | 12 +++++-------
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/Cslib/Algorithms/MergeSort/MergeSort.lean b/Cslib/Algorithms/MergeSort/MergeSort.lean
index cee93ffa..9dd84051 100644
--- a/Cslib/Algorithms/MergeSort/MergeSort.lean
+++ b/Cslib/Algorithms/MergeSort/MergeSort.lean
@@ -48,7 +48,8 @@ def split (xs : List Nat) : List Nat × List Nat :=
     | x :: y :: xs, accL, accR => go xs (x :: accL) (y :: accR)
   go xs [] []
 
-/-- Merge sort expressed as a program in the query model. -/
+/-- Merge sort expressed as a program in the query model.
+TODO: Working version without partial -/
 partial def mergeSort : List Nat → Prog (List Nat)
   | []      => pure []
   | [x]     => pure [x]
@@ -63,4 +64,3 @@ partial def mergeSort : List Nat → Prog (List Nat)
 #eval timeProg (mergeSort [5,3,8,6,2,7,4,1])
 
 end Cslib.Algorithms.MergeSort.QueryBased
-
diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index d89c55d5..c8ef4433 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -1,7 +1,7 @@
 /-
 Copyright (c) 2025 Tanner Duve. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Tanner Duve, Sorrachai Yingchareonthawornhcai
+Authors: Tanner Duve
 -/
 
 import Cslib.Foundations.Control.Monad.Free.Effects
@@ -9,10 +9,9 @@ import Cslib.Foundations.Control.Monad.Free.Fold
 import Cslib.Foundations.Control.Monad.Time
 
 /-
-# Query model for comparison-based algorithms
+# Query model
 
-This file defines a simple query language for reading, writing, and comparing natural numbers,
-modeled as a free monad over primitive query operations.
+This file defines a simple query language modeled as a free monad over primitive 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
@@ -20,10 +19,9 @@ are performed. An example algorithm (merge sort) is implemented in
 
 ## Main definitions
 
-- `QueryF`   : functor of primitive query operations
-- `Prog`     : free monad of query programs
+- `QueryF`, `Prog` : query language and programs
 - `timeOfQuery`, `timeInterp`, `timeProg` : cost model for programs
-- `evalProg` : concrete execution semantics for programs
+- `evalQuery`, `evalProg` : concrete execution semantics
 
 ## Tags
 

From a5389ab2b3ab428a54d889be8caa3c11dc0ef152 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:16:06 -0800
Subject: [PATCH 03/13] docs

---
 Cslib/Foundations/Control/Monad/Time.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
index ca2edb1b..d9d36995 100644
--- a/Cslib/Foundations/Control/Monad/Time.lean
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -1,7 +1,7 @@
 /-
 Copyright (c) 2025 Tanner Duve. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Tanner Duve
+Authors: Tanner Duve, Sorrachai Yingchareonthawornhcai
 -/
 
 import Mathlib.Control.Monad.Writer

From a32b00d3f9c5a93630b9f20dca34cbecd4dd95bd Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:39:28 -0800
Subject: [PATCH 04/13] fix up

---
 Cslib/Algorithms/QueryModel.lean              | 14 +++-
 .../Control/Monad/Free/Effects.lean           |  2 +-
 Cslib/Foundations/Control/Monad/Time.lean     | 69 ++++++++-----------
 3 files changed, 42 insertions(+), 43 deletions(-)

diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index c8ef4433..e242916b 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -28,7 +28,9 @@ are performed. An example algorithm (merge sort) is implemented in
 query model, free monad, time complexity, merge sort
 -/
 
-open Cslib
+namespace Cslib
+
+namespace Algorithms
 
 /-- Primitive queries on natural-number registers. -/
 inductive QueryF : Type → Type where
@@ -54,11 +56,13 @@ def cond {α} (b : Prog Bool) (t e : Prog α) : Prog α :=
 
 /-- A counting loop from `0` to `n - 1`, sequencing the body. -/
 def forLoop (n : Nat) (body : Nat → Prog PUnit) : Prog PUnit :=
-  let rec go : Nat → Prog PUnit
+  go n
+where
+  /-- Auxiliary recursive worker for `forLoop`. -/
+  go : Nat → Prog PUnit
     | 0       => pure ()
     | i + 1   =>
       body i >>= fun _ => go i
-  go n
 
 end Prog
 
@@ -94,3 +98,7 @@ def evalProg {α : Type} (p : Prog α) : α :=
     (fun {ι} (op : QueryF ι) (k : ι → α) =>
       k (evalQuery op))
     p
+
+end Algorithms
+
+end Cslib
diff --git a/Cslib/Foundations/Control/Monad/Free/Effects.lean b/Cslib/Foundations/Control/Monad/Free/Effects.lean
index 39ae3449..ea56e221 100644
--- a/Cslib/Foundations/Control/Monad/Free/Effects.lean
+++ b/Cslib/Foundations/Control/Monad/Free/Effects.lean
@@ -272,7 +272,7 @@ abbrev FreeTime := FreeWriter Nat
 
 namespace FreeTime
 
-variable {α : Type u}
+variable {α : Type}
 
 /-- Emit a time cost of `c` units (default `1`). -/
 def tick (c : Nat := 1) : FreeTime PUnit :=
diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
index d9d36995..1fe8e0c4 100644
--- a/Cslib/Foundations/Control/Monad/Time.lean
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -1,82 +1,73 @@
 /-
-Copyright (c) 2025 Tanner Duve. All rights reserved.
+Copyright (c) 2025 Sorrachai Yingchareonthawornhcai. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Tanner Duve, Sorrachai Yingchareonthawornhcai
+Authors: Sorrachai Yingchareonthawornhcai, Tanner Duve
 -/
 
 import Mathlib.Control.Monad.Writer
 
 /-!
-# Time monad
+# Time Monad
 
-This file defines a simple monad `TimeM` that pairs a value with a natural number
-representing an accumulated time/cost. As plain types it is isomorphic to `Writer Nat`.
-
-## Main definitions
-
-- `TimeM`          : computations with a time component
-- `TimeM.equivWriter` : equivalence with `Writer Nat`
+`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`.
 -/
 
-namespace Cslib
-
-universe u
+set_option tactic.hygienic false
+set_option autoImplicit false
 
-/-- A computation that returns a value of type `α` together with an accumulated
-time cost (a natural number). -/
-structure TimeM (α : Type u) where
+structure TimeM (α : Type) where
   /-- The result of the computation. -/
-  val : α
+  ret : α
   /-- The accumulated time cost. -/
   time : Nat
 
 namespace TimeM
 
-variable {α β : Type u}
-
-/-- Return a value with zero time cost. -/
-def pure (a : α) : TimeM α :=
+def pure {α} (a : α) : TimeM α :=
   ⟨a, 0⟩
 
-/-- Sequence two computations, adding their time components. -/
-def bind (m : TimeM α) (f : α → TimeM β) : TimeM β :=
-  let r := f m.val
-  ⟨r.val, m.time + r.time⟩
+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
 
-/-- Construct a value that costs `c` units of time. -/
-def tick (a : α) (c : Nat := 1) : TimeM α :=
+@[simp] def tick {α : Type} (a : α) (c : ℕ := 1) : TimeM α :=
   ⟨a, c⟩
 
-@[simp] theorem time_of_pure (a : α) : (pure a).time = 0 := rfl
+scoped notation "✓" a:arg ", " c:arg => tick a c
+scoped notation "✓" a:arg => tick a  -- Default case with only one argument
 
-@[simp] theorem time_of_bind (m : TimeM α) (f : α → TimeM β) :
-    (bind m f).time = m.time + (f m.val).time := rfl
+def tickUnit : TimeM Unit :=
+  ✓ () -- This uses the default time increment of 1
 
-@[simp] theorem time_of_tick (a : α) (c : Nat) : (tick a c).time = c := rfl
+@[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
 
-@[simp] theorem val_bind (m : TimeM α) (f : α → TimeM β) :
-    (bind m f).val = (f m.val).val := 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 `Writer Nat`. -/
+/-- `TimeM` is (definitionally) the same as the writer monad `WriterT Nat Id`. -/
 abbrev WriterNat (α : Type) := WriterT Nat Id α
 
-/-- Equivalence between `TimeM α` and `Writer Nat α` as plain types. -/
+/-- Equivalence between `TimeM α` and `WriterT Nat Id α` as plain types. -/
 def equivWriter (α : Type) : TimeM α ≃ WriterNat α where
-  toFun m := (m.val, m.time)
+  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.val, m.time) := rfl
+    (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
-
-end Cslib

From c8d3206965a25a076e2d2725206a96337b55ebf9 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:49:27 -0800
Subject: [PATCH 05/13] effects fix

---
 Cslib/Foundations/Control/Monad/Free/Effects.lean | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/Cslib/Foundations/Control/Monad/Free/Effects.lean b/Cslib/Foundations/Control/Monad/Free/Effects.lean
index 9b6ea47b..63bda555 100644
--- a/Cslib/Foundations/Control/Monad/Free/Effects.lean
+++ b/Cslib/Foundations/Control/Monad/Free/Effects.lean
@@ -7,6 +7,7 @@ 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
@@ -278,9 +279,14 @@ variable {α : Type}
 def tick (c : Nat := 1) : FreeTime PUnit :=
   FreeWriter.tell c
 
-/-- Run a `FreeTime` computation, returning the result and total time cost. -/
-def run (comp : FreeTime α) : α × Nat :=
-  FreeWriter.run (ω := Nat) comp
+/-- 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 α :=
@@ -289,7 +295,7 @@ def toTimeM (comp : FreeTime α) : TimeM α :=
 
 @[simp]
 lemma run_pure (a : α) :
-    run (.pure a : FreeTime α) = (a, 1) := rfl
+    run (.pure a : FreeTime α) = (a, 0) := rfl
 
 @[simp]
 lemma tick_def (c : Nat) :

From 9c3c26cb73cd26762516d6721a696d720e429601 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:50:20 -0800
Subject: [PATCH 06/13] Add TimeM time/cost monad

Co-authored-by: Sorrachai Yingchareonthawornhcai <sorrachai@users.noreply.github.com>
---
 Cslib/Foundations/Control/Monad/Time.lean | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
index 1fe8e0c4..51fa7245 100644
--- a/Cslib/Foundations/Control/Monad/Time.lean
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -13,8 +13,8 @@ import Mathlib.Control.Monad.Writer
 as a simple cost model. As plain types it is isomorphic to `WriterT Nat Id`.
 -/
 
-set_option tactic.hygienic false
-set_option autoImplicit false
+
+
 
 structure TimeM (α : Type) where
   /-- The result of the computation. -/

From 2136afe1679aadc2f51acb0a94de3cc3ececb614 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:52:58 -0800
Subject: [PATCH 07/13] Add TimeM time/cost monad

Co-authored-by: Sorrachai Yingchareonthawornhcai <sorrachai@users.noreply.github.com>
---
 Cslib/Foundations/Control/Monad/Time.lean  |  3 -
 Cslib/Foundations/Control/Monad/TimeM.lean | 70 ++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 Cslib/Foundations/Control/Monad/TimeM.lean

diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
index 51fa7245..72ed077c 100644
--- a/Cslib/Foundations/Control/Monad/Time.lean
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -13,9 +13,6 @@ import Mathlib.Control.Monad.Writer
 as a simple cost model. As plain types it is isomorphic to `WriterT Nat Id`.
 -/
 
-
-
-
 structure TimeM (α : Type) where
   /-- The result of the computation. -/
   ret : α
diff --git a/Cslib/Foundations/Control/Monad/TimeM.lean b/Cslib/Foundations/Control/Monad/TimeM.lean
new file mode 100644
index 00000000..72ed077c
--- /dev/null
+++ b/Cslib/Foundations/Control/Monad/TimeM.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) 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

From f5397b060857292a5fd9a833a4558578d0846484 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:53:13 -0800
Subject: [PATCH 08/13] rename

---
 Cslib/Foundations/Control/Monad/TimeM.lean | 70 ----------------------
 1 file changed, 70 deletions(-)
 delete mode 100644 Cslib/Foundations/Control/Monad/TimeM.lean

diff --git a/Cslib/Foundations/Control/Monad/TimeM.lean b/Cslib/Foundations/Control/Monad/TimeM.lean
deleted file mode 100644
index 72ed077c..00000000
--- a/Cslib/Foundations/Control/Monad/TimeM.lean
+++ /dev/null
@@ -1,70 +0,0 @@
-/-
-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) 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

From 9e6fcd4114fe189947fea09d1939ab61ed7f5862 Mon Sep 17 00:00:00 2001
From: Tanner Duve <tannerduve@gmail.com>
Date: Fri, 5 Dec 2025 15:58:35 -0800
Subject: [PATCH 09/13] update cslib.lean

---
 Cslib.lean | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/Cslib.lean b/Cslib.lean
index 17891097..e7ae752e 100644
--- a/Cslib.lean
+++ b/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

From 9381b9aab11a148c73313efbd2bcf3c332c61c5b Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasss94@gmail.com>
Date: Sat, 6 Dec 2025 01:36:12 +0100
Subject: [PATCH 10/13] Test PR to PR ability with doc comment

---
 Cslib/Algorithms/QueryModel.lean | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index e242916b..ec0a5726 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -11,7 +11,8 @@ import Cslib.Foundations.Control.Monad.Time
 /-
 # Query model
 
-This file defines a simple query language modeled as a free monad over primitive query operations.
+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

From aeed9a74db5c974f18dcfe916fbef92f47b5a0e2 Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasss94@gmail.com>
Date: Sat, 6 Dec 2025 02:42:36 +0100
Subject: [PATCH 11/13] Yet to fix monad instance issue in timeProg function

---
 Cslib/Algorithms/MergeSort/MergeSort.lean | 33 ++++++++++-
 Cslib/Algorithms/QueryModel.lean          | 69 +++++++++++------------
 2 files changed, 64 insertions(+), 38 deletions(-)

diff --git a/Cslib/Algorithms/MergeSort/MergeSort.lean b/Cslib/Algorithms/MergeSort/MergeSort.lean
index 9dd84051..222619fb 100644
--- a/Cslib/Algorithms/MergeSort/MergeSort.lean
+++ b/Cslib/Algorithms/MergeSort/MergeSort.lean
@@ -27,8 +27,37 @@ 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 → Prog (List Nat)
+def merge : List Nat → List Nat → SortProg (List Nat)
   | [], ys => pure ys
   | xs, [] => pure xs
   | x :: xs', y :: ys' => do
@@ -50,7 +79,7 @@ def split (xs : List Nat) : List Nat × List Nat :=
 
 /-- Merge sort expressed as a program in the query model.
 TODO: Working version without partial -/
-partial def mergeSort : List Nat → Prog (List Nat)
+partial def mergeSort : List Nat → SortProg (List Nat)
   | []      => pure []
   | [x]     => pure [x]
   | xs      =>
diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index ec0a5726..4b604e93 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -33,70 +33,67 @@ namespace Cslib
 
 namespace Algorithms
 
-/-- Primitive queries on 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
+
 
 /-- Programs built as the free monad over `QueryF`. -/
-abbrev Prog (α : Type) := FreeM QueryF α
+abbrev Prog (QType : Type u → Type u) (α : Type v) := FreeM QType α
+
+instance {QType : Type u → Type u} : Bind (Prog QType) := inferInstance
 
+instance {QType : Type u → Type u} : Monad (Prog QType) := inferInstance
 namespace Prog
 
-/-- Lift a comparison on values into the free monad. -/
-def cmpVal (x y : Nat) : Prog Bool :=
-  FreeM.lift (QueryF.cmp x y)
+
 
 /-- Conditional branching on a boolean program. -/
-def cond {α} (b : Prog Bool) (t e : Prog α) : Prog α :=
+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 (n : Nat) (body : Nat → Prog PUnit) : Prog PUnit :=
+def forLoop {QType} (n : Nat) (body : Nat → Prog QType PUnit) : Prog QType PUnit :=
   go n
 where
   /-- Auxiliary recursive worker for `forLoop`. -/
-  go : Nat → Prog PUnit
+  go : Nat → Prog QType PUnit
     | 0       => pure ()
     | i + 1   =>
       body i >>= fun _ => go i
 
 end Prog
 
-/-- Constant time cost assigned to each primitive query. -/
-def timeOfQuery : {ι : Type} → QueryF ι → Nat
-  | _, .read _       => 1
-  | _, .write _ _    => 1
-  | _, .cmp _ _      => 1
+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] {ι : Type u} (q : QF ι) : Nat :=
+  Query.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))
+-- /-- 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))
 
 /-- Total time cost of running a program under the interpreter `timeInterp`. -/
-def timeProg {α : Type} (p : Prog α) : Nat :=
+def timeProg [Query QF] {α : Type u}  (p : Prog QF α) : Nat :=
   (p.liftM timeInterp).time
 
-/-- 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
 
 /-- Evaluate a query program to a pure value using `evalQuery`. -/
-def evalProg {α : Type} (p : Prog α) : α :=
+def evalProg [Query QF] {α : Type} (p : Prog QF α) : α :=
   FreeM.foldFreeM id
-    (fun {ι} (op : QueryF ι) (k : ι → α) =>
+    (fun {ι} (op : QF ι) (k : ι → α) =>
       k (evalQuery op))
     p
 

From b8c9a1efcb45a274236a4d0b32467797d7a9b8c5 Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasss94@gmail.com>
Date: Sat, 6 Dec 2025 02:44:12 +0100
Subject: [PATCH 12/13] Remove unnecessary instance

---
 Cslib/Algorithms/QueryModel.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index 4b604e93..9ee2f717 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -38,7 +38,6 @@ 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} : Bind (Prog QType) := inferInstance
 
 instance {QType : Type u → Type u} : Monad (Prog QType) := inferInstance
 namespace Prog
@@ -84,6 +83,7 @@ def timeInterp [Query QF] {ι : Type u} (q : QF ι) : Nat :=
 --   | _, .write i v   => TimeM.tick PUnit.unit (timeOfQuery (.write i v))
 --   | _, .cmp i j     => TimeM.tick false (timeOfQuery (.cmp i j))
 
+set_option diagnostics true
 /-- 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

From 9f1ba8ac5a0bb04153c2e736e9979b45c525fcce Mon Sep 17 00:00:00 2001
From: Shreyas <shreyasss94@gmail.com>
Date: Sat, 6 Dec 2025 03:05:40 +0100
Subject: [PATCH 13/13] monadLift instance needed

---
 Cslib/Algorithms/QueryModel.lean          | 14 +++++++++++---
 Cslib/Foundations/Control/Monad/Time.lean |  2 +-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/Cslib/Algorithms/QueryModel.lean b/Cslib/Algorithms/QueryModel.lean
index 9ee2f717..5a769901 100644
--- a/Cslib/Algorithms/QueryModel.lean
+++ b/Cslib/Algorithms/QueryModel.lean
@@ -40,6 +40,10 @@ abbrev Prog (QType : Type u → Type u) (α : Type v) := FreeM QType α
 
 
 instance {QType : Type u → Type u} : Monad (Prog QType) := inferInstance
+
+
+
+
 namespace Prog
 
 
@@ -74,8 +78,8 @@ open Query
 /--
 Interpret primitive queries into the time-counting monad `TimeM`.
 -/
-def timeInterp [Query QF] {ι : Type u} (q : QF ι) : Nat :=
-  Query.timeOfQuery q
+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 ι
@@ -83,7 +87,11 @@ def timeInterp [Query QF] {ι : Type u} (q : QF ι) : Nat :=
 --   | _, .write i v   => TimeM.tick PUnit.unit (timeOfQuery (.write i v))
 --   | _, .cmp i j     => TimeM.tick false (timeOfQuery (.cmp i j))
 
-set_option diagnostics true
+
+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
diff --git a/Cslib/Foundations/Control/Monad/Time.lean b/Cslib/Foundations/Control/Monad/Time.lean
index 72ed077c..41bce8f5 100644
--- a/Cslib/Foundations/Control/Monad/Time.lean
+++ b/Cslib/Foundations/Control/Monad/Time.lean
@@ -13,7 +13,7 @@ import Mathlib.Control.Monad.Writer
 as a simple cost model. As plain types it is isomorphic to `WriterT Nat Id`.
 -/
 
-structure TimeM (α : Type) where
+structure TimeM (α : Type u) where
   /-- The result of the computation. -/
   ret : α
   /-- The accumulated time cost. -/