Skip to main content

okane_core/report/
balance.rs

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/// Error related to [Balance] operations.
12#[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/// Accumulated balance of accounts.
37#[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    /// Converts into the underlying [`HashMap`].
60    pub fn into_map(self) -> HashMap<Account<'ctx>, Amount<'ctx>> {
61        self.accounts
62    }
63
64    /// Constructs sorted vec of account and commodity tuple.
65    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    /// Adds a particular account value, and returns the updated balance.
72    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    /// Adds a particular account value with the specified commodity, and returns the updated balance.
80    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    /// Tries to set the particular account's balance with the specified commodity,
92    /// and returns the delta which should have caused the difference.
93    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    /// Gets the balance of the given account.
124    pub fn get(&self, account: Account<'ctx>) -> Option<&Amount<'ctx>> {
125        self.accounts.get(&account)
126    }
127
128    /// Rounds the balance following the context.
129    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    /// Returns unordered iterator for the account and the amount.
136    pub(crate) fn iter(&self) -> impl Iterator<Item = (&Account<'ctx>, &Amount<'ctx>)> {
137        self.accounts.iter()
138    }
139
140    /// Returns `true` when no account has any amount recorded.
141    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        // Note it won't be PostingAmount::zero(),
283        // as set_partial is called with commodity amount.
284        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}