1use std::collections::HashMap;
2
3use crate::report::eval::EvalError;
4
5use super::{
6 ReportContext,
7 account::Account,
8 eval::{Amount, OwnedEvalError, PostingAmount},
9};
10
11#[derive(Debug, thiserror::Error, PartialEq, Eq)]
13pub enum BalanceError {
14 #[error("balance = 0 should be used on single commodity balance")]
15 MultiCommodityWithPartialSet(#[source] OwnedEvalError, String),
16}
17
18impl BalanceError {
19 pub(super) fn note(&self) -> impl std::fmt::Display + '_ {
20 BalanceErrorNote(self)
21 }
22}
23
24struct BalanceErrorNote<'a>(&'a BalanceError);
25
26impl std::fmt::Display for BalanceErrorNote<'_> {
27 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 match self.0 {
29 BalanceError::MultiCommodityWithPartialSet(_, balance) => {
30 write!(f, "actual: {balance}")
31 }
32 }
33 }
34}
35
36#[derive(Debug, Default, PartialEq, Eq, Clone)]
38pub struct Balance<'ctx> {
39 accounts: HashMap<Account<'ctx>, Amount<'ctx>>,
40}
41
42impl<'ctx> FromIterator<(Account<'ctx>, Amount<'ctx>)> for Balance<'ctx> {
43 fn from_iter<T>(iter: T) -> Self
44 where
45 T: IntoIterator<Item = (Account<'ctx>, Amount<'ctx>)>,
46 {
47 let iter = iter.into_iter();
48 let mut value = Self {
49 accounts: HashMap::with_capacity(iter.size_hint().0),
50 };
51 for (account, amount) in iter {
52 value.add_amount(account, amount);
53 }
54 value
55 }
56}
57
58impl<'ctx> Balance<'ctx> {
59 pub fn into_map(self) -> HashMap<Account<'ctx>, Amount<'ctx>> {
61 self.accounts
62 }
63
64 pub fn into_vec(self) -> Vec<(Account<'ctx>, Amount<'ctx>)> {
66 let mut ret: Vec<(Account<'ctx>, Amount<'ctx>)> = self.accounts.into_iter().collect();
67 ret.sort_unstable_by_key(|(a, _)| a.as_str());
68 ret
69 }
70
71 pub fn add_amount(&mut self, account: Account<'ctx>, amount: Amount<'ctx>) -> &Amount<'ctx> {
73 let curr: &mut Amount = self.accounts.entry(account).or_default();
74 *curr += amount;
75 curr.remove_zero_entries();
76 curr
77 }
78
79 pub(super) fn add_posting_amount(
81 &mut self,
82 account: Account<'ctx>,
83 amount: PostingAmount<'ctx>,
84 ) -> &Amount<'ctx> {
85 let curr: &mut Amount = self.accounts.entry(account).or_default();
86 *curr += amount;
87 curr.remove_zero_entries();
88 curr
89 }
90
91 pub(super) fn set_partial(
94 &mut self,
95 ctx: &ReportContext<'ctx>,
96 account: Account<'ctx>,
97 amount: PostingAmount<'ctx>,
98 ) -> Result<PostingAmount<'ctx>, BalanceError> {
99 match amount {
100 PostingAmount::Zero => {
101 let prev: Amount<'ctx> = self
102 .accounts
103 .insert(account, Amount::zero())
104 .unwrap_or_default();
105 (&prev).try_into().map_err(|e: EvalError<'_>| {
106 BalanceError::MultiCommodityWithPartialSet(
107 e.into_owned(ctx),
108 prev.as_inline_display(ctx).to_string(),
109 )
110 })
111 }
112 PostingAmount::Single(single_amount) => {
113 let prev = self
114 .accounts
115 .entry(account)
116 .or_default()
117 .set_partial(single_amount);
118 Ok(PostingAmount::Single(prev))
119 }
120 }
121 }
122
123 pub fn get(&self, account: Account<'ctx>) -> Option<&Amount<'ctx>> {
125 self.accounts.get(&account)
126 }
127
128 pub fn round(&mut self, ctx: &ReportContext<'ctx>) {
130 for amount in self.accounts.values_mut() {
131 amount.round_mut(ctx);
132 }
133 }
134
135 pub(crate) fn iter(&self) -> impl Iterator<Item = (&Account<'ctx>, &Amount<'ctx>)> {
137 self.accounts.iter()
138 }
139
140 pub fn is_empty(&self) -> bool {
142 self.accounts.is_empty()
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149
150 use bumpalo::Bump;
151 use maplit::hashmap;
152 use pretty_assertions::assert_eq;
153 use rust_decimal::Decimal;
154 use rust_decimal_macros::dec;
155
156 use super::super::context::ReportContext;
157
158 #[test]
159 fn from_iter_to_map_round_trip() {
160 let arena = Bump::new();
161 let mut ctx = ReportContext::new(&arena);
162 let expenses = ctx.accounts.ensure("Expenses");
163
164 let m = hashmap! {
165 expenses =>
166 Amount::from_value(ctx.commodities.ensure("JPY"), dec!(10)),
167 ctx.accounts.ensure("Income") =>
168 Amount::from_iter([
169 (ctx.commodities.ensure("CHF"), dec!(15)),
170 (ctx.commodities.ensure("USD"), dec!(-5)),
171 ]),
172 };
173
174 let b = Balance::from_iter(m.clone());
175 assert_eq!(
176 b.get(expenses),
177 Some(&Amount::from_value(ctx.commodities.ensure("JPY"), dec!(10)))
178 );
179
180 let m2 = b.into_map();
181
182 assert_eq!(m, m2);
183 }
184
185 #[test]
186 fn from_iter_removes_zero_valued_commodity() {
187 let arena = Bump::new();
188 let mut ctx = ReportContext::new(&arena);
189 let expenses = ctx.accounts.ensure("Expenses");
190 let income = ctx.accounts.ensure("Income");
191 let chf = ctx.commodities.ensure("CHF");
192
193 let b = Balance::from_iter(hashmap! {
194 expenses =>
195 Amount::from_value(ctx.commodities.ensure("JPY"), Decimal::ZERO),
196 income =>
197 Amount::from_iter([
198 (chf, dec!(15)),
199 (ctx.commodities.ensure("USD"), dec!(0)),
200 ]),
201 });
202
203 let m2 = b.into_map();
204 assert_eq!(
205 m2,
206 hashmap! {
207 income => Amount::from_value(chf, dec!(15)),
208 expenses => Amount::zero(),
209 }
210 );
211 }
212
213 #[test]
214 fn balance_is_empty_returns_true_for_empty() {
215 let balance = Balance::default();
216 assert!(balance.is_empty());
217 }
218
219 #[test]
220 fn balance_is_empty_returns_false_for_non_empty() {
221 let arena = Bump::new();
222 let mut ctx = ReportContext::new(&arena);
223
224 let balance = Balance::from_iter(hashmap! {
225 ctx.accounts.ensure("Expenses") =>
226 Amount::from_value(ctx.commodities.ensure("JPY"), Decimal::ZERO),
227 });
228 assert!(!balance.is_empty());
229 }
230
231 #[test]
232 fn balance_gives_zero_amount_when_not_initalized() {
233 let arena = Bump::new();
234 let mut ctx = ReportContext::new(&arena);
235
236 let balance = Balance::default();
237 assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), None);
238 }
239
240 #[test]
241 fn test_balance_increment_adds_value() {
242 let arena = Bump::new();
243 let mut ctx = ReportContext::new(&arena);
244
245 let mut balance = Balance::default();
246 let updated = balance
247 .add_posting_amount(
248 ctx.accounts.ensure("Expenses"),
249 PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(1000)),
250 )
251 .clone();
252
253 assert_eq!(
254 updated,
255 Amount::from_value(ctx.commodities.ensure("JPY"), dec!(1000))
256 );
257 assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), Some(&updated));
258
259 let updated = balance
260 .add_posting_amount(
261 ctx.accounts.ensure("Expenses"),
262 PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(-1000)),
263 )
264 .clone();
265
266 assert_eq!(updated, Amount::zero());
267 assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), Some(&updated));
268 }
269
270 #[test]
271 fn test_balance_set_partial_from_absolute_zero() {
272 let arena = Bump::new();
273 let mut ctx = ReportContext::new(&arena);
274 let mut balance = Balance::default();
275
276 let expenses = ctx.accounts.ensure("Expenses");
277 let jpy = ctx.commodities.insert("JPY").unwrap();
278 let prev = balance
279 .set_partial(&ctx, expenses, PostingAmount::from_value(jpy, dec!(1000)))
280 .unwrap();
281
282 assert_eq!(prev, PostingAmount::from_value(jpy, dec!(0)));
285 assert_eq!(
286 balance.get(expenses),
287 Some(&Amount::from_value(jpy, dec!(1000)))
288 );
289 }
290
291 #[test]
292 fn test_balance_set_partial_hit_same_commodity() {
293 let arena = Bump::new();
294 let mut ctx = ReportContext::new(&arena);
295 let mut balance = Balance::default();
296 let jpy = ctx.commodities.ensure("JPY");
297 balance.add_posting_amount(
298 ctx.accounts.ensure("Expenses"),
299 PostingAmount::from_value(jpy, dec!(1000)),
300 );
301
302 let expenses = ctx.accounts.ensure("Expenses");
303
304 let prev = balance
305 .set_partial(&ctx, expenses, PostingAmount::from_value(jpy, dec!(-1000)))
306 .unwrap();
307
308 assert_eq!(prev, PostingAmount::from_value(jpy, dec!(1000)));
309 assert_eq!(
310 balance.get(ctx.accounts.ensure("Expenses")),
311 Some(&Amount::from_value(jpy, dec!(-1000)))
312 );
313 }
314
315 #[test]
316 fn test_balance_set_partial_multi_commodities() {
317 let arena = Bump::new();
318 let mut ctx = ReportContext::new(&arena);
319 let mut balance = Balance::default();
320 let jpy = ctx.commodities.ensure("JPY");
321 let chf = ctx.commodities.ensure("CHF");
322 balance.add_posting_amount(
323 ctx.accounts.ensure("Expenses"),
324 PostingAmount::from_value(jpy, dec!(1000)),
325 );
326 balance.add_posting_amount(
327 ctx.accounts.ensure("Expenses"),
328 PostingAmount::from_value(chf, dec!(200)),
329 );
330
331 let expenses = ctx.accounts.ensure("Expenses");
332
333 let prev = balance
334 .set_partial(&ctx, expenses, PostingAmount::from_value(chf, dec!(100)))
335 .unwrap();
336
337 assert_eq!(prev, PostingAmount::from_value(chf, dec!(200)));
338 assert_eq!(
339 balance.get(ctx.accounts.ensure("Expenses")),
340 Some(&Amount::from_iter([(jpy, dec!(1000)), (chf, dec!(100)),]))
341 );
342 }
343
344 #[test]
345 fn test_balance_set_partial_zero_on_zero() {
346 let arena = Bump::new();
347 let mut ctx = ReportContext::new(&arena);
348 let mut balance = Balance::default();
349
350 let expenses = ctx.accounts.ensure("Expenses");
351
352 let prev = balance
353 .set_partial(&ctx, expenses, PostingAmount::zero())
354 .unwrap();
355
356 assert_eq!(prev, PostingAmount::zero());
357 assert_eq!(
358 balance.get(ctx.accounts.ensure("Expenses")),
359 Some(&Amount::zero())
360 );
361 }
362
363 #[test]
364 fn test_balance_set_partial_zero_on_single_commodity() {
365 let arena = Bump::new();
366 let mut ctx = ReportContext::new(&arena);
367 let mut balance = Balance::default();
368 let jpy = ctx.commodities.ensure("JPY");
369 balance.add_posting_amount(
370 ctx.accounts.ensure("Expenses"),
371 PostingAmount::from_value(jpy, dec!(1000)),
372 );
373
374 let expenses = ctx.accounts.ensure("Expenses");
375
376 let prev = balance
377 .set_partial(&ctx, expenses, PostingAmount::zero())
378 .unwrap();
379
380 assert_eq!(prev, PostingAmount::from_value(jpy, dec!(1000)));
381 assert_eq!(
382 balance.get(ctx.accounts.ensure("Expenses")),
383 Some(&Amount::zero())
384 );
385 }
386
387 #[test]
388 fn test_balance_set_partial_zero_fails_on_multi_commodities() {
389 let arena = Bump::new();
390 let mut ctx = ReportContext::new(&arena);
391 let mut balance = Balance::default();
392 balance.add_posting_amount(
393 ctx.accounts.ensure("Expenses"),
394 PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(1000)),
395 );
396 balance.add_posting_amount(
397 ctx.accounts.ensure("Expenses"),
398 PostingAmount::from_value(ctx.commodities.ensure("CHF"), dec!(200)),
399 );
400
401 let expenses = ctx.accounts.ensure("Expenses");
402
403 let err = balance
404 .set_partial(&ctx, expenses, PostingAmount::zero())
405 .unwrap_err();
406
407 assert_eq!(
408 err,
409 BalanceError::MultiCommodityWithPartialSet(
410 OwnedEvalError::PostingAmountRequired,
411 "(1000 JPY + 200 CHF)".to_string()
412 )
413 );
414 }
415}