1use std::collections::BTreeMap;
10use std::fmt;
11
12use rustev_contract::Decimal;
13use rustev_contract::run::{Charge, CostBound};
14
15const SCALE: i128 = 1_000_000_000;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum CostError {
20 Syntax(String),
22 Negative,
23 Overflow,
24 Unpriced(String),
26}
27
28impl fmt::Display for CostError {
29 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
30 match self {
31 CostError::Syntax(s) => write!(f, "not a decimal amount: {s}"),
32 CostError::Negative => write!(f, "negative amount or rate"),
33 CostError::Overflow => write!(f, "amount out of range"),
34 CostError::Unpriced(u) => write!(f, "no price for usage `{u}`"),
35 }
36 }
37}
38
39impl std::error::Error for CostError {}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub struct UnitRate(Decimal);
45
46impl UnitRate {
47 pub const ONE: UnitRate = UnitRate(Decimal::from_i64(1));
49
50 pub fn new(rate: Decimal) -> Result<Self, CostError> {
51 if rate.is_negative() {
52 return Err(CostError::Negative);
53 }
54 Ok(UnitRate(rate))
55 }
56
57 pub fn parse(s: &str) -> Result<Self, CostError> {
58 Self::new(Decimal::parse(s).map_err(|e| CostError::Syntax(e.to_string()))?)
59 }
60
61 pub fn decimal(&self) -> Decimal {
62 self.0
63 }
64}
65
66fn ceil_div(n: i128, d: i128) -> i128 {
68 n / d + i128::from(n % d != 0)
69}
70
71fn to_u64(v: i128) -> Result<u64, CostError> {
72 u64::try_from(v).map_err(|_| CostError::Overflow)
73}
74
75pub fn units_from_amount(amount: &str, rate: UnitRate) -> Result<u64, CostError> {
77 let a = Decimal::parse(amount).map_err(|e| CostError::Syntax(e.to_string()))?;
78 if a.is_negative() {
79 return Err(CostError::Negative);
80 }
81 let p = a
82 .units()
83 .checked_mul(rate.0.units())
84 .ok_or(CostError::Overflow)?;
85 to_u64(ceil_div(p, SCALE * SCALE))
86}
87
88pub fn convert_units(units: u64, rate: UnitRate) -> Result<u64, CostError> {
90 let p = i128::from(units)
91 .checked_mul(rate.0.units())
92 .ok_or(CostError::Overflow)?;
93 to_u64(ceil_div(p, SCALE))
94}
95
96pub fn convert_charge(charge: Charge, rate: UnitRate) -> Charge {
99 match charge {
100 Charge::Observed { units } => convert_units(units, rate)
101 .map(|units| Charge::Observed { units })
102 .unwrap_or(Charge::Unknown),
103 Charge::Estimated { units } => convert_units(units, rate)
104 .map(|units| Charge::Estimated { units })
105 .unwrap_or(Charge::Unknown),
106 Charge::Unknown => Charge::Unknown,
107 }
108}
109
110pub fn convert_bound(bound: CostBound, rate: UnitRate) -> CostBound {
112 match bound {
113 CostBound::Bounded { max_units } => convert_units(max_units, rate)
114 .map(|max_units| CostBound::Bounded { max_units })
115 .unwrap_or(CostBound::Unknown),
116 CostBound::Estimated { units } => convert_units(units, rate)
117 .map(|units| CostBound::Estimated { units })
118 .unwrap_or(CostBound::Unknown),
119 CostBound::Unknown => CostBound::Unknown,
120 }
121}
122
123pub fn observed_from_amount(amount: &str, rate: UnitRate) -> Charge {
126 units_from_amount(amount, rate)
127 .map(|units| Charge::Observed { units })
128 .unwrap_or(Charge::Unknown)
129}
130
131#[derive(Debug, Clone, Default, PartialEq, Eq)]
134pub struct PriceTable(pub BTreeMap<String, Decimal>);
135
136impl PriceTable {
137 pub fn estimate(&self, usage: &BTreeMap<String, u64>) -> Charge {
140 match self.estimate_units(usage) {
141 Ok(units) => Charge::Estimated { units },
142 Err(_) => Charge::Unknown,
143 }
144 }
145
146 pub fn estimate_units(&self, usage: &BTreeMap<String, u64>) -> Result<u64, CostError> {
147 let mut total: i128 = 0;
148 for (name, n) in usage {
149 let price = self
150 .0
151 .get(name)
152 .ok_or_else(|| CostError::Unpriced(name.clone()))?;
153 if price.is_negative() {
154 return Err(CostError::Negative);
155 }
156 let p = i128::from(*n)
157 .checked_mul(price.units())
158 .ok_or(CostError::Overflow)?;
159 total = total.checked_add(p).ok_or(CostError::Overflow)?;
160 }
161 to_u64(ceil_div(total, SCALE))
162 }
163}
164
165pub fn attribute(total: Charge, received: &[bool]) -> Vec<Charge> {
172 let n = received.iter().filter(|r| **r).count() as u64;
173 let (units, make): (u64, fn(u64) -> Charge) = match total {
174 Charge::Observed { units } => (units, |units| Charge::Observed { units }),
175 Charge::Estimated { units } => (units, |units| Charge::Estimated { units }),
176 Charge::Unknown => return vec![Charge::Unknown; received.len()],
177 };
178 if n == 0 {
179 return vec![Charge::Unknown; received.len()];
180 }
181 let (base, mut extra) = (units / n, units % n);
182 received
183 .iter()
184 .map(|r| {
185 if !*r {
186 return Charge::Unknown;
187 }
188 let share = base + u64::from(extra > 0);
189 extra = extra.saturating_sub(1);
190 make(share)
191 })
192 .collect()
193}
194
195pub fn sum_charges(charges: &[Charge]) -> Charge {
198 let mut total: u64 = 0;
199 let mut estimated = false;
200 for c in charges {
201 match c {
202 Charge::Observed { units } => total = total.saturating_add(*units),
203 Charge::Estimated { units } => {
204 estimated = true;
205 total = total.saturating_add(*units)
206 }
207 Charge::Unknown => return Charge::Unknown,
208 }
209 }
210 if total == u64::MAX {
211 return Charge::Unknown;
212 }
213 if estimated {
214 Charge::Estimated { units: total }
215 } else {
216 Charge::Observed { units: total }
217 }
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223
224 fn rate(s: &str) -> UnitRate {
225 UnitRate::parse(s).unwrap()
226 }
227
228 #[test]
229 fn amounts_are_converted_and_rounded_up() {
230 assert_eq!(units_from_amount("0.000125", rate("10000")), Ok(2));
232 assert_eq!(units_from_amount("0.0002", rate("10000")), Ok(2));
233 assert_eq!(units_from_amount("0", rate("10000")), Ok(0));
234 assert_eq!(units_from_amount("0.000000001", rate("0.000000001")), Ok(1));
236 assert!(matches!(
237 units_from_amount("-1", UnitRate::ONE),
238 Err(CostError::Negative)
239 ));
240 assert!(matches!(
241 units_from_amount("1e3", UnitRate::ONE),
242 Err(CostError::Syntax(_))
243 ));
244 assert_eq!(observed_from_amount("x", UnitRate::ONE), Charge::Unknown);
245 assert!(UnitRate::parse("-1").is_err());
246 }
247
248 #[test]
249 fn remote_units_convert_by_the_declared_rate() {
250 assert_eq!(
251 convert_charge(Charge::Observed { units: 3 }, rate("1.5")),
252 Charge::Observed { units: 5 }
253 );
254 assert_eq!(
255 convert_charge(Charge::Estimated { units: 3 }, UnitRate::ONE),
256 Charge::Estimated { units: 3 }
257 );
258 assert_eq!(
259 convert_charge(Charge::Observed { units: u64::MAX }, rate("2")),
260 Charge::Unknown
261 );
262 assert_eq!(
263 convert_bound(CostBound::Bounded { max_units: 3 }, rate("0.5")),
264 CostBound::Bounded { max_units: 2 }
265 );
266 }
267
268 #[test]
269 fn usage_is_estimated_from_the_price_table() {
270 let t = PriceTable(BTreeMap::from([
271 ("input".to_string(), Decimal::parse("0.001").unwrap()),
272 ("output".to_string(), Decimal::parse("0.002").unwrap()),
273 ]));
274 let usage = BTreeMap::from([("input".to_string(), 1000), ("output".to_string(), 1)]);
275 assert_eq!(t.estimate(&usage), Charge::Estimated { units: 2 });
276 let unpriced = BTreeMap::from([("cached".to_string(), 1)]);
277 assert_eq!(t.estimate(&unpriced), Charge::Unknown);
278 }
279
280 #[test]
281 fn shares_sum_exactly_to_the_total_in_attempt_id_order() {
282 let got = attribute(Charge::Observed { units: 10 }, &[true, false, true, true]);
283 assert_eq!(
284 got,
285 vec![
286 Charge::Observed { units: 4 },
287 Charge::Unknown,
288 Charge::Observed { units: 3 },
289 Charge::Observed { units: 3 },
290 ]
291 );
292 for total in 0..20u64 {
293 for n in 1..6 {
294 let shares = attribute(Charge::Observed { units: total }, &vec![true; n]);
295 let sum: u64 = shares
296 .iter()
297 .map(|c| match c {
298 Charge::Observed { units } => *units,
299 _ => panic!(),
300 })
301 .sum();
302 assert_eq!(sum, total);
303 }
304 }
305 assert_eq!(
306 attribute(Charge::Unknown, &[true, true]),
307 vec![Charge::Unknown; 2]
308 );
309 assert_eq!(
310 attribute(Charge::Observed { units: 5 }, &[false, false]),
311 vec![Charge::Unknown; 2]
312 );
313 assert_eq!(
314 attribute(Charge::Estimated { units: 3 }, &[true, true]),
315 vec![
316 Charge::Estimated { units: 2 },
317 Charge::Estimated { units: 1 }
318 ]
319 );
320 }
321
322 #[test]
323 fn sums_are_known_only_when_every_part_is() {
324 let o = |u| Charge::Observed { units: u };
325 assert_eq!(sum_charges(&[o(1), o(2)]), o(3));
326 assert_eq!(
327 sum_charges(&[o(1), Charge::Estimated { units: 2 }]),
328 Charge::Estimated { units: 3 }
329 );
330 assert_eq!(sum_charges(&[o(1), Charge::Unknown]), Charge::Unknown);
331 assert_eq!(sum_charges(&[]), o(0));
332 }
333}