bee_ledger/types/
balance_diff.rs

1// Copyright 2020-2021 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::types::Error;
5
6use bee_message::{
7    address::Address,
8    constants::IOTA_SUPPLY,
9    output::{Output, DUST_THRESHOLD},
10};
11
12use std::collections::{
13    hash_map::{IntoIter, Iter, IterMut},
14    HashMap,
15};
16
17/// Records a balance difference to apply to an address.
18#[derive(Clone, Debug, Default)]
19pub struct BalanceDiff {
20    amount: i64,
21    dust_allowance: i64,
22    dust_outputs: i64,
23}
24
25impl BalanceDiff {
26    /// Creates a new `BalanceDiff`.
27    pub fn new(amount: i64, dust_allowance: i64, dust_outputs: i64) -> Result<Self, Error> {
28        if amount.abs() as u64 > IOTA_SUPPLY {
29            Err(Error::InvalidBalanceDiff(amount))
30        } else if dust_allowance.abs() as u64 > IOTA_SUPPLY {
31            Err(Error::InvalidBalanceDiff(dust_allowance))
32        } else if dust_outputs.abs() as u64 > IOTA_SUPPLY {
33            Err(Error::InvalidBalanceDiff(dust_outputs))
34        } else {
35            Ok(Self {
36                amount,
37                dust_allowance,
38                dust_outputs,
39            })
40        }
41    }
42
43    /// Returns the amount of a `BalanceDiff`.
44    pub fn amount(&self) -> i64 {
45        self.amount
46    }
47
48    /// Returns the dust allowance of a `BalanceDiff`.
49    pub fn dust_allowance(&self) -> i64 {
50        self.dust_allowance
51    }
52
53    /// Returns the number of dust outputs of a `BalanceDiff`.
54    pub fn dust_outputs(&self) -> i64 {
55        self.dust_outputs
56    }
57
58    /// Returns whether dust allowance has been decreased or dust outputs has been increased.
59    pub fn is_dust_mutating(&self) -> bool {
60        self.dust_allowance < 0 || self.dust_outputs > 0
61    }
62
63    /// Negates a `BalanceDiff`.
64    pub fn negate(&mut self) {
65        self.amount = -self.amount;
66        self.dust_allowance = -self.dust_allowance;
67        self.dust_outputs = -self.dust_outputs;
68    }
69}
70
71/// Records a balance differences to apply to addresses.
72#[derive(Clone, Debug, Default)]
73pub struct BalanceDiffs(HashMap<Address, BalanceDiff>);
74
75impl BalanceDiffs {
76    /// Creates a new `BalanceDiffs`.
77    pub fn new() -> Self {
78        Self::default()
79    }
80
81    /// Merges a `BalanceDiffs` into another.
82    pub fn merge(&mut self, other: Self) -> Result<(), Error> {
83        for (address, diff) in other.0 {
84            let e = self.0.entry(address).or_default();
85            e.amount = e
86                .amount
87                .checked_add(diff.amount)
88                .ok_or(Error::BalanceDiffOverflow(e.amount as i128 + diff.amount as i128))?;
89            e.dust_allowance = e
90                .dust_allowance
91                .checked_add(diff.dust_allowance)
92                .ok_or(Error::BalanceDiffOverflow(
93                    e.dust_allowance as i128 + diff.dust_allowance as i128,
94                ))?;
95            e.dust_outputs = e
96                .dust_outputs
97                .checked_add(diff.dust_outputs)
98                .ok_or(Error::BalanceDiffOverflow(
99                    e.dust_outputs as i128 + diff.dust_outputs as i128,
100                ))?;
101        }
102
103        Ok(())
104    }
105
106    /// Gets the `BalanceDiff` of a given address.
107    pub fn get(&self, address: &Address) -> Option<&BalanceDiff> {
108        self.0.get(address)
109    }
110
111    /// Negates a `BalanceDiffs`.
112    pub fn negate(&mut self) {
113        for (_, diff) in self.iter_mut() {
114            diff.negate();
115        }
116    }
117
118    /// Creates a new negated version of a `BalanceDiffs`.
119    pub fn negated(&self) -> Self {
120        let mut new = self.clone();
121        new.negate();
122        new
123    }
124
125    /// Adds an output to a `BalanceDiffs`.
126    pub fn output_add(&mut self, output: &Output) -> Result<(), Error> {
127        match output {
128            Output::SignatureLockedSingle(output) => {
129                self.amount_add(*output.address(), output.amount())?;
130                if output.amount() < DUST_THRESHOLD {
131                    self.dust_outputs_inc(*output.address())?;
132                }
133            }
134            Output::SignatureLockedDustAllowance(output) => {
135                self.amount_add(*output.address(), output.amount())?;
136                self.dust_allowance_add(*output.address(), output.amount())?;
137            }
138            Output::Treasury(_) => return Err(Error::UnsupportedOutputKind(output.kind())),
139        }
140
141        Ok(())
142    }
143
144    /// Subtracts an output from a BalanceDiffs`.
145    pub fn output_sub(&mut self, output: &Output) -> Result<(), Error> {
146        match output {
147            Output::SignatureLockedSingle(output) => {
148                self.amount_sub(*output.address(), output.amount())?;
149                if output.amount() < DUST_THRESHOLD {
150                    self.dust_outputs_dec(*output.address())?;
151                }
152            }
153            Output::SignatureLockedDustAllowance(output) => {
154                self.amount_sub(*output.address(), output.amount())?;
155                self.dust_allowance_sub(*output.address(), output.amount())?;
156            }
157            Output::Treasury(_) => return Err(Error::UnsupportedOutputKind(output.kind())),
158        }
159
160        Ok(())
161    }
162
163    /// Adds a given amount to a given address.
164    pub fn amount_add(&mut self, address: Address, amount: u64) -> Result<(), Error> {
165        let entry = self.0.entry(address).or_default();
166        entry.amount = entry
167            .amount
168            .checked_add(amount as i64)
169            .ok_or(Error::BalanceDiffOverflow(entry.amount as i128 + amount as i128))?;
170        Ok(())
171    }
172
173    /// Subtracts a given amount from a given address.
174    pub fn amount_sub(&mut self, address: Address, amount: u64) -> Result<(), Error> {
175        let entry = self.0.entry(address).or_default();
176        entry.amount = entry
177            .amount
178            .checked_sub(amount as i64)
179            .ok_or(Error::BalanceDiffOverflow(entry.amount as i128 + amount as i128))?;
180        Ok(())
181    }
182
183    /// Adds a given dust allowance to a given address.
184    pub fn dust_allowance_add(&mut self, address: Address, amount: u64) -> Result<(), Error> {
185        let entry = self.0.entry(address).or_default();
186        entry.dust_allowance = entry
187            .dust_allowance
188            .checked_add(amount as i64)
189            .ok_or(Error::BalanceDiffOverflow(
190                entry.dust_allowance as i128 + amount as i128,
191            ))?;
192        Ok(())
193    }
194
195    /// Subtracts a given dust allowance from a given address.
196    pub fn dust_allowance_sub(&mut self, address: Address, amount: u64) -> Result<(), Error> {
197        let entry = self.0.entry(address).or_default();
198        entry.dust_allowance = entry
199            .dust_allowance
200            .checked_sub(amount as i64)
201            .ok_or(Error::BalanceDiffOverflow(
202                entry.dust_allowance as i128 + amount as i128,
203            ))?;
204        Ok(())
205    }
206
207    /// Increments the number of dust outputs of a given address.
208    pub fn dust_outputs_inc(&mut self, address: Address) -> Result<(), Error> {
209        let entry = self.0.entry(address).or_default();
210        entry.dust_outputs = entry
211            .dust_outputs
212            .checked_add(1)
213            .ok_or(Error::BalanceDiffOverflow(entry.dust_outputs as i128 + 1))?;
214        Ok(())
215    }
216
217    /// Decrements the number of dust outputs of a given address.
218    pub fn dust_outputs_dec(&mut self, address: Address) -> Result<(), Error> {
219        let entry = self.0.entry(address).or_default();
220        entry.dust_outputs = entry
221            .dust_outputs
222            .checked_sub(1)
223            .ok_or(Error::BalanceDiffOverflow(entry.dust_outputs as i128 + 1))?;
224        Ok(())
225    }
226
227    /// Creates an iterator over the balance diffs.
228    pub fn iter(&self) -> Iter<'_, Address, BalanceDiff> {
229        self.0.iter()
230    }
231
232    /// Creates a mutable iterator over the balance diffs.
233    pub fn iter_mut(&mut self) -> IterMut<'_, Address, BalanceDiff> {
234        self.0.iter_mut()
235    }
236}
237
238impl IntoIterator for BalanceDiffs {
239    type Item = (Address, BalanceDiff);
240    type IntoIter = IntoIter<Address, BalanceDiff>;
241
242    fn into_iter(self) -> Self::IntoIter {
243        self.0.into_iter()
244    }
245}
bee_ledger/types/balance_diff.rs

bee_ledger/types/
balance_diff.rs
