use core::cell::RefCell;
use core::hash::BuildHasher as _;
use core::hash::Hash;
use beancount_types::Acc;
use beancount_types::Account;
use beancount_types::Amount;
use beancount_types::Balance;
use beancount_types::Commodity;
use beancount_types::Directive;
use beancount_types::Transaction;
use delegate::delegate;
use hashbrown::hash_map::RawEntryMut;
use hashbrown::HashMap;
use time::Date;
use xxhash_rust::xxh3::Xxh3;
use crate::ImporterProtocol;
pub struct ApplyTransactionHook<I, F>
where
I: ImporterProtocol,
F: FnMut(&mut Transaction),
{
inner: I,
hook: RefCell<F>,
}
impl<I, F> ApplyTransactionHook<I, F>
where
I: ImporterProtocol,
F: FnMut(&mut Transaction),
{
pub fn new(inner: I, hook: F) -> Self {
let hook = RefCell::new(hook);
Self { inner, hook }
}
}
impl<I, F> ImporterProtocol for ApplyTransactionHook<I, F>
where
I: ImporterProtocol,
F: FnMut(&mut Transaction),
{
type Error = I::Error;
delegate! {
to (self.inner) {
fn account(&self, buffer: &[u8]) -> Result<Account, Self::Error>;
fn date(&self, buffer: &[u8]) -> Option<Result<Date, Self::Error>>;
fn filename(&self, buffer: &[u8]) -> Option<Result<String, Self::Error>>;
fn identify(&self, buffer: &[u8]) -> Result<bool, Self::Error>;
fn name(&self) -> &'static str;
fn typetag_deserialize(&self);
}
}
fn extract(
&self,
buffer: &[u8],
existing: &[Directive],
) -> Result<Vec<Directive>, Self::Error> {
let mut directives = self.inner.extract(buffer, existing)?;
let mut hook = self.hook.borrow_mut();
for directive in &mut directives {
let Directive::Transaction(transaction) = directive else {
continue;
};
hook(transaction)
}
Ok(directives)
}
}
pub struct DeduplicateBalances<I, T, F>
where
I: ImporterProtocol,
T: Ord,
F: Fn(&Balance) -> T,
{
inner: I,
key: F,
}
impl<I, T, F> DeduplicateBalances<I, T, F>
where
I: ImporterProtocol,
T: Ord,
F: Fn(&Balance) -> T,
{
pub fn new(inner: I, key: F) -> Self {
Self { inner, key }
}
}
impl<I, T, F> DeduplicateBalances<I, T, F>
where
I: ImporterProtocol,
T: Ord,
F: Fn(&Balance) -> T,
{
fn key(&self, balance: &Balance) -> T {
(self.key)(balance)
}
fn upsert(&self, map: &mut HashMap<StorageKey, Balance>, balance: Balance) {
let Balance {
date,
ref account,
amount: Amount { ref commodity, .. },
..
} = balance;
let query = QueryKey {
date,
account,
commodity,
};
let hash = map.hasher().hash_one(&query);
let entry = map
.raw_entry_mut()
.from_hash(hash, |storage| storage == query);
match entry {
RawEntryMut::Occupied(mut entry) => {
if self.key(entry.get()) < self.key(&balance) {
entry.insert(balance);
}
}
RawEntryMut::Vacant(entry) => {
entry.insert(query.into(), balance);
}
}
}
}
impl<I, T, F> ImporterProtocol for DeduplicateBalances<I, T, F>
where
F: Fn(&Balance) -> T,
I: ImporterProtocol,
T: Ord,
{
type Error = I::Error;
delegate! {
to (self.inner) {
fn account(&self, buffer: &[u8]) -> Result<Account, Self::Error>;
fn date(&self, buffer: &[u8]) -> Option<Result<Date, Self::Error>>;
fn filename(&self, buffer: &[u8]) -> Option<Result<String, Self::Error>>;
fn identify(&self, buffer: &[u8]) -> Result<bool, Self::Error>;
fn name(&self) -> &'static str;
fn typetag_deserialize(&self);
}
}
fn extract(
&self,
buffer: &[u8],
existing: &[Directive],
) -> Result<Vec<Directive>, Self::Error> {
let directives = self.inner.extract(buffer, existing)?;
let mut balances = HashMap::new();
let mut directives: Vec<_> = directives
.into_iter()
.filter_map(|directive| {
if let Directive::Balance(balance) = directive {
self.upsert(&mut balances, balance);
None
} else {
Some(directive)
}
})
.collect();
directives.extend(balances.into_values().map(Directive::from));
Ok(directives)
}
}
pub trait ImporterProtocolExt {
fn deduplicate_balances_by<T, F>(self, key: F) -> DeduplicateBalances<Self, T, F>
where
Self: ImporterProtocol + Sized,
T: Ord,
F: Fn(&Balance) -> T,
{
DeduplicateBalances::new(self, key)
}
fn apply_transaction_hook<F>(self, hook: F) -> ApplyTransactionHook<Self, F>
where
Self: Sized + ImporterProtocol,
F: FnMut(&mut Transaction),
{
ApplyTransactionHook::new(self, hook)
}
}
impl<I> ImporterProtocolExt for I where I: ImporterProtocol + Sized {}
pub fn hash<T: Hash>(data: T) -> u128 {
let mut hasher = Xxh3::new();
data.hash(&mut hasher);
hasher.digest128()
}
#[derive(Debug, Eq, Hash, PartialEq)]
struct StorageKey {
date: Date,
account: Account,
commodity: Commodity,
}
impl From<QueryKey<'_>> for StorageKey {
fn from(query: QueryKey) -> Self {
let QueryKey {
date,
account,
commodity,
} = query;
let account = account.to_owned();
let commodity = *commodity;
Self {
date,
account,
commodity,
}
}
}
impl PartialEq<QueryKey<'_>> for &StorageKey {
fn eq(&self, other: &QueryKey) -> bool {
self.date == other.date
&& self.account == other.account
&& &self.commodity == other.commodity
}
}
#[derive(Debug, Hash)]
struct QueryKey<'q> {
date: Date,
account: &'q Acc,
commodity: &'q Commodity,
}