khora_control/cost_model.rs
1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Empirical cost model — fits a subsystem's measured runtime to a complexity
16//! class so the DCC can *anticipate* a budget breach instead of only reacting to
17//! it.
18//!
19//! Big-O cannot be derived from arbitrary code, so the engine **measures** it:
20//! given `(n, time)` samples (workload size vs observed milliseconds), it fits a
21//! constant factor `c` to each candidate class `f(n)` by least squares and keeps
22//! the best-fitting one. The prediction is then `cost(n) ≈ c · f(n)` — the same
23//! shape a database query optimiser uses (cardinality × per-row cost). This
24//! learns the cost *on the hardware it actually runs on*, instead of trusting a
25//! compile-time guess, and lets the cold path forecast "at this growth rate the
26//! frame budget breaks at ~N entities".
27//!
28//! This module is the pure, dependency-free *algorithm*. Feeding it live samples
29//! from telemetry and consuming its forecast in arbitration is the integration
30//! step performed by the DCC.
31
32/// A candidate complexity class the cost model can fit a measured workload to.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ComplexityClass {
35 /// `f(n) = 1` — cost independent of workload size.
36 Constant,
37 /// `f(n) = n` — linear in workload size.
38 Linear,
39 /// `f(n) = n · log₂(n)` — linearithmic.
40 Linearithmic,
41 /// `f(n) = n²` — quadratic.
42 Quadratic,
43}
44
45impl ComplexityClass {
46 /// Every class considered by the fitter.
47 pub const ALL: [ComplexityClass; 4] = [
48 ComplexityClass::Constant,
49 ComplexityClass::Linear,
50 ComplexityClass::Linearithmic,
51 ComplexityClass::Quadratic,
52 ];
53
54 /// Evaluates the (unscaled) growth function `f(n)` for this class.
55 pub fn f(self, n: f64) -> f64 {
56 let n = n.max(0.0);
57 match self {
58 ComplexityClass::Constant => 1.0,
59 ComplexityClass::Linear => n,
60 ComplexityClass::Linearithmic => {
61 if n <= 1.0 {
62 n
63 } else {
64 n * n.log2()
65 }
66 }
67 ComplexityClass::Quadratic => n * n,
68 }
69 }
70}
71
72/// One observation: a workload size and the time it took (milliseconds).
73#[derive(Debug, Clone, Copy)]
74pub struct CostSample {
75 /// Workload size (entities, bodies, draws, …).
76 pub n: f64,
77 /// Observed cost in milliseconds.
78 pub time_ms: f64,
79}
80
81/// A rolling empirical cost model for a single subsystem (agent / strategy).
82///
83/// Records `(n, time)` samples in a bounded ring buffer and fits `c · f(n)` over
84/// the candidate [`ComplexityClass`]es. Pure and dependency-free.
85#[derive(Debug, Clone)]
86pub struct CostModel {
87 samples: Vec<CostSample>,
88 capacity: usize,
89 next: usize,
90}
91
92impl CostModel {
93 /// Creates a model that keeps the most recent `capacity` samples (min 1).
94 pub fn new(capacity: usize) -> Self {
95 Self {
96 samples: Vec::new(),
97 capacity: capacity.max(1),
98 next: 0,
99 }
100 }
101
102 /// Records one `(n, time_ms)` observation. Overwrites the oldest sample once
103 /// the ring is full. Non-finite inputs are ignored.
104 pub fn record(&mut self, n: f64, time_ms: f64) {
105 if !n.is_finite() || !time_ms.is_finite() {
106 return;
107 }
108 let sample = CostSample { n, time_ms };
109 if self.samples.len() < self.capacity {
110 self.samples.push(sample);
111 } else {
112 self.samples[self.next] = sample;
113 self.next = (self.next + 1) % self.capacity;
114 }
115 }
116
117 /// Number of samples currently held.
118 pub fn len(&self) -> usize {
119 self.samples.len()
120 }
121
122 /// Whether no samples have been recorded.
123 pub fn is_empty(&self) -> bool {
124 self.samples.is_empty()
125 }
126
127 /// Fits each complexity class by least squares and returns the best
128 /// `(class, c)` — the class with the smallest residual and the constant
129 /// factor that scales its `f(n)` to the data.
130 ///
131 /// Returns `None` until there are at least two samples with distinct `n`
132 /// (a single workload size can't distinguish the classes).
133 pub fn best_fit(&self) -> Option<(ComplexityClass, f64)> {
134 if self.samples.len() < 2 {
135 return None;
136 }
137 let n0 = self.samples[0].n;
138 if self.samples.iter().all(|s| s.n == n0) {
139 return None;
140 }
141
142 // (class, c, residual)
143 let mut best: Option<(ComplexityClass, f64, f64)> = None;
144 for class in ComplexityClass::ALL {
145 // Least-squares `c` minimising Σ(c·fᵢ − tᵢ)² ⇒ c = Σ(fᵢ·tᵢ) / Σ(fᵢ²).
146 let mut num = 0.0;
147 let mut den = 0.0;
148 for s in &self.samples {
149 let f = class.f(s.n);
150 num += f * s.time_ms;
151 den += f * f;
152 }
153 if den <= f64::EPSILON {
154 continue;
155 }
156 let c = num / den;
157 let residual: f64 = self
158 .samples
159 .iter()
160 .map(|s| {
161 let err = c * class.f(s.n) - s.time_ms;
162 err * err
163 })
164 .sum();
165 if best.is_none_or(|(_, _, r)| residual < r) {
166 best = Some((class, c, residual));
167 }
168 }
169 best.map(|(class, c, _)| (class, c))
170 }
171
172 /// Predicts the cost in milliseconds at workload size `n` from the best fit.
173 /// `None` until [`best_fit`](Self::best_fit) is available.
174 pub fn predict_ms(&self, n: f64) -> Option<f64> {
175 self.best_fit().map(|(class, c)| c * class.f(n))
176 }
177
178 /// The most recently recorded observation's time in milliseconds, if any.
179 ///
180 /// Fallback cost signal when [`predict_ms`](Self::predict_ms) has no fit yet
181 /// (a stable workload never produces two distinct `n` values, so the model
182 /// can't pick a complexity class — but the raw measurement is still the best
183 /// available anchor).
184 pub fn latest_ms(&self) -> Option<f64> {
185 if self.samples.is_empty() {
186 return None;
187 }
188 let idx = if self.samples.len() < self.capacity {
189 self.samples.len() - 1
190 } else {
191 (self.next + self.capacity - 1) % self.capacity
192 };
193 Some(self.samples[idx].time_ms)
194 }
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200
201 #[test]
202 fn latest_ms_returns_most_recent_sample_across_ring_wrap() {
203 let mut m = CostModel::new(2);
204 assert!(m.latest_ms().is_none());
205 m.record(10.0, 1.0);
206 assert_eq!(m.latest_ms(), Some(1.0));
207 m.record(10.0, 2.0);
208 m.record(10.0, 3.0); // ring is full — overwrites the oldest slot
209 assert_eq!(m.latest_ms(), Some(3.0));
210 }
211
212 #[test]
213 fn fits_linear_and_predicts() {
214 let mut m = CostModel::new(16);
215 for n in [10.0, 20.0, 40.0, 80.0] {
216 m.record(n, 3.0 * n);
217 }
218 let (class, c) = m.best_fit().expect("fit available");
219 assert_eq!(class, ComplexityClass::Linear);
220 assert!((c - 3.0).abs() < 1e-6, "c = {c}");
221 // Extrapolate beyond the observed range.
222 assert!((m.predict_ms(100.0).unwrap() - 300.0).abs() < 1e-3);
223 }
224
225 #[test]
226 fn fits_quadratic() {
227 let mut m = CostModel::new(16);
228 for n in [2.0, 4.0, 8.0, 16.0] {
229 m.record(n, 0.5 * n * n);
230 }
231 assert_eq!(m.best_fit().unwrap().0, ComplexityClass::Quadratic);
232 }
233
234 #[test]
235 fn needs_two_distinct_sizes() {
236 let mut m = CostModel::new(16);
237 m.record(10.0, 5.0);
238 m.record(10.0, 5.0);
239 assert!(m.best_fit().is_none());
240 }
241
242 #[test]
243 fn ring_buffer_is_bounded() {
244 let mut m = CostModel::new(2);
245 m.record(1.0, 1.0);
246 m.record(2.0, 2.0);
247 m.record(3.0, 3.0);
248 assert_eq!(m.len(), 2);
249 }
250}