feat: even stricter word ram

Algolean.lean

@@ -2,6 +2,7 @@ module  -- shake: keep-all --deprecated_module: ignore
 
 public import Algolean.AddWriter.Basic
 public import Algolean.AddWriter.Transformer
+public import Algolean.AddWriter.WP
 public import Algolean.Algorithms.BoyerMooreMajorityVote
 public import Algolean.Algorithms.Circuits.FanInTwo.LogAnd
 public import Algolean.Algorithms.KMPPatternSearch
@@ -13,13 +14,22 @@ public import Algolean.Algorithms.MergeSort
 public import Algolean.Algorithms.NaivePatternSearch
 public import Algolean.Algorithms.VecBubbleSort
 public import Algolean.Algorithms.VecSearch
-public import Algolean.Algorithms.WordRAMLinearSearch
+public import Algolean.Algorithms.WordRAM.Basic
+public import Algolean.Algorithms.WordRAM.BinarySearch.Algorithm
+public import Algolean.Algorithms.WordRAM.BinarySearch.Common
+public import Algolean.Algorithms.WordRAM.BinarySearch.Complexity
+public import Algolean.Algorithms.WordRAM.BinarySearch.Correctness
+public import Algolean.Algorithms.WordRAM.LinearSearch.Algorithm
+public import Algolean.Algorithms.WordRAM.LinearSearch.Common
+public import Algolean.Algorithms.WordRAM.LinearSearch.Complexity
+public import Algolean.Algorithms.WordRAM.LinearSearch.Correctness
 public import Algolean.Complexity.Basic
 public import Algolean.Complexity.PolytimeBasicClasses
 public import Algolean.FreeWP.Effects
 public import Algolean.FreeWP.WP
 public import Algolean.LowerBounds.ComparisonSort
 public import Algolean.ModelM
+public import Algolean.ModelStateM
 public import Algolean.Models.Arithmetic
 public import Algolean.Models.Circuits
 public import Algolean.Models.Comparison
@@ -36,5 +46,8 @@ public import Algolean.Models.ReadWriteVec
 public import Algolean.Models.RobertsonWebb
 public import Algolean.Models.SingleTapeTM
 public import Algolean.Models.WordRAM
+public import Algolean.Models.WordRAMSyntax
+public import Algolean.Problems.Basic
+public import Algolean.Problems.Search
 public import Algolean.QueryComposition
 public import Algolean.QueryModel

Algolean/AddWriter/Transformer.lean

@@ -120,6 +120,25 @@ theorem cost_bind [Monad m] [LawfulMonad m] [Add Cost]
       pure (a.tell + b.tell)) := by
   simp only [cost, run_bind, map_bind, map_pure]
 
+/-- Joint writer/state bind at a concrete state, without opaque intermediate pair matches. -/
+@[simp, grind =] theorem run_bind_state [Add Cost]
+    (x : AddWriterT Cost (StateM σ) α) (f : α → AddWriterT Cost (StateM σ) β) (s : σ) :
+    (x >>= f).run s =
+      let first := x.run s
+      let rest := (f first.fst.ret).run first.snd
+      ((⟨rest.fst.ret, first.fst.tell + rest.fst.tell⟩ : AddWriter Cost β), rest.snd) := rfl
+
+/-- Pure writer/state execution produces no cost. -/
+@[simp, grind =] theorem run_pure_state [Zero Cost] (a : α) (s : σ) :
+    (pure a : AddWriterT Cost (StateM σ) α).run s =
+      ((⟨a, 0⟩ : AddWriter Cost α), s) := rfl
+
+/-- Mapping changes the result while preserving the cost and state. -/
+@[simp, grind =] theorem run_map_state (f : α → β)
+    (x : AddWriterT Cost (StateM σ) α) (s : σ) :
+    (f <$> x).run s =
+      ((⟨f (x.run s).fst.ret, (x.run s).fst.tell⟩ : AddWriter Cost β), (x.run s).snd) := rfl
+
 @[ext] protected theorem ext
     (x y : AddWriterT Cost m α) (h : x.run = y.run) : x = y := h
 

Algolean/AddWriter/WP.lean

@@ -0,0 +1,87 @@
+/-
+Copyright (c) 2026 Shreyas Srinivas. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Shreyas Srinivas
+-/
+
+module
+
+public import Algolean.AddWriter.Transformer
+public import Std.Do
+
+/-!
+# Weakest preconditions for additive writer computations
+
+The extra postcondition argument is an accumulated cost. Starting it at zero gives the cost
+reported by `run`; an arbitrary initial cost permits compositional reasoning across binds.
+-/
+
+@[expose] public section
+
+namespace Algolean.AddWriterT
+
+open Std.Do Std.Do.WPMonad
+
+variable {Cost : Type u} {m : Type u → Type v} {ps : PostShape.{u}}
+
+/-- Interpret the writer output as an increment to a logical cost accumulator. -/
+def toStateT [Functor m] [Add Cost] (x : AddWriterT Cost m α) : StateT Cost m α :=
+  fun initial => (fun a => (a.ret, initial + a.tell)) <$> x.run
+
+/-- Expose the accumulated cost without unfolding the state transformer. -/
+@[simp] theorem toStateT_run [Functor m] [Add Cost]
+    (x : AddWriterT Cost m α) (initial : Cost) :
+    x.toStateT.run initial = (fun a => (a.ret, initial + a.tell)) <$> x.run := rfl
+
+@[simp] theorem toStateT_pure [Monad m] [LawfulMonad m] [AddZeroClass Cost] (a : α) :
+    toStateT (pure a : AddWriterT Cost m α) = pure a := by
+  funext initial
+  simp [toStateT]
+  rfl
+
+@[simp] theorem toStateT_bind [Monad m] [LawfulMonad m] [AddSemigroup Cost]
+    (x : AddWriterT Cost m α) (f : α → AddWriterT Cost m β) :
+    toStateT (x >>= f) = (toStateT x >>= fun a => toStateT (f a)) := by
+  apply StateT.ext
+  intro initial
+  rw [StateT.run_bind]
+  simp [toStateT, StateT.run, bind_map_left, add_assoc]
+
+/-- Writer weakest preconditions expose accumulated cost before the underlying post-shape. -/
+instance [Functor m] [Add Cost] [WP m ps] : WP (AddWriterT Cost m) (.arg Cost ps) where
+  wp x := wp x.toStateT
+
+/-- The cost-accumulator interpretation respects pure and bind. -/
+instance [Monad m] [AddMonoid Cost] [WPMonad m ps] :
+    WPMonad (AddWriterT Cost m) (.arg Cost ps) where
+  wp_pure a := by
+    change wp (toStateT (pure a : AddWriterT Cost m _)) = _
+    rw [toStateT_pure, wp_pure]
+  wp_bind x f := by
+    change wp (toStateT (x >>= f)) = _
+    rw [toStateT_bind, wp_bind]
+    rfl
+
+/-- Expose the underlying state interpretation to verification condition generation. -/
+theorem wp_eq_wp_toStateT [Functor m] [Add Cost] [WP m ps]
+    (x : AddWriterT Cost m α) : wp x = wp x.toStateT := rfl
+
+/-- A writer over state exposes its result, accumulated cost, and final physical state. -/
+@[simp] theorem wp_apply_state [Add Cost] (x : AddWriterT Cost (StateM σ) α)
+    (Q : PostCond α (.arg Cost (.arg σ .pure))) (initial : Cost) (s : σ) :
+    (wp x).apply Q initial s =
+      Q.fst (x.run s).fst.ret (initial + (x.run s).fst.tell) (x.run s).snd := rfl
+
+/-- Optional state execution exposes either the exact joint result or its failure postcondition. -/
+@[simp] theorem wp_apply_state_option [Add Cost] (x : AddWriterT Cost (StateT σ Option) α)
+    (Q : PostCond α (.arg Cost (.arg σ (.except PUnit .pure)))) (initial : Cost) (s : σ) :
+    (wp x).apply Q initial s =
+      match x.run s with
+      | none => Q.snd.fst PUnit.unit
+      | some (result, final) => Q.fst result.ret (initial + result.tell) final := by
+  dsimp [wp_eq_wp_toStateT, wp, toStateT, PredTrans.apply]
+  simp only [StateT.run_map]
+  dsimp only [StateT.run]
+  cases x.run s <;> rfl
+
+end Algolean.AddWriterT

Algolean/Algorithms/WordRAM/Basic.lean

@@ -0,0 +1,110 @@
+/-
+Copyright (c) 2026 Shreyas Srinivas. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Shreyas Srinivas
+-/
+
+module
+
+public import Algolean.Models.WordRAM
+public import Algolean.Problems.Search
+
+/-!
+# Supporting definitions for WordRAM search algorithms
+
+Search representations, the linear-search input layout and state constructor, output decoding,
+and input memory regions are defined together in `Algolean.Problems.Search`.
+
+This file supplies the bounds and initial memory layout used by binary search, word-address
+arithmetic, sorted-array lemmas, and `Executes`, the completed-execution relation used by
+both algorithms' correctness and complexity proofs.
+-/
+
+@[expose] public section
+
+namespace Algolean.Algorithms.WordRAM
+
+/-- The abstract sortedness relation instantiated with unsigned word order. -/
+abbrev SortedWords (input : Array (Word w)) : Prop :=
+  Search.SortedBy (fun a b => a.toNat ≤ b.toNat) input
+
+/-- Search inputs constrain the array and key register, not scratch registers or flags. -/
+structure RepresentsSearchInput (input : Search.Input (Word w)) (key : Register k)
+    (s : RAMState w k) : Prop extends RepresentsArray input.data s.Memory where
+  /-- The designated register contains the abstract key. -/
+  key_eq : s.Registers key = input.key
+
+/-- Runtime bounds for a search: an inclusive last address and a nonempty flag.
+This represents empty arrays and all `2^w` cells, even when `w = 0`. -/
+structure RepresentsBoundedSearchInput (input : Search.Input (Word w))
+    (key last : Register k) (s : RAMState w k) : Prop
+    extends RepresentsSearchInput input key s where
+  /-- Last input address; ignored for an empty input. -/
+  last_eq : s.Registers last = BitVec.ofNat w (input.data.size - 1)
+  /-- The initial less-than flag indicates whether there is an interval to search. -/
+  nonempty_eq : s.Flags .ult = decide (input.data.size ≠ 0)
+
+/-- Array layout used by the initial machine state. -/
+def arrayMemory (input : Array (BitVec w)) : Memory w :=
+  fun addr => input[addr.toNat]?.getD 0
+
+@[grind =] theorem wordAddress_toNat (i : Nat) (hi : i < 2 ^ w) :
+    (BitVec.ofNat w i).toNat = i := Nat.mod_eq_of_lt hi
+
+@[simp, grind =] theorem arrayMemory_ofNat (input : Array (BitVec w))
+    (hfits : input.size ≤ 2 ^ w) (i : Nat) (hi : i < input.size) :
+    arrayMemory input (BitVec.ofNat w i) = input[i] := by
+  simp [arrayMemory, BitVec.toNat_ofNat, Nat.mod_eq_of_lt (lt_of_lt_of_le hi hfits), hi]
+
+/-- If the array fits in memory, `arrayMemory` stores each element at its index. -/
+@[simp] theorem arrayMemory_represents (input : Array (Word w)) (hfits : input.size ≤ 2 ^ w) :
+    RepresentsArray input (arrayMemory input) :=
+  ⟨hfits, fun i hi => arrayMemory_ofNat input hfits i hi⟩
+
+@[grind =] theorem wordAddress_succ (i : Nat) :
+    BitVec.ofNat w i + 1 = BitVec.ofNat w (i + 1) := (BitVec.ofNat_add i 1).symm
+
+/-- In a sorted word array, words at or before a value below the key cannot match it. -/
+@[grind →] theorem SortedWords.exclude_left {input : Array (Word w)} (h : SortedWords input)
+    {target : Word w} {pivot : Nat} (hp : pivot < input.size)
+    (hlt : input[pivot].toNat < target.toNat) (i : Nat) (hi : i ≤ pivot) :
+    input[i]? ≠ some target := by
+  have hib : i < input.size := by lia
+  have hs := h i pivot hib hp hi
+  intro heq
+  have heq' : input[i] = target := by simpa [hib] using heq
+  rw [heq'] at hs
+  lia
+
+/-- In a sorted word array, words at or after a value above the key cannot match it. -/
+@[grind →] theorem SortedWords.exclude_right {input : Array (Word w)} (h : SortedWords input)
+    {target : Word w} {pivot : Nat} (hp : pivot < input.size)
+    (hlt : target.toNat < input[pivot].toNat) (i : Nat) (hi : pivot ≤ i)
+    (hib : i < input.size) : input[i]? ≠ some target := by
+  have hs := h pivot i hp hib hi
+  intro heq
+  have heq' : input[i] = target := by simpa [hib] using heq
+  rw [heq'] at hs
+  lia
+
+/-- Subtracting one converts a positive address to the preceding index without wrapping. -/
+theorem word_pred_toNat (x : BitVec w) (hx : 0 < x.toNat) :
+    (x - 1).toNat = x.toNat - 1 := by
+  have h := BitVec.ofNat_sub_ofNat_of_le (w := w) x.toNat 1 (by have := x.isLt; lia) hx
+  have h' := congrArg BitVec.toNat h
+  simpa [Nat.mod_eq_of_lt (show x.toNat - 1 < 2 ^ w by have := x.isLt; lia)] using h'
+
+/-- Completed execution in the time-and-space model, hiding interpreter fuel.
+Unused fuel is allowed and is not charged as time. -/
+def Executes (program : Prog (WordRAM w k) Unit) (s : RAMState w k)
+    (cost : RAMCost w k) (t : RAMState w k) : Prop :=
+  ∃ fuel remaining, execute fuel program s = some (⟨(), cost⟩, ⟨t, remaining⟩)
+
+/-- If a program's instructions finish with the stated cost and final state,
+the program satisfies `Executes` with that same cost and state. -/
+theorem Completes.executes {program : Prog (WordRAM w k) Unit}
+    (h : Completes (instructions program) s cost t) : Executes program s cost t := by
+  obtain ⟨fuel, hr⟩ := h.execute
+  exact ⟨fuel, 0, hr⟩
+
+end Algolean.Algorithms.WordRAM

Algolean/Algorithms/WordRAM/BinarySearch/Algorithm.lean

@@ -0,0 +1,91 @@
+/-
+Copyright (c) 2026 Shreyas Srinivas. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Shreyas Srinivas
+-/
+
+module
+
+public import Algolean.Algorithms.WordRAM.Basic
+public import Algolean.Models.WordRAMSyntax
+
+/-!
+# Binary search in the word-RAM model
+
+The program searches a sorted array using six registers and no extra memory.
+The initial state supplies the search key, the last array address, and a flag indicating
+whether the array is nonempty. The same program handles every input size that fits in
+memory at word width `w`, including arrays that use all `2 ^ w` cells.
+
+The midpoint is `lo + (hi - lo) / 2`. Bounds checks prevent address arithmetic from wrapping.
+
+Adapted from https://github.com/Shreyas4991/Algolean/pull/89.
+-/
+
+@[expose] public section
+
+namespace Algolean.Algorithms.WordRAM
+
+open scoped WordRAM Prog
+
+namespace BinarySearch
+
+/-- Register holding the first address still to search. -/
+abbrev lower : Register 6 := 0
+/-- Register holding the last address still to search. -/
+abbrev upper : Register 6 := 1
+/-- Register holding the midpoint address, or a matching address when found. -/
+abbrev middle : Register 6 := 2
+/-- Register holding the word loaded from the midpoint address. -/
+abbrev value : Register 6 := 3
+/-- Register holding the search key. -/
+abbrev key : Register 6 := 4
+/-- Register holding the constant one. -/
+abbrev one : Register 6 := 5
+
+/-- One machine iteration, with its continuation indicated by the less-than flag. -/
+def body (w : Nat) : Prog (WordRAM w 6) Unit := do [WordRAM w 6]
+  middle ←ᵣ upper - lower
+  middle ←ᵣ middle >>> one
+  middle ←ᵣ lower + middle
+  value ←ᵣ mem[middle]
+  ifₚ value =ᵣ key then
+    reset .ult
+  else
+    ifₚ value <ᵣ key then
+      ifₚ middle <ᵣ upper then
+        lower ←ᵣ middle + one
+      else
+        nop
+    else
+      ifₚ lower <ᵣ middle then
+        upper ←ᵣ middle - one
+      else
+        nop
+
+/-- Initialize the lower endpoint and increment constant; the upper endpoint is runtime input. -/
+def setup (w : Nat) : Prog (WordRAM w 6) Unit := do [WordRAM w 6]
+  lower ←ᵣ imm[0]
+  one ←ᵣ imm[1]
+
+end BinarySearch
+
+/-- Uniform binary search: width determines code; memory, key, last address, and the
+nonempty flag supply the runtime input. -/
+def binarySearch (w : Nat) : Prog (WordRAM w 6) Unit := do [WordRAM w 6]
+  reset .eq
+  ifₚ flag .ult then
+    BinarySearch.setup w
+    whileₚ .ult do
+      BinarySearch.body w
+  else
+    nop
+
+/-- Store the input array and search key, and initialize the bounds and flags for binary search. -/
+def binarySearchState (input : Array (Word w)) (target : Word w) : RAMState w 6 :=
+  ⟨arrayMemory input,
+    fun r => if r = BinarySearch.key then target
+      else if r = BinarySearch.upper then BitVec.ofNat w (input.size - 1) else 0,
+    fun op => if op = .ult then decide (input.size ≠ 0) else false⟩
+
+end Algolean.Algorithms.WordRAM