Skip to main content

okane_core/report/
balance.rs

1use std::collections::HashMap;
2
3use crate::report::eval::EvalError;
4
5use super::{
6    context::Account,
7    eval::{Amount, OwnedEvalError, PostingAmount},
8    ReportContext,
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        Self {
48            accounts: iter.into_iter().collect(),
49        }
50    }
51}
52
53impl<'ctx> Balance<'ctx> {
54    /// Constructs an instance directly from the map.
55    pub fn from_map(values: HashMap<Account<'ctx>, Amount<'ctx>>) -> Self {
56        Self { accounts: values }
57    }
58
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 to_from_map() {
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_map(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 balance_is_empty_returns_true_for_empty() {
187        let balance = Balance::default();
188        assert!(balance.is_empty());
189    }
190
191    #[test]
192    fn balance_is_empty_returns_false_for_non_empty() {
193        let arena = Bump::new();
194        let mut ctx = ReportContext::new(&arena);
195
196        let balance = Balance::from_map(hashmap! {
197            ctx.accounts.ensure("Expenses") =>
198                Amount::from_value(ctx.commodities.ensure("JPY"), Decimal::ZERO),
199        });
200        assert!(!balance.is_empty());
201    }
202
203    #[test]
204    fn balance_gives_zero_amount_when_not_initalized() {
205        let arena = Bump::new();
206        let mut ctx = ReportContext::new(&arena);
207
208        let balance = Balance::default();
209        assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), None);
210    }
211
212    #[test]
213    fn test_balance_increment_adds_value() {
214        let arena = Bump::new();
215        let mut ctx = ReportContext::new(&arena);
216
217        let mut balance = Balance::default();
218        let updated = balance
219            .add_posting_amount(
220                ctx.accounts.ensure("Expenses"),
221                PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(1000)),
222            )
223            .clone();
224
225        assert_eq!(
226            updated,
227            Amount::from_value(ctx.commodities.ensure("JPY"), dec!(1000))
228        );
229        assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), Some(&updated));
230
231        let updated = balance
232            .add_posting_amount(
233                ctx.accounts.ensure("Expenses"),
234                PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(-1000)),
235            )
236            .clone();
237
238        assert_eq!(updated, Amount::zero());
239        assert_eq!(balance.get(ctx.accounts.ensure("Expenses")), Some(&updated));
240    }
241
242    #[test]
243    fn test_balance_set_partial_from_absolute_zero() {
244        let arena = Bump::new();
245        let mut ctx = ReportContext::new(&arena);
246        let mut balance = Balance::default();
247
248        let expenses = ctx.accounts.ensure("Expenses");
249        let jpy = ctx.commodities.insert("JPY").unwrap();
250        let prev = balance
251            .set_partial(&ctx, expenses, PostingAmount::from_value(jpy, dec!(1000)))
252            .unwrap();
253
254        // Note it won't be PostingAmount::zero(),
255        // as set_partial is called with commodity amount.
256        assert_eq!(prev, PostingAmount::from_value(jpy, dec!(0)));
257        assert_eq!(
258            balance.get(expenses),
259            Some(&Amount::from_value(jpy, dec!(1000)))
260        );
261    }
262
263    #[test]
264    fn test_balance_set_partial_hit_same_commodity() {
265        let arena = Bump::new();
266        let mut ctx = ReportContext::new(&arena);
267        let mut balance = Balance::default();
268        let jpy = ctx.commodities.ensure("JPY");
269        balance.add_posting_amount(
270            ctx.accounts.ensure("Expenses"),
271            PostingAmount::from_value(jpy, dec!(1000)),
272        );
273
274        let expenses = ctx.accounts.ensure("Expenses");
275
276        let prev = balance
277            .set_partial(&ctx, expenses, PostingAmount::from_value(jpy, dec!(-1000)))
278            .unwrap();
279
280        assert_eq!(prev, PostingAmount::from_value(jpy, dec!(1000)));
281        assert_eq!(
282            balance.get(ctx.accounts.ensure("Expenses")),
283            Some(&Amount::from_value(jpy, dec!(-1000)))
284        );
285    }
286
287    #[test]
288    fn test_balance_set_partial_multi_commodities() {
289        let arena = Bump::new();
290        let mut ctx = ReportContext::new(&arena);
291        let mut balance = Balance::default();
292        let jpy = ctx.commodities.ensure("JPY");
293        let chf = ctx.commodities.ensure("CHF");
294        balance.add_posting_amount(
295            ctx.accounts.ensure("Expenses"),
296            PostingAmount::from_value(jpy, dec!(1000)),
297        );
298        balance.add_posting_amount(
299            ctx.accounts.ensure("Expenses"),
300            PostingAmount::from_value(chf, dec!(200)),
301        );
302
303        let expenses = ctx.accounts.ensure("Expenses");
304
305        let prev = balance
306            .set_partial(&ctx, expenses, PostingAmount::from_value(chf, dec!(100)))
307            .unwrap();
308
309        assert_eq!(prev, PostingAmount::from_value(chf, dec!(200)));
310        assert_eq!(
311            balance.get(ctx.accounts.ensure("Expenses")),
312            Some(&Amount::from_iter([(jpy, dec!(1000)), (chf, dec!(100)),]))
313        );
314    }
315
316    #[test]
317    fn test_balance_set_partial_zero_on_zero() {
318        let arena = Bump::new();
319        let mut ctx = ReportContext::new(&arena);
320        let mut balance = Balance::default();
321
322        let expenses = ctx.accounts.ensure("Expenses");
323
324        let prev = balance
325            .set_partial(&ctx, expenses, PostingAmount::zero())
326            .unwrap();
327
328        assert_eq!(prev, PostingAmount::zero());
329        assert_eq!(
330            balance.get(ctx.accounts.ensure("Expenses")),
331            Some(&Amount::zero())
332        );
333    }
334
335    #[test]
336    fn test_balance_set_partial_zero_on_single_commodity() {
337        let arena = Bump::new();
338        let mut ctx = ReportContext::new(&arena);
339        let mut balance = Balance::default();
340        let jpy = ctx.commodities.ensure("JPY");
341        balance.add_posting_amount(
342            ctx.accounts.ensure("Expenses"),
343            PostingAmount::from_value(jpy, dec!(1000)),
344        );
345
346        let expenses = ctx.accounts.ensure("Expenses");
347
348        let prev = balance
349            .set_partial(&ctx, expenses, PostingAmount::zero())
350            .unwrap();
351
352        assert_eq!(prev, PostingAmount::from_value(jpy, dec!(1000)));
353        assert_eq!(
354            balance.get(ctx.accounts.ensure("Expenses")),
355            Some(&Amount::zero())
356        );
357    }
358
359    #[test]
360    fn test_balance_set_partial_zero_fails_on_multi_commodities() {
361        let arena = Bump::new();
362        let mut ctx = ReportContext::new(&arena);
363        let mut balance = Balance::default();
364        balance.add_posting_amount(
365            ctx.accounts.ensure("Expenses"),
366            PostingAmount::from_value(ctx.commodities.ensure("JPY"), dec!(1000)),
367        );
368        balance.add_posting_amount(
369            ctx.accounts.ensure("Expenses"),
370            PostingAmount::from_value(ctx.commodities.ensure("CHF"), dec!(200)),
371        );
372
373        let expenses = ctx.accounts.ensure("Expenses");
374
375        let err = balance
376            .set_partial(&ctx, expenses, PostingAmount::zero())
377            .unwrap_err();
378
379        assert_eq!(
380            err,
381            BalanceError::MultiCommodityWithPartialSet(
382                OwnedEvalError::PostingAmountRequired,
383                "(1000 JPY + 200 CHF)".to_string()
384            )
385        );
386    }
387}