1use crate::{
2 key::{StableKey, StableKeyError},
3 schema::{SchemaMetadata, SchemaMetadataError},
4 slot::{AllocationSlotDescriptor, MemoryManagerSlotError},
5};
6use serde::{Deserialize, Serialize};
7use std::collections::BTreeSet;
8
9const DIAGNOSTIC_STRING_MAX_BYTES: usize = 256;
10
11#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
21pub struct AllocationDeclaration {
22 pub(crate) stable_key: StableKey,
24 pub(crate) slot: AllocationSlotDescriptor,
26 pub(crate) label: Option<String>,
28 pub(crate) schema: SchemaMetadata,
30}
31
32impl AllocationDeclaration {
33 pub fn new(
35 stable_key: impl AsRef<str>,
36 slot: AllocationSlotDescriptor,
37 label: Option<String>,
38 schema: SchemaMetadata,
39 ) -> Result<Self, DeclarationSnapshotError> {
40 let stable_key = StableKey::parse(stable_key).map_err(DeclarationSnapshotError::Key)?;
41 validate_label(label.as_deref())?;
42 schema
43 .validate()
44 .map_err(DeclarationSnapshotError::SchemaMetadata)?;
45 Ok(Self {
46 stable_key,
47 slot,
48 label,
49 schema,
50 })
51 }
52
53 pub fn memory_manager(
55 stable_key: impl AsRef<str>,
56 id: u8,
57 label: impl Into<String>,
58 ) -> Result<Self, DeclarationSnapshotError> {
59 Self::memory_manager_with_schema(stable_key, id, label, SchemaMetadata::default())
60 }
61
62 pub fn memory_manager_unlabeled(
64 stable_key: impl AsRef<str>,
65 id: u8,
66 ) -> Result<Self, DeclarationSnapshotError> {
67 Self::memory_manager_unlabeled_with_schema(stable_key, id, SchemaMetadata::default())
68 }
69
70 pub fn memory_manager_with_schema(
72 stable_key: impl AsRef<str>,
73 id: u8,
74 label: impl Into<String>,
75 schema: SchemaMetadata,
76 ) -> Result<Self, DeclarationSnapshotError> {
77 let slot = AllocationSlotDescriptor::memory_manager(id)
78 .map_err(DeclarationSnapshotError::MemoryManagerSlot)?;
79 Self::new(stable_key, slot, Some(label.into()), schema)
80 }
81
82 pub fn memory_manager_unlabeled_with_schema(
84 stable_key: impl AsRef<str>,
85 id: u8,
86 schema: SchemaMetadata,
87 ) -> Result<Self, DeclarationSnapshotError> {
88 let slot = AllocationSlotDescriptor::memory_manager(id)
89 .map_err(DeclarationSnapshotError::MemoryManagerSlot)?;
90 Self::new(stable_key, slot, None, schema)
91 }
92
93 #[must_use]
95 pub const fn stable_key(&self) -> &StableKey {
96 &self.stable_key
97 }
98
99 #[must_use]
101 pub const fn slot(&self) -> &AllocationSlotDescriptor {
102 &self.slot
103 }
104
105 #[must_use]
107 pub fn label(&self) -> Option<&str> {
108 self.label.as_deref()
109 }
110
111 #[must_use]
113 pub const fn schema(&self) -> &SchemaMetadata {
114 &self.schema
115 }
116}
117
118#[derive(Clone, Debug, Default)]
127pub struct DeclarationCollector {
128 declarations: Vec<AllocationDeclaration>,
129}
130
131impl DeclarationCollector {
132 #[must_use]
134 pub const fn new() -> Self {
135 Self {
136 declarations: Vec::new(),
137 }
138 }
139
140 pub fn push(&mut self, declaration: AllocationDeclaration) {
142 self.declarations.push(declaration);
143 }
144
145 pub fn declare(&mut self, declaration: AllocationDeclaration) -> &mut Self {
147 self.push(declaration);
148 self
149 }
150
151 #[must_use]
153 pub fn with_declaration(mut self, declaration: AllocationDeclaration) -> Self {
154 self.push(declaration);
155 self
156 }
157
158 pub fn declare_memory_manager(
160 &mut self,
161 stable_key: impl AsRef<str>,
162 id: u8,
163 label: impl Into<String>,
164 ) -> Result<&mut Self, DeclarationSnapshotError> {
165 self.declare_memory_manager_with_schema(stable_key, id, label, SchemaMetadata::default())
166 }
167
168 pub fn declare_memory_manager_unlabeled(
170 &mut self,
171 stable_key: impl AsRef<str>,
172 id: u8,
173 ) -> Result<&mut Self, DeclarationSnapshotError> {
174 self.declare_memory_manager_unlabeled_with_schema(stable_key, id, SchemaMetadata::default())
175 }
176
177 pub fn declare_memory_manager_with_schema(
179 &mut self,
180 stable_key: impl AsRef<str>,
181 id: u8,
182 label: impl Into<String>,
183 schema: SchemaMetadata,
184 ) -> Result<&mut Self, DeclarationSnapshotError> {
185 self.push(AllocationDeclaration::memory_manager_with_schema(
186 stable_key, id, label, schema,
187 )?);
188 Ok(self)
189 }
190
191 pub fn declare_memory_manager_unlabeled_with_schema(
193 &mut self,
194 stable_key: impl AsRef<str>,
195 id: u8,
196 schema: SchemaMetadata,
197 ) -> Result<&mut Self, DeclarationSnapshotError> {
198 self.push(AllocationDeclaration::memory_manager_unlabeled_with_schema(
199 stable_key, id, schema,
200 )?);
201 Ok(self)
202 }
203
204 pub fn with_memory_manager(
206 mut self,
207 stable_key: impl AsRef<str>,
208 id: u8,
209 label: impl Into<String>,
210 ) -> Result<Self, DeclarationSnapshotError> {
211 self.declare_memory_manager(stable_key, id, label)?;
212 Ok(self)
213 }
214
215 pub fn with_memory_manager_unlabeled(
217 mut self,
218 stable_key: impl AsRef<str>,
219 id: u8,
220 ) -> Result<Self, DeclarationSnapshotError> {
221 self.declare_memory_manager_unlabeled(stable_key, id)?;
222 Ok(self)
223 }
224
225 pub fn with_memory_manager_schema(
227 mut self,
228 stable_key: impl AsRef<str>,
229 id: u8,
230 label: impl Into<String>,
231 schema: SchemaMetadata,
232 ) -> Result<Self, DeclarationSnapshotError> {
233 self.declare_memory_manager_with_schema(stable_key, id, label, schema)?;
234 Ok(self)
235 }
236
237 pub fn with_memory_manager_unlabeled_schema(
239 mut self,
240 stable_key: impl AsRef<str>,
241 id: u8,
242 schema: SchemaMetadata,
243 ) -> Result<Self, DeclarationSnapshotError> {
244 self.declare_memory_manager_unlabeled_with_schema(stable_key, id, schema)?;
245 Ok(self)
246 }
247
248 pub fn seal(self) -> Result<DeclarationSnapshot, DeclarationSnapshotError> {
250 DeclarationSnapshot::new(self.declarations)
251 }
252}
253
254#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
263pub struct DeclarationSnapshot {
264 declarations: Vec<AllocationDeclaration>,
266 runtime_fingerprint: Option<String>,
268}
269
270impl DeclarationSnapshot {
271 pub fn new(declarations: Vec<AllocationDeclaration>) -> Result<Self, DeclarationSnapshotError> {
273 for declaration in &declarations {
274 validate_label(declaration.label.as_deref())?;
275 declaration
276 .schema
277 .validate()
278 .map_err(DeclarationSnapshotError::SchemaMetadata)?;
279 }
280 reject_duplicates(&declarations)?;
281 Ok(Self {
282 declarations,
283 runtime_fingerprint: None,
284 })
285 }
286
287 pub fn with_runtime_fingerprint(
289 mut self,
290 fingerprint: impl Into<String>,
291 ) -> Result<Self, DeclarationSnapshotError> {
292 let fingerprint = fingerprint.into();
293 validate_runtime_fingerprint(Some(&fingerprint))?;
294 self.runtime_fingerprint = Some(fingerprint);
295 Ok(self)
296 }
297
298 #[must_use]
300 pub fn is_empty(&self) -> bool {
301 self.declarations.is_empty()
302 }
303
304 #[must_use]
306 pub fn len(&self) -> usize {
307 self.declarations.len()
308 }
309
310 #[must_use]
312 pub fn declarations(&self) -> &[AllocationDeclaration] {
313 &self.declarations
314 }
315
316 #[must_use]
318 pub fn runtime_fingerprint(&self) -> Option<&str> {
319 self.runtime_fingerprint.as_deref()
320 }
321
322 pub(crate) fn into_parts(self) -> (Vec<AllocationDeclaration>, Option<String>) {
323 (self.declarations, self.runtime_fingerprint)
324 }
325}
326
327#[derive(Clone, Debug, Eq, thiserror::Error, PartialEq)]
332pub enum DeclarationSnapshotError {
333 #[error(transparent)]
335 Key(StableKeyError),
336 #[error(transparent)]
338 MemoryManagerSlot(MemoryManagerSlotError),
339 #[error(transparent)]
341 SchemaMetadata(SchemaMetadataError),
342 #[error("stable key '{0}' is declared more than once")]
344 DuplicateStableKey(StableKey),
345 #[error("allocation slot '{0:?}' is declared more than once")]
347 DuplicateSlot(AllocationSlotDescriptor),
348 #[error("allocation declaration label must not be empty when present")]
350 EmptyLabel,
351 #[error("allocation declaration label must be at most 256 bytes")]
353 LabelTooLong,
354 #[error("allocation declaration label must be ASCII")]
356 NonAsciiLabel,
357 #[error("allocation declaration label must not contain ASCII control characters")]
359 ControlCharacterLabel,
360 #[error("runtime_fingerprint must not be empty when present")]
362 EmptyRuntimeFingerprint,
363 #[error("runtime_fingerprint must be at most 256 bytes")]
365 RuntimeFingerprintTooLong,
366 #[error("runtime_fingerprint must be ASCII")]
368 NonAsciiRuntimeFingerprint,
369 #[error("runtime_fingerprint must not contain ASCII control characters")]
371 ControlCharacterRuntimeFingerprint,
372}
373
374fn validate_label(label: Option<&str>) -> Result<(), DeclarationSnapshotError> {
375 let Some(label) = label else {
376 return Ok(());
377 };
378 if label.is_empty() {
379 return Err(DeclarationSnapshotError::EmptyLabel);
380 }
381 if label.len() > DIAGNOSTIC_STRING_MAX_BYTES {
382 return Err(DeclarationSnapshotError::LabelTooLong);
383 }
384 if !label.is_ascii() {
385 return Err(DeclarationSnapshotError::NonAsciiLabel);
386 }
387 if label.bytes().any(|byte| byte.is_ascii_control()) {
388 return Err(DeclarationSnapshotError::ControlCharacterLabel);
389 }
390 Ok(())
391}
392
393pub(crate) fn validate_runtime_fingerprint(
394 fingerprint: Option<&str>,
395) -> Result<(), DeclarationSnapshotError> {
396 let Some(fingerprint) = fingerprint else {
397 return Ok(());
398 };
399 if fingerprint.is_empty() {
400 return Err(DeclarationSnapshotError::EmptyRuntimeFingerprint);
401 }
402 if fingerprint.len() > DIAGNOSTIC_STRING_MAX_BYTES {
403 return Err(DeclarationSnapshotError::RuntimeFingerprintTooLong);
404 }
405 if !fingerprint.is_ascii() {
406 return Err(DeclarationSnapshotError::NonAsciiRuntimeFingerprint);
407 }
408 if fingerprint.bytes().any(|byte| byte.is_ascii_control()) {
409 return Err(DeclarationSnapshotError::ControlCharacterRuntimeFingerprint);
410 }
411 Ok(())
412}
413
414fn reject_duplicates(
415 declarations: &[AllocationDeclaration],
416) -> Result<(), DeclarationSnapshotError> {
417 let mut keys = BTreeSet::new();
418 let mut slots = BTreeSet::new();
419
420 for declaration in declarations {
421 if !slots.insert(declaration.slot.clone()) {
422 return Err(DeclarationSnapshotError::DuplicateSlot(
423 declaration.slot.clone(),
424 ));
425 }
426 if !keys.insert(declaration.stable_key.clone()) {
427 return Err(DeclarationSnapshotError::DuplicateStableKey(
428 declaration.stable_key.clone(),
429 ));
430 }
431 }
432
433 Ok(())
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439 use crate::slot::AllocationSlotDescriptor;
440
441 fn declaration(key: &str, id: u8) -> AllocationDeclaration {
442 AllocationDeclaration::new(
443 key,
444 AllocationSlotDescriptor::memory_manager(id).expect("usable slot"),
445 None,
446 SchemaMetadata::default(),
447 )
448 .expect("declaration")
449 }
450
451 #[test]
452 fn declaration_rejects_unbounded_label_metadata() {
453 let err = AllocationDeclaration::new(
454 "app.users.v1",
455 AllocationSlotDescriptor::memory_manager(100).expect("usable slot"),
456 Some("x".repeat(257)),
457 SchemaMetadata::default(),
458 )
459 .expect_err("label too long");
460
461 assert_eq!(err, DeclarationSnapshotError::LabelTooLong);
462 }
463
464 #[test]
465 fn memory_manager_declaration_constructor_builds_common_declaration() {
466 let declaration = AllocationDeclaration::memory_manager("app.orders.v1", 100, "orders")
467 .expect("declaration");
468
469 assert_eq!(declaration.stable_key.as_str(), "app.orders.v1");
470 assert_eq!(
471 declaration.slot,
472 AllocationSlotDescriptor::memory_manager(100).expect("usable slot")
473 );
474 assert_eq!(declaration.label.as_deref(), Some("orders"));
475 assert_eq!(declaration.schema, SchemaMetadata::default());
476 }
477
478 #[test]
479 fn memory_manager_declaration_constructor_rejects_invalid_slot() {
480 let err = AllocationDeclaration::memory_manager("app.orders.v1", u8::MAX, "orders")
481 .expect_err("sentinel must fail");
482
483 assert!(matches!(
484 err,
485 DeclarationSnapshotError::MemoryManagerSlot(_)
486 ));
487 }
488
489 #[test]
490 fn declaration_collector_declares_memory_manager_allocations() {
491 let mut declarations = DeclarationCollector::new();
492 declarations
493 .declare_memory_manager("app.orders.v1", 100, "orders")
494 .expect("orders declaration")
495 .declare_memory_manager_unlabeled("app.users.v1", 101)
496 .expect("users declaration");
497
498 let snapshot = declarations.seal().expect("snapshot");
499
500 assert_eq!(snapshot.len(), 2);
501 assert_eq!(
502 snapshot.declarations()[0].slot,
503 AllocationSlotDescriptor::memory_manager(100).expect("usable slot")
504 );
505 assert_eq!(snapshot.declarations()[0].label.as_deref(), Some("orders"));
506 assert_eq!(snapshot.declarations()[1].label, None);
507 }
508
509 #[test]
510 fn declaration_collector_builder_declares_memory_manager_allocations() {
511 let snapshot = DeclarationCollector::new()
512 .with_memory_manager("app.orders.v1", 100, "orders")
513 .expect("orders declaration")
514 .with_memory_manager_unlabeled("app.users.v1", 101)
515 .expect("users declaration")
516 .seal()
517 .expect("snapshot");
518
519 assert_eq!(snapshot.len(), 2);
520 }
521
522 #[test]
523 fn snapshot_rejects_unbounded_runtime_fingerprint() {
524 let snapshot =
525 DeclarationSnapshot::new(vec![declaration("app.users.v1", 100)]).expect("snapshot");
526
527 let err = snapshot
528 .with_runtime_fingerprint("x".repeat(257))
529 .expect_err("fingerprint too long");
530
531 assert_eq!(err, DeclarationSnapshotError::RuntimeFingerprintTooLong);
532 }
533
534 #[test]
535 fn rejects_duplicate_keys() {
536 let err = DeclarationSnapshot::new(vec![
537 declaration("app.users.v1", 100),
538 declaration("app.users.v1", 101),
539 ])
540 .expect_err("duplicate key");
541
542 assert!(matches!(
543 err,
544 DeclarationSnapshotError::DuplicateStableKey(_)
545 ));
546 }
547
548 #[test]
549 fn rejects_duplicate_slots() {
550 let err = DeclarationSnapshot::new(vec![
551 declaration("app.users.v1", 100),
552 declaration("app.orders.v1", 100),
553 ])
554 .expect_err("duplicate slot");
555
556 assert!(matches!(err, DeclarationSnapshotError::DuplicateSlot(_)));
557 }
558}