use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use core::fmt;
pub const MAX_BUDGET_SHARE_BPS: u16 = 10_000;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChildAdmission {
pub share_bps: u16,
pub holders: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BudgetSplit {
pub parent_token_id: String,
pub parent_share_bps: u16,
pub children: BTreeMap<String, ChildAdmission>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BudgetSplitError {
ChildShareExceedsCap {
child_id: String,
share_bps: u16,
},
OversubscribedSiblings {
child_id: String,
share_bps: u16,
current_total_child_bps: u32,
parent_share_bps: u16,
},
UnknownParent {
parent_token_id: String,
},
DuplicateChild {
child_id: String,
},
ReleaseShareMismatch {
child_id: String,
expected_share_bps: u16,
actual_share_bps: u16,
},
}
impl fmt::Display for BudgetSplitError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
BudgetSplitError::ChildShareExceedsCap {
child_id,
share_bps,
} => write!(
f,
"child {child_id} share {share_bps} bps exceeds the {} bps per-token cap",
MAX_BUDGET_SHARE_BPS
),
BudgetSplitError::OversubscribedSiblings {
child_id,
share_bps,
current_total_child_bps,
parent_share_bps,
} => write!(
f,
"child {child_id} share {share_bps} bps + sibling sum {current_total_child_bps} bps > parent share {parent_share_bps} bps"
),
BudgetSplitError::UnknownParent { parent_token_id } => {
write!(f, "parent capability {parent_token_id} is not registered")
}
BudgetSplitError::DuplicateChild { child_id } => write!(
f,
"child capability {child_id} already admitted under this parent with a different share"
),
BudgetSplitError::ReleaseShareMismatch {
child_id,
expected_share_bps,
actual_share_bps,
} => write!(
f,
"child capability {child_id} release share {expected_share_bps} bps does not match admitted share {actual_share_bps} bps"
),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdmitMode {
Lease,
VerifyOnly,
}
impl BudgetSplit {
#[must_use]
pub fn new(parent_token_id: String, parent_share_bps: u16) -> Self {
Self {
parent_token_id,
parent_share_bps,
children: BTreeMap::new(),
}
}
#[must_use]
pub fn current_total_child_bps(&self) -> u32 {
self.children
.values()
.map(|admission| u32::from(admission.share_bps))
.sum()
}
#[must_use]
pub fn child_share_bps(&self, child_id: &str) -> Option<u16> {
self.children
.get(child_id)
.map(|admission| admission.share_bps)
}
#[must_use]
pub fn child_holders(&self, child_id: &str) -> Option<usize> {
self.children
.get(child_id)
.map(|admission| admission.holders)
}
pub fn try_admit_child(
&mut self,
child_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError> {
self.admit_child(child_id, share_bps, AdmitMode::Lease)
}
pub fn verify_child_admission(
&mut self,
child_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError> {
self.admit_child(child_id, share_bps, AdmitMode::VerifyOnly)
}
fn admit_child(
&mut self,
child_id: String,
share_bps: u16,
mode: AdmitMode,
) -> Result<(), BudgetSplitError> {
if share_bps > MAX_BUDGET_SHARE_BPS {
return Err(BudgetSplitError::ChildShareExceedsCap {
child_id,
share_bps,
});
}
if let Some(existing) = self.children.get_mut(&child_id) {
if existing.share_bps == share_bps {
if let AdmitMode::Lease = mode {
existing.holders = existing.holders.saturating_add(1);
}
return Ok(());
}
return Err(BudgetSplitError::DuplicateChild { child_id });
}
let running = self.current_total_child_bps();
let proposed_total = running.saturating_add(u32::from(share_bps));
if proposed_total > u32::from(self.parent_share_bps) {
return Err(BudgetSplitError::OversubscribedSiblings {
child_id,
share_bps,
current_total_child_bps: running,
parent_share_bps: self.parent_share_bps,
});
}
let holders = match mode {
AdmitMode::Lease => 1,
AdmitMode::VerifyOnly => 0,
};
self.children
.insert(child_id, ChildAdmission { share_bps, holders });
Ok(())
}
pub fn release_child(
&mut self,
child_id: &str,
expected_share_bps: u16,
) -> Result<(), BudgetSplitError> {
let Some(admission) = self.children.get_mut(child_id) else {
return Ok(());
};
if admission.share_bps != expected_share_bps {
return Err(BudgetSplitError::ReleaseShareMismatch {
child_id: child_id.to_string(),
expected_share_bps,
actual_share_bps: admission.share_bps,
});
}
admission.holders = admission.holders.saturating_sub(1);
if admission.holders == 0 {
let _ = self.children.remove(child_id);
}
Ok(())
}
}
pub trait BudgetRegistry {
fn register_parent(
&mut self,
parent_token_id: String,
parent_share_bps: u16,
) -> Result<(), BudgetSplitError>;
fn try_admit_child(
&mut self,
parent_token_id: &str,
child_token_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError>;
fn verify_child_admission(
&mut self,
parent_token_id: &str,
child_token_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError>;
fn release_child(
&mut self,
parent_token_id: &str,
child_token_id: &str,
expected_share_bps: u16,
) -> Result<(), BudgetSplitError>;
fn evict_parent(&mut self, parent_token_id: &str);
}
#[derive(Debug, Default, Clone, Copy)]
pub struct NoopBudgetRegistry;
impl BudgetRegistry for NoopBudgetRegistry {
fn register_parent(
&mut self,
_parent_token_id: String,
_parent_share_bps: u16,
) -> Result<(), BudgetSplitError> {
Ok(())
}
fn try_admit_child(
&mut self,
_parent_token_id: &str,
_child_token_id: String,
_share_bps: u16,
) -> Result<(), BudgetSplitError> {
Ok(())
}
fn verify_child_admission(
&mut self,
_parent_token_id: &str,
_child_token_id: String,
_share_bps: u16,
) -> Result<(), BudgetSplitError> {
Ok(())
}
fn release_child(
&mut self,
_parent_token_id: &str,
_child_token_id: &str,
_expected_share_bps: u16,
) -> Result<(), BudgetSplitError> {
Ok(())
}
fn evict_parent(&mut self, _parent_token_id: &str) {}
}
#[derive(Debug, Default, Clone)]
pub struct InMemoryBudgetRegistry {
splits: BTreeMap<String, BudgetSplit>,
}
impl InMemoryBudgetRegistry {
#[must_use]
pub fn new() -> Self {
Self {
splits: BTreeMap::new(),
}
}
#[must_use]
pub fn split(&self, parent_token_id: &str) -> Option<&BudgetSplit> {
self.splits.get(parent_token_id)
}
}
impl BudgetRegistry for InMemoryBudgetRegistry {
fn register_parent(
&mut self,
parent_token_id: String,
parent_share_bps: u16,
) -> Result<(), BudgetSplitError> {
if parent_share_bps > MAX_BUDGET_SHARE_BPS {
return Err(BudgetSplitError::ChildShareExceedsCap {
child_id: parent_token_id,
share_bps: parent_share_bps,
});
}
if let Some(existing) = self.splits.get(&parent_token_id) {
if existing.parent_share_bps == parent_share_bps {
return Ok(());
}
return Err(BudgetSplitError::DuplicateChild {
child_id: parent_token_id,
});
}
self.splits.insert(
parent_token_id.clone(),
BudgetSplit::new(parent_token_id, parent_share_bps),
);
Ok(())
}
fn try_admit_child(
&mut self,
parent_token_id: &str,
child_token_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError> {
match self.splits.get_mut(parent_token_id) {
Some(split) => split.try_admit_child(child_token_id, share_bps),
None => Err(BudgetSplitError::UnknownParent {
parent_token_id: parent_token_id.into(),
}),
}
}
fn verify_child_admission(
&mut self,
parent_token_id: &str,
child_token_id: String,
share_bps: u16,
) -> Result<(), BudgetSplitError> {
match self.splits.get_mut(parent_token_id) {
Some(split) => split.verify_child_admission(child_token_id, share_bps),
None => Err(BudgetSplitError::UnknownParent {
parent_token_id: parent_token_id.into(),
}),
}
}
fn release_child(
&mut self,
parent_token_id: &str,
child_token_id: &str,
expected_share_bps: u16,
) -> Result<(), BudgetSplitError> {
match self.splits.get_mut(parent_token_id) {
Some(split) => split.release_child(child_token_id, expected_share_bps),
None => Ok(()),
}
}
fn evict_parent(&mut self, parent_token_id: &str) {
let _ = self.splits.remove(parent_token_id);
}
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::string::ToString;
#[test]
fn empty_split_total_is_zero() {
let split = BudgetSplit::new("parent".to_string(), 5_000);
assert_eq!(split.current_total_child_bps(), 0);
}
#[test]
fn admit_first_child_under_parent_succeeds() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("first child fits");
assert_eq!(split.current_total_child_bps(), 4_000);
}
#[test]
fn second_oversubscribed_sibling_is_rejected() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("first child fits");
let err = split
.try_admit_child("child-b".to_string(), 4_000)
.expect_err("second child must oversubscribe");
match err {
BudgetSplitError::OversubscribedSiblings {
child_id,
share_bps,
current_total_child_bps,
parent_share_bps,
} => {
assert_eq!(child_id, "child-b");
assert_eq!(share_bps, 4_000);
assert_eq!(current_total_child_bps, 4_000);
assert_eq!(parent_share_bps, 5_000);
}
other => panic!("unexpected error: {other:?}"),
}
assert_eq!(split.current_total_child_bps(), 4_000);
}
#[test]
fn child_share_exceeding_cap_is_rejected() {
let mut split = BudgetSplit::new("parent".to_string(), MAX_BUDGET_SHARE_BPS);
let err = split
.try_admit_child("child-a".to_string(), MAX_BUDGET_SHARE_BPS + 1)
.expect_err("over-cap share must fail");
assert!(matches!(err, BudgetSplitError::ChildShareExceedsCap { .. }));
}
#[test]
fn duplicate_child_with_same_share_is_idempotent() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("first admit");
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("idempotent re-admit");
assert_eq!(split.current_total_child_bps(), 4_000);
}
#[test]
fn duplicate_child_with_different_share_fails() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 2_000)
.expect("first admit");
let err = split
.try_admit_child("child-a".to_string(), 3_000)
.expect_err("different share must fail");
assert!(matches!(err, BudgetSplitError::DuplicateChild { .. }));
assert_eq!(split.current_total_child_bps(), 2_000);
}
#[test]
fn release_child_restores_parent_headroom() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("child admits");
assert_eq!(split.current_total_child_bps(), 4_000);
split
.release_child("child-a", 4_000)
.expect("matching release succeeds");
assert_eq!(split.current_total_child_bps(), 0);
split
.try_admit_child("child-b".to_string(), 5_000)
.expect("released child should restore full headroom");
}
#[test]
fn release_child_rejects_share_mismatch() {
let mut split = BudgetSplit::new("parent".to_string(), 5_000);
split
.try_admit_child("child-a".to_string(), 4_000)
.expect("child admits");
let error = split
.release_child("child-a", 3_000)
.expect_err("mismatched release must fail");
assert!(matches!(
error,
BudgetSplitError::ReleaseShareMismatch { .. }
));
assert_eq!(split.current_total_child_bps(), 4_000);
}
#[test]
fn in_memory_registry_rejects_unknown_parent() {
let mut registry = InMemoryBudgetRegistry::new();
let err = registry
.try_admit_child("missing", "child".to_string(), 1_000)
.expect_err("unknown parent must fail closed");
assert!(matches!(err, BudgetSplitError::UnknownParent { .. }));
assert!(registry.split("missing").is_none());
}
#[test]
fn registered_parent_enforces_sibling_sum_across_calls() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("register parent snapshot");
registry
.try_admit_child("p", "child-a".to_string(), 6_000)
.expect("first child admits under registered parent");
let err = registry
.try_admit_child("p", "child-b".to_string(), 5_000)
.expect_err("second sibling must oversubscribe (6_000 + 5_000 > 10_000)");
assert!(matches!(
err,
BudgetSplitError::OversubscribedSiblings { .. }
));
}
#[test]
fn in_memory_registry_release_is_idempotent_and_restores_headroom() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), 5_000)
.expect("register parent snapshot");
registry
.try_admit_child("p", "child-a".to_string(), 4_000)
.expect("child admits");
registry
.release_child("p", "child-a", 4_000)
.expect("matching release succeeds");
registry
.release_child("p", "child-a", 4_000)
.expect("missing child release is idempotent");
registry
.release_child("missing-parent", "child-a", 4_000)
.expect("missing parent release is idempotent");
registry
.try_admit_child("p", "child-b".to_string(), 5_000)
.expect("released child should restore full headroom");
}
#[test]
fn overlapping_holders_release_frees_only_on_last() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("register parent");
registry
.try_admit_child("p", "child".to_string(), 6_000)
.expect("A admits the child");
registry
.try_admit_child("p", "child".to_string(), 6_000)
.expect("B re-admits the same child");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
Some(2),
"both overlapping evaluations must hold the single child edge"
);
registry
.release_child("p", "child", 6_000)
.expect("A releases its lease");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
Some(1),
"A's cleanup must NOT free the edge B still holds"
);
let denied = registry
.try_admit_child("p", "sibling".to_string(), 5_000)
.expect_err("an oversubscribing sibling must stay denied while B holds the child");
assert!(matches!(
denied,
BudgetSplitError::OversubscribedSiblings { .. }
));
registry
.release_child("p", "child", 6_000)
.expect("B releases the last lease");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
None,
"the last release must free the edge and return the share to the parent"
);
registry
.try_admit_child("p", "sibling".to_string(), 5_000)
.expect("the sibling admits once both holders have released");
}
#[test]
fn single_holder_release_frees_edge_no_leak() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("register parent");
registry
.try_admit_child("p", "child".to_string(), 6_000)
.expect("child admits");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
Some(1)
);
registry
.release_child("p", "child", 6_000)
.expect("single holder release frees the edge");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
None,
"the sole holder's release must free the edge (no leak)"
);
registry
.try_admit_child("p", "sibling".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("full parent headroom is restored after the release");
}
#[test]
fn verify_only_readmit_does_not_change_holders() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("register parent");
registry
.try_admit_child("p", "child".to_string(), 6_000)
.expect("real dispatch leases the child");
for _ in 0..3 {
registry
.verify_child_admission("p", "child".to_string(), 6_000)
.expect("verify-only re-admit is admissible");
}
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
Some(1),
"verifier-only re-admits must leave holders at the real-holder count"
);
registry
.release_child("p", "child", 6_000)
.expect("the real holder releases its only lease");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
None,
"the last REAL holder's release must free the edge (verify-only did not leak)"
);
registry
.try_admit_child("p", "sibling".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("full headroom restored: no verify-only lease was leaked");
}
#[test]
fn verify_only_fresh_admit_commits_share_without_a_holder() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), 5_000)
.expect("register parent");
registry
.verify_child_admission("p", "child-a".to_string(), 4_000)
.expect("fresh verify-only admit is admissible");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_share_bps("child-a")),
Some(4_000),
"a fresh verify-only admit must commit the child's share"
);
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child-a")),
Some(0),
"a verify-only commit records no releasable holder"
);
let denied = registry
.verify_child_admission("p", "child-b".to_string(), 4_000)
.expect_err("4_000 + 4_000 > 5_000 parent share must be denied");
assert!(matches!(
denied,
BudgetSplitError::OversubscribedSiblings { .. }
));
}
#[test]
fn verify_only_commit_is_superseded_by_real_lease_lifecycle() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("register parent");
registry
.verify_child_admission("p", "child".to_string(), 6_000)
.expect("verify-only commits the fresh edge");
registry
.try_admit_child("p", "child".to_string(), 6_000)
.expect("real dispatch leases the committed edge");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
Some(1),
"the real lease is the sole holder on the verify-only-committed edge"
);
registry
.release_child("p", "child", 6_000)
.expect("the real holder releases");
assert_eq!(
registry
.split("p")
.and_then(|split| split.child_holders("child")),
None,
"releasing the last real holder frees the edge (no verify-only pin)"
);
}
#[test]
fn in_memory_registry_evict_is_idempotent() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), 5_000)
.expect("register");
registry.evict_parent("p");
registry.evict_parent("p");
assert!(registry.split("p").is_none());
}
#[test]
fn in_memory_registry_admits_under_registered_parent() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), 5_000)
.expect("register");
registry
.try_admit_child("p", "c1".to_string(), 4_000)
.expect("first child fits");
let err = registry
.try_admit_child("p", "c2".to_string(), 4_000)
.expect_err("second child must oversubscribe");
assert!(matches!(
err,
BudgetSplitError::OversubscribedSiblings { .. }
));
}
#[test]
fn registered_parent_share_takes_precedence_over_default() {
let mut registry = InMemoryBudgetRegistry::new();
registry
.register_parent("p".to_string(), 5_000)
.expect("register");
let err = registry
.try_admit_child("p", "c1".to_string(), 6_000)
.expect_err("child larger than registered parent share must fail");
assert!(matches!(
err,
BudgetSplitError::OversubscribedSiblings { .. }
));
}
#[test]
fn noop_registry_admits_anything() {
let mut registry = NoopBudgetRegistry;
registry
.register_parent("p".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("noop register");
registry
.try_admit_child("p", "c1".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("noop admit 1");
registry
.try_admit_child("p", "c2".to_string(), MAX_BUDGET_SHARE_BPS)
.expect("noop admit 2");
}
#[test]
fn current_total_avoids_u16_overflow() {
let mut split = BudgetSplit {
parent_token_id: "parent".to_string(),
parent_share_bps: MAX_BUDGET_SHARE_BPS,
children: BTreeMap::new(),
};
split.children.insert(
"c1".to_string(),
ChildAdmission {
share_bps: MAX_BUDGET_SHARE_BPS,
holders: 1,
},
);
split.children.insert(
"c2".to_string(),
ChildAdmission {
share_bps: MAX_BUDGET_SHARE_BPS,
holders: 1,
},
);
assert_eq!(
split.current_total_child_bps(),
u32::from(MAX_BUDGET_SHARE_BPS) * 2
);
}
}