use std::fmt;
use std::time::Duration;
use rust_decimal::Decimal;
use rust_decimal::prelude::ToPrimitive;
use serde::{Deserialize, Serialize};
pub const MICROS_PER_USD: u32 = 1_000_000;
pub fn micros_of_usd(usd: Decimal) -> u64 {
(usd * Decimal::from(MICROS_PER_USD))
.round()
.to_u64()
.unwrap_or(0)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Unit {
Tokens,
Usd,
Wall,
Calls,
}
impl Unit {
pub const ALL: [Self; 4] = [Self::Tokens, Self::Usd, Self::Wall, Self::Calls];
pub const fn name(self) -> &'static str {
match self {
Self::Tokens => "tokens",
Self::Usd => "usd",
Self::Wall => "wall",
Self::Calls => "calls",
}
}
pub fn from_name(name: &str) -> Option<Self> {
Self::ALL.into_iter().find(|u| u.name() == name)
}
pub const fn index(self) -> usize {
match self {
Self::Tokens => 0,
Self::Usd => 1,
Self::Wall => 2,
Self::Calls => 3,
}
}
pub fn render(self, amount: u64) -> String {
match self {
Self::Tokens => format!("{amount} tokens"),
Self::Usd => {
let mut usd = (Decimal::from(amount) / Decimal::from(MICROS_PER_USD)).normalize();
if usd.scale() < 2 {
usd.rescale(2);
}
format!("${usd}")
}
Self::Wall => format!("{:.3}s", amount as f64 / 1_000.0),
Self::Calls => format!("{amount} calls"),
}
}
}
impl fmt::Display for Unit {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(self.name())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Excess {
pub unit: Unit,
pub declared: u64,
pub permitted: u64,
}
impl Excess {
pub const fn overage(&self) -> u64 {
self.declared.saturating_sub(self.permitted)
}
}
impl fmt::Display for Excess {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"{} against a limit of {}",
self.unit.render(self.declared),
self.unit.render(self.permitted)
)
}
}
#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct Budget {
pub tokens: Option<u64>,
pub usd: Option<Decimal>,
pub wall: Option<Duration>,
pub calls: Option<u32>,
}
impl Budget {
pub const UNBOUNDED: Self = Self {
tokens: None,
usd: None,
wall: None,
calls: None,
};
#[must_use]
pub fn with_tokens(mut self, tokens: u64) -> Self {
self.tokens = Some(tokens);
self
}
#[must_use]
pub fn with_usd(mut self, usd: Decimal) -> Self {
self.usd = Some(usd.normalize());
self
}
#[must_use]
pub fn with_wall(mut self, wall: Duration) -> Self {
self.wall = Some(wall);
self
}
#[must_use]
pub fn with_calls(mut self, calls: u32) -> Self {
self.calls = Some(calls);
self
}
pub fn is_unbounded(&self) -> bool {
*self == Self::UNBOUNDED
}
pub fn amount(&self, unit: Unit) -> Option<u64> {
match unit {
Unit::Tokens => self.tokens,
Unit::Usd => self.usd.map(micros_of_usd),
Unit::Wall => self
.wall
.map(|w| w.as_millis().min(u128::from(u64::MAX)) as u64),
Unit::Calls => self.calls.map(u64::from),
}
}
pub fn bounded_units(&self) -> Vec<Unit> {
Unit::ALL
.into_iter()
.filter(|u| self.amount(*u).is_some())
.collect()
}
pub fn excesses(&self, parent: &Self) -> Vec<Excess> {
Unit::ALL
.into_iter()
.filter_map(|unit| {
let declared = self.amount(unit)?;
let permitted = parent.amount(unit)?;
(declared > permitted).then_some(Excess {
unit,
declared,
permitted,
})
})
.collect()
}
pub fn fits_within(&self, parent: &Self) -> bool {
self.excesses(parent).is_empty()
}
#[must_use]
pub fn clamped_to(&self, parent: &Self) -> Self {
Self {
tokens: tighter(self.tokens, parent.tokens),
usd: tighter(self.usd, parent.usd),
wall: tighter(self.wall, parent.wall),
calls: tighter(self.calls, parent.calls),
}
}
pub fn declared_units(&self) -> Vec<&'static str> {
let mut units = Vec::with_capacity(4);
if self.tokens.is_some() {
units.push("tokens");
}
if self.usd.is_some() {
units.push("usd");
}
if self.wall.is_some() {
units.push("wall");
}
if self.calls.is_some() {
units.push("calls");
}
units
}
}
fn tighter<T: Ord>(child: Option<T>, parent: Option<T>) -> Option<T> {
match (child, parent) {
(Some(a), Some(b)) => Some(a.min(b)),
(Some(a), None) => Some(a),
(None, b) => b,
}
}
impl fmt::Display for Budget {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if self.is_unbounded() {
return f.write_str("unbounded");
}
let mut first = true;
let mut part = |f: &mut fmt::Formatter<'_>, text: &str| {
let sep = if first { "" } else { ", " };
first = false;
write!(f, "{sep}{text}")
};
if let Some(t) = self.tokens {
part(f, &format!("{t} tokens"))?;
}
if let Some(u) = self.usd {
part(f, &format!("${u}"))?;
}
if let Some(w) = self.wall {
part(f, &format!("{:.3}s", w.as_secs_f64()))?;
}
if let Some(c) = self.calls {
part(f, &format!("{c} calls"))?;
}
Ok(())
}
}
impl fmt::Debug for Budget {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Budget({self})")
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::canonical::Canonical;
#[test]
fn the_default_is_unbounded() {
assert!(Budget::default().is_unbounded());
assert_eq!(Budget::default(), Budget::UNBOUNDED);
assert!(Budget::UNBOUNDED.declared_units().is_empty());
}
#[test]
fn units_are_independent() {
let b = Budget::UNBOUNDED.with_tokens(1_000);
assert_eq!(b.tokens, Some(1_000));
assert_eq!(b.usd, None);
assert_eq!(b.wall, None);
assert_eq!(b.calls, None);
assert_eq!(b.declared_units(), ["tokens"]);
}
#[test]
fn all_four_units_can_be_declared_together() {
let b = Budget::UNBOUNDED
.with_tokens(10)
.with_usd(Decimal::new(41, 2))
.with_wall(Duration::from_millis(1_500))
.with_calls(3);
assert_eq!(b.declared_units(), ["tokens", "usd", "wall", "calls"]);
assert!(!b.is_unbounded());
}
#[test]
fn currency_scale_does_not_change_identity() {
let a = Budget::UNBOUNDED.with_usd("1.50".parse().unwrap());
let b = Budget::UNBOUNDED.with_usd("1.5".parse().unwrap());
assert_eq!(a, b);
assert_eq!(a.digest().unwrap(), b.digest().unwrap());
}
#[test]
fn round_trips_through_canonical_cbor() {
let b = Budget::UNBOUNDED
.with_tokens(64)
.with_usd("0.50".parse().unwrap())
.with_wall(Duration::from_secs(30))
.with_calls(2);
let bytes = crate::to_canonical_cbor(&b).unwrap();
let back: Budget = crate::from_canonical_cbor(&bytes).unwrap();
assert_eq!(b, back);
}
#[test]
fn display_names_every_declared_unit() {
assert_eq!(Budget::UNBOUNDED.to_string(), "unbounded");
let b = Budget::UNBOUNDED
.with_tokens(100)
.with_usd("0.41".parse().unwrap());
assert_eq!(b.to_string(), "100 tokens, $0.41");
}
#[test]
fn unit_names_round_trip() {
for unit in Unit::ALL {
assert_eq!(Unit::from_name(unit.name()), Some(unit));
}
assert_eq!(Unit::from_name("furlongs"), None);
}
#[test]
fn unit_indices_are_stable_and_distinct() {
let indices: Vec<_> = Unit::ALL.iter().map(|u| u.index()).collect();
assert_eq!(indices, [0, 1, 2, 3]);
}
#[test]
fn amounts_render_in_the_unit_a_person_reads() {
assert_eq!(Unit::Tokens.render(1_500), "1500 tokens");
assert_eq!(Unit::Usd.render(500_000), "$0.50");
assert_eq!(Unit::Usd.render(2_000_000), "$2.00");
assert_eq!(
Unit::Usd.render(3),
"$0.000003",
"sub-cent precision survives"
);
assert_eq!(Unit::Wall.render(1_500), "1.500s");
assert_eq!(Unit::Calls.render(3), "3 calls");
}
#[test]
fn currency_normalises_to_whole_micro_dollars() {
let b = Budget::UNBOUNDED.with_usd("0.50".parse().unwrap());
assert_eq!(b.amount(Unit::Usd), Some(500_000));
let sum: Decimal = "0.1".parse::<Decimal>().unwrap() + "0.2".parse::<Decimal>().unwrap();
assert_eq!(
Budget::UNBOUNDED.with_usd(sum).amount(Unit::Usd),
Some(300_000)
);
}
#[test]
fn time_normalises_to_milliseconds() {
let b = Budget::UNBOUNDED.with_wall(Duration::from_millis(1_500));
assert_eq!(b.amount(Unit::Wall), Some(1_500));
}
#[test]
fn an_unbounded_unit_is_not_a_zero_bound() {
let unbounded = Budget::UNBOUNDED;
let zero = Budget::UNBOUNDED.with_tokens(0);
assert_eq!(unbounded.amount(Unit::Tokens), None);
assert_eq!(zero.amount(Unit::Tokens), Some(0));
assert_ne!(unbounded, zero);
}
#[test]
fn bounded_units_reports_only_what_is_declared() {
let b = Budget::UNBOUNDED.with_tokens(1).with_calls(2);
assert_eq!(b.bounded_units(), [Unit::Tokens, Unit::Calls]);
assert!(Budget::UNBOUNDED.bounded_units().is_empty());
}
#[test]
fn a_child_asking_for_more_than_its_parent_does_not_fit() {
let parent = Budget::UNBOUNDED.with_tokens(1_000);
let ok = Budget::UNBOUNDED.with_tokens(400);
let greedy = Budget::UNBOUNDED.with_tokens(1_400);
assert!(ok.fits_within(&parent));
assert!(!greedy.fits_within(&parent));
let excesses = greedy.excesses(&parent);
assert_eq!(excesses.len(), 1);
assert_eq!(excesses[0].unit, Unit::Tokens);
assert_eq!(excesses[0].declared, 1_400);
assert_eq!(excesses[0].permitted, 1_000);
assert_eq!(excesses[0].overage(), 400);
}
#[test]
fn equalling_the_parent_fits() {
let parent = Budget::UNBOUNDED.with_tokens(1_000);
assert!(Budget::UNBOUNDED.with_tokens(1_000).fits_within(&parent));
}
#[test]
fn a_child_may_leave_a_unit_its_parent_bounds_unbounded() {
let parent = Budget::UNBOUNDED.with_tokens(1_000).with_calls(5);
let child = Budget::UNBOUNDED.with_tokens(100);
assert!(child.fits_within(&parent));
assert_eq!(child.clamped_to(&parent).calls, Some(5), "it inherits");
}
#[test]
fn a_child_may_bound_a_unit_its_parent_leaves_open() {
let parent = Budget::UNBOUNDED.with_tokens(1_000);
let child = Budget::UNBOUNDED.with_usd("0.10".parse().unwrap());
assert!(child.fits_within(&parent));
assert_eq!(
child.clamped_to(&parent).usd,
Some("0.10".parse().unwrap()),
"the child's own bound survives"
);
}
#[test]
fn every_exceeded_unit_is_reported() {
let parent = Budget::UNBOUNDED
.with_tokens(100)
.with_usd("0.10".parse().unwrap())
.with_calls(10);
let child = Budget::UNBOUNDED
.with_tokens(200)
.with_usd("0.50".parse().unwrap())
.with_calls(1);
let units: Vec<_> = child.excesses(&parent).iter().map(|e| e.unit).collect();
assert_eq!(
units,
[Unit::Tokens, Unit::Usd],
"in unit order, and calls fits"
);
}
#[test]
fn clamping_takes_the_tighter_bound_in_every_unit() {
let parent = Budget::UNBOUNDED.with_tokens(1_000).with_calls(2);
let child = Budget::UNBOUNDED.with_tokens(4_000).with_calls(1);
let effective = child.clamped_to(&parent);
assert_eq!(effective.tokens, Some(1_000), "the parent's is tighter");
assert_eq!(effective.calls, Some(1), "and here the child's is");
}
#[test]
fn clamping_twice_changes_nothing() {
let parent = Budget::UNBOUNDED.with_tokens(1_000);
let child = Budget::UNBOUNDED.with_tokens(4_000).with_calls(3);
let once = child.clamped_to(&parent);
assert_eq!(once.clamped_to(&parent), once);
}
#[test]
fn clamping_composes_down_a_chain() {
let grandparent = Budget::UNBOUNDED.with_tokens(500);
let parent = Budget::UNBOUNDED.with_tokens(2_000);
let child = Budget::UNBOUNDED.with_tokens(4_000);
let stepwise = child.clamped_to(&parent).clamped_to(&grandparent);
assert_eq!(stepwise.tokens, Some(500));
assert_eq!(stepwise, child.clamped_to(&grandparent));
}
#[test]
fn an_unbounded_parent_constrains_nothing() {
let child = Budget::UNBOUNDED.with_tokens(9_000);
assert!(child.fits_within(&Budget::UNBOUNDED));
assert_eq!(child.clamped_to(&Budget::UNBOUNDED), child);
}
#[test]
fn an_excess_names_both_figures() {
let parent = Budget::UNBOUNDED.with_usd("0.10".parse().unwrap());
let child = Budget::UNBOUNDED.with_usd("0.50".parse().unwrap());
let excess = child.excesses(&parent).pop().unwrap();
let printed = excess.to_string();
assert!(printed.contains("$0.5"), "{printed}");
assert!(printed.contains("$0.1"), "{printed}");
}
}