krabka_verified/
throttle.rs

1//! Pure token-bucket consume arithmetic.
2
3use creusot_std::prelude::*;
4#[cfg(not(creusot))]
5use derive_more::{Display, From, Into};
6
7macro_rules! token_count {
8    ($name:ident, $doc:literal) => {
9        #[doc = $doc]
10        #[krabka_macros::model_type(
11            proof(DeepModel),
12            runtime(
13                Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Display, From, Into
14            )
15        )]
16        pub struct $name(pub u64);
17    };
18}
19
20token_count!(
21    AvailableTokens,
22    "Tokens currently sitting in the bucket, available to grant."
23);
24token_count!(
25    RefillTokens,
26    "Tokens accrued since the last refill, to be added to `available` this call."
27);
28token_count!(
29    BurstCapacity,
30    "The burst cap: the maximum the bucket may hold after a refill."
31);
32token_count!(RequestedTokens, "Tokens the caller is asking to consume.");
33token_count!(
34    GrantedTokens,
35    "Tokens actually granted by a consume call (`<= requested`)."
36);
37token_count!(
38    NewAvailable,
39    "The bucket's new `available` count after a consume call commits."
40);
41
42// cargo-mutants: #[cfg(creusot)] spec function; not compiled outside Creusot, so no test can tell.
43#[cfg(creusot)]
44#[cfg_attr(test, mutants::skip)]
45#[logic(open)]
46/// `min(available + refill, burst)` in unbounded integers.
47///
48/// This equals the executable `available.saturating_add(refill).min(burst)`
49/// for every `u64` input. Without overflow the saturating sum is the exact
50/// sum. With overflow the exact sum exceeds `u64::MAX`, which is at least
51/// `burst`, so both sides are `burst`. [`plan_consume`] proves this equality
52/// for its own inputs.
53pub fn capped(available: Int, refill: Int, burst: Int) -> Int {
54    if available + refill <= burst {
55        available + refill
56    } else {
57        burst
58    }
59}
60
61/// Caps a refill, grants at most the request, and returns the new balance.
62#[ensures(result.0.0@ <= requested.0@)]
63#[ensures(result.1.0@ <= burst.0@)]
64#[ensures(result.0.0@ + result.1.0@ == capped(available.0@, refill.0@, burst.0@))]
65#[ensures(result.0.0@ == if requested.0@ <= capped(available.0@, refill.0@, burst.0@) {
66    requested.0@
67} else {
68    capped(available.0@, refill.0@, burst.0@)
69})]
70#[must_use]
71pub fn plan_consume(
72    available: AvailableTokens,
73    refill: RefillTokens,
74    burst: BurstCapacity,
75    requested: RequestedTokens,
76) -> (GrantedTokens, NewAvailable) {
77    let total = available.0.saturating_add(refill.0).min(burst.0);
78    proof_assert!(total@ == capped(available.0@, refill.0@, burst.0@));
79    let grant = requested.0.min(total);
80    (GrantedTokens(grant), NewAvailable(total - grant))
81}
82
83#[cfg(test)]
84mod tests {
85    use {super::*, proptest::prelude::*};
86
87    #[yare::parameterized(
88        available = { 100, 0, 1000, 50, 50, 50 },
89        request_capped_by_available = { 100, 0, 1000, 200, 100, 0 },
90        request_after_refill = { 900, 500, 1000, 200, 200, 800 },
91        empty_bucket = { 0, 0, 1000, 100, 0, 0 },
92        saturating_refill = { u64::MAX, u64::MAX, 1000, 1000, 1000, 0 },
93    )]
94    fn plan_consume_grants_and_caps(
95        available: u64,
96        refill: u64,
97        burst: u64,
98        requested: u64,
99        grant: u64,
100        new_available: u64,
101    ) {
102        assert2::assert!(
103            plan_consume(
104                AvailableTokens(available),
105                RefillTokens(refill),
106                BurstCapacity(burst),
107                RequestedTokens(requested)
108            ) == (GrantedTokens(grant), NewAvailable(new_available))
109        );
110    }
111
112    proptest! {
113        #[test]
114        fn plan_consume_invariants(
115            available in 0u64..=u64::MAX,
116            refill in 0u64..=u64::MAX,
117            burst in 0u64..1_000_000,
118            requested in 0u64..=u64::MAX,
119        ) {
120            let (grant, new) = plan_consume(
121                AvailableTokens(available),
122                RefillTokens(refill),
123                BurstCapacity(burst),
124                RequestedTokens(requested),
125            );
126            let capped = available.saturating_add(refill).min(burst);
127            prop_assert!(grant.0 <= requested);
128            prop_assert!(grant.0 <= capped);
129            prop_assert_eq!(new.0, capped - grant.0);
130            prop_assert!(new.0 <= burst, "burst cap");
131        }
132    }
133}
krabka_verified/throttle.rs

krabka_verified/
throttle.rs
