1use core::sync::atomic::{AtomicU64, Ordering::Relaxed};
16
17use crate::memory::tag::{MemTag, Realm};
18
19struct Cell {
22 bytes: AtomicU64,
23 peak: AtomicU64,
24 budget: AtomicU64,
25}
26
27impl Cell {
28 const fn new() -> Self {
29 Self {
30 bytes: AtomicU64::new(0),
31 peak: AtomicU64::new(0),
32 budget: AtomicU64::new(0),
33 }
34 }
35
36 fn observe(&self, bytes: u64) {
37 self.peak.fetch_max(bytes, Relaxed);
38 }
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub struct TagUsage {
44 pub tag: MemTag,
45 pub realm: Realm,
46 pub bytes: u64,
47 pub peak_bytes: u64,
49 pub budget: Option<u64>,
51}
52
53impl TagUsage {
54 pub const fn empty(tag: MemTag, realm: Realm) -> Self {
58 Self {
59 tag,
60 realm,
61 bytes: 0,
62 peak_bytes: 0,
63 budget: None,
64 }
65 }
66
67 pub(crate) fn is_reported(&self) -> bool {
68 self.bytes > 0 || self.peak_bytes > 0 || self.budget.is_some()
69 }
70
71 pub fn over_budget(&self) -> bool {
72 self.budget.is_some_and(|b| self.bytes > b)
73 }
74
75 pub fn fraction(&self) -> Option<f32> {
77 self.budget
78 .filter(|&b| b > 0)
79 .map(|b| self.bytes as f32 / b as f32)
80 }
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub struct LedgerSnapshot {
86 entries: [TagUsage; MemTag::COUNT * Realm::COUNT],
87}
88
89impl LedgerSnapshot {
90 pub fn get(&self, tag: MemTag, realm: Realm) -> TagUsage {
92 self.entries[entry_index(tag, realm)]
93 }
94
95 pub fn reported(&self, realm: Realm) -> impl Iterator<Item = TagUsage> + '_ {
98 MemTag::ALL
99 .into_iter()
100 .map(move |tag| self.get(tag, realm))
101 .filter(TagUsage::is_reported)
102 }
103
104 pub fn realm_bytes(&self, realm: Realm) -> u64 {
108 MemTag::ALL
109 .into_iter()
110 .map(|tag| self.get(tag, realm).bytes)
111 .sum()
112 }
113
114 pub fn is_empty(&self) -> bool {
116 !self.entries.iter().any(TagUsage::is_reported)
117 }
118}
119
120impl Default for LedgerSnapshot {
121 fn default() -> Self {
122 Self {
123 entries: core::array::from_fn(|i| {
124 TagUsage::empty(MemTag::ALL[i / Realm::COUNT], Realm::ALL[i % Realm::COUNT])
125 }),
126 }
127 }
128}
129
130const fn entry_index(tag: MemTag, realm: Realm) -> usize {
131 tag.index() * Realm::COUNT + realm.index()
132}
133
134pub struct Ledger {
138 cells: [Cell; MemTag::COUNT * Realm::COUNT],
139}
140
141impl Ledger {
142 pub const fn new() -> Self {
144 Self {
145 cells: [const { Cell::new() }; MemTag::COUNT * Realm::COUNT],
146 }
147 }
148
149 fn cell(&self, tag: MemTag, realm: Realm) -> &Cell {
150 &self.cells[entry_index(tag, realm)]
151 }
152
153 pub fn add(&self, tag: MemTag, realm: Realm, bytes: u64) {
156 let cell = self.cell(tag, realm);
157 let total = cell.bytes.fetch_add(bytes, Relaxed).saturating_add(bytes);
158 cell.observe(total);
159 }
160
161 pub fn release(&self, tag: MemTag, realm: Realm, bytes: u64) {
165 let cell = self.cell(tag, realm);
166 let _ = cell
167 .bytes
168 .fetch_update(Relaxed, Relaxed, |held| Some(held.saturating_sub(bytes)));
169 }
170
171 pub fn set(&self, tag: MemTag, realm: Realm, bytes: u64) {
175 let cell = self.cell(tag, realm);
176 cell.bytes.store(bytes, Relaxed);
177 cell.observe(bytes);
178 }
179
180 pub fn set_budget(&self, tag: MemTag, realm: Realm, budget: Option<u64>) {
182 self.cell(tag, realm)
183 .budget
184 .store(budget.unwrap_or(0), Relaxed);
185 }
186
187 pub fn usage(&self, tag: MemTag, realm: Realm) -> TagUsage {
189 let cell = self.cell(tag, realm);
190 let budget = cell.budget.load(Relaxed);
191 TagUsage {
192 tag,
193 realm,
194 bytes: cell.bytes.load(Relaxed),
195 peak_bytes: cell.peak.load(Relaxed),
196 budget: (budget > 0).then_some(budget),
197 }
198 }
199
200 pub fn snapshot(&self) -> LedgerSnapshot {
202 LedgerSnapshot {
203 entries: core::array::from_fn(|i| {
204 self.usage(MemTag::ALL[i / Realm::COUNT], Realm::ALL[i % Realm::COUNT])
205 }),
206 }
207 }
208
209 pub fn clear(&self) {
212 for cell in &self.cells {
213 cell.bytes.store(0, Relaxed);
214 cell.peak.store(0, Relaxed);
215 cell.budget.store(0, Relaxed);
216 }
217 }
218}
219
220impl Default for Ledger {
221 fn default() -> Self {
222 Self::new()
223 }
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229
230 #[test]
231 fn an_untouched_ledger_reports_nothing() {
232 let ledger = Ledger::new();
233 let snap = ledger.snapshot();
234 assert!(snap.is_empty());
235 assert_eq!(snap.reported(Realm::Host).count(), 0);
236 assert_eq!(snap.realm_bytes(Realm::Device), 0);
237 }
238
239 #[test]
240 fn increments_accumulate_and_release_gives_them_back() {
241 let ledger = Ledger::new();
242 ledger.add(MemTag::Audio, Realm::Host, 1024);
243 ledger.add(MemTag::Audio, Realm::Host, 512);
244 assert_eq!(ledger.usage(MemTag::Audio, Realm::Host).bytes, 1536);
245
246 ledger.release(MemTag::Audio, Realm::Host, 512);
247 assert_eq!(ledger.usage(MemTag::Audio, Realm::Host).bytes, 1024);
248 }
249
250 #[test]
253 fn releasing_more_than_held_saturates_at_zero() {
254 let ledger = Ledger::new();
255 ledger.add(MemTag::Meshes, Realm::Device, 100);
256 ledger.release(MemTag::Meshes, Realm::Device, 4096);
257 assert_eq!(ledger.usage(MemTag::Meshes, Realm::Device).bytes, 0);
258 }
259
260 #[test]
263 fn set_replaces_rather_than_accumulates() {
264 let ledger = Ledger::new();
265 ledger.set(MemTag::Textures, Realm::Device, 8_000);
266 ledger.set(MemTag::Textures, Realm::Device, 6_000);
267 assert_eq!(ledger.usage(MemTag::Textures, Realm::Device).bytes, 6_000);
268 }
269
270 #[test]
271 fn peak_holds_the_high_water_mark_across_both_report_styles() {
272 let ledger = Ledger::new();
273 ledger.set(MemTag::Textures, Realm::Device, 9_000);
274 ledger.set(MemTag::Textures, Realm::Device, 1_000);
275 ledger.add(MemTag::Textures, Realm::Device, 500);
276
277 let usage = ledger.usage(MemTag::Textures, Realm::Device);
278 assert_eq!(usage.bytes, 1_500);
279 assert_eq!(usage.peak_bytes, 9_000);
280 }
281
282 #[test]
283 fn a_budget_reads_back_and_clears() {
284 let ledger = Ledger::new();
285 ledger.set_budget(MemTag::Chunks, Realm::Device, Some(4_096));
286 assert_eq!(
287 ledger.usage(MemTag::Chunks, Realm::Device).budget,
288 Some(4_096)
289 );
290
291 ledger.set_budget(MemTag::Chunks, Realm::Device, None);
292 assert_eq!(ledger.usage(MemTag::Chunks, Realm::Device).budget, None);
293 }
294
295 #[test]
296 fn over_budget_and_fraction_track_the_ceiling() {
297 let ledger = Ledger::new();
298 ledger.set_budget(MemTag::Meshes, Realm::Device, Some(1_000));
299 ledger.set(MemTag::Meshes, Realm::Device, 500);
300 let under = ledger.usage(MemTag::Meshes, Realm::Device);
301 assert!(!under.over_budget());
302 assert_eq!(under.fraction(), Some(0.5));
303
304 ledger.set(MemTag::Meshes, Realm::Device, 1_200);
305 assert!(ledger.usage(MemTag::Meshes, Realm::Device).over_budget());
306 }
307
308 #[test]
311 fn an_unbudgeted_tag_has_no_fraction() {
312 let ledger = Ledger::new();
313 ledger.set(MemTag::Scratch, Realm::Host, 4_096);
314 let usage = ledger.usage(MemTag::Scratch, Realm::Host);
315 assert_eq!(usage.fraction(), None);
316 assert!(!usage.over_budget());
317 }
318
319 #[test]
322 fn the_realms_are_counted_separately() {
323 let ledger = Ledger::new();
324 ledger.set(MemTag::Textures, Realm::Device, 2_048);
325 ledger.set(MemTag::Textures, Realm::Host, 64);
326
327 let snap = ledger.snapshot();
328 assert_eq!(snap.get(MemTag::Textures, Realm::Device).bytes, 2_048);
329 assert_eq!(snap.get(MemTag::Textures, Realm::Host).bytes, 64);
330 assert_eq!(snap.realm_bytes(Realm::Device), 2_048);
331 assert_eq!(snap.realm_bytes(Realm::Host), 64);
332 }
333
334 #[test]
337 fn reported_lists_touched_tags_in_vocabulary_order() {
338 let ledger = Ledger::new();
339 ledger.set(MemTag::Chunks, Realm::Device, 3);
340 ledger.set(MemTag::Textures, Realm::Device, 1);
341 ledger.set(MemTag::Meshes, Realm::Device, 2);
342 ledger.set(MemTag::Audio, Realm::Host, 9);
343
344 let tags: std::vec::Vec<MemTag> = ledger
345 .snapshot()
346 .reported(Realm::Device)
347 .map(|u| u.tag)
348 .collect();
349 assert_eq!(tags, [MemTag::Textures, MemTag::Meshes, MemTag::Chunks]);
350 }
351
352 #[test]
355 fn a_tag_that_dropped_to_zero_stays_reported() {
356 let ledger = Ledger::new();
357 ledger.set(MemTag::Textures, Realm::Device, 4_096);
358 ledger.set(MemTag::Textures, Realm::Device, 0);
359
360 let snap = ledger.snapshot();
361 assert!(snap.get(MemTag::Textures, Realm::Device).is_reported());
362 assert_eq!(snap.reported(Realm::Device).count(), 1);
363 }
364
365 #[test]
366 fn clear_drops_every_report_and_budget() {
367 let ledger = Ledger::new();
368 ledger.set(MemTag::Textures, Realm::Device, 4_096);
369 ledger.set_budget(MemTag::Textures, Realm::Device, Some(8_192));
370 ledger.clear();
371
372 assert!(ledger.snapshot().is_empty());
373 assert_eq!(ledger.usage(MemTag::Textures, Realm::Device).peak_bytes, 0);
374 }
375
376 #[test]
379 fn every_tag_and_realm_addresses_its_own_cell() {
380 let ledger = Ledger::new();
381 for (i, tag) in MemTag::ALL.into_iter().enumerate() {
382 for (j, realm) in Realm::ALL.into_iter().enumerate() {
383 ledger.set(tag, realm, (i * Realm::COUNT + j + 1) as u64);
384 }
385 }
386 let snap = ledger.snapshot();
387 for (i, tag) in MemTag::ALL.into_iter().enumerate() {
388 for (j, realm) in Realm::ALL.into_iter().enumerate() {
389 let usage = snap.get(tag, realm);
390 assert_eq!(usage.bytes, (i * Realm::COUNT + j + 1) as u64);
391 assert_eq!(usage.tag, tag);
392 assert_eq!(usage.realm, realm);
393 }
394 }
395 }
396}