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#[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 Self {
48 accounts: iter.into_iter().collect(),
49 }
50 }
51}
52
53impl<'ctx> Balance<'ctx> {
54 pub fn from_map(values: HashMap<Account<'ctx>, Amount<'ctx>>) -> Self {
56 Self { accounts: values }
57 }
58
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 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 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}