use petgraph::dot::Dot;
use petgraph::graph::{DiGraph, NodeIndex};
use serde::{Deserialize, Serialize};
use sn_transfers::{NanoTokens, SignedSpend, SpendAddress};
use std::collections::BTreeMap;
use crate::error::{Error, Result};
#[derive(Default, Debug, Clone, Serialize, Deserialize)]
pub struct SpendDag {
dag: DiGraph<SpendAddress, NanoTokens>,
spends: BTreeMap<SpendAddress, Vec<(Option<SignedSpend>, usize)>>,
}
impl SpendDag {
pub fn new() -> Self {
Self {
dag: DiGraph::new(),
spends: BTreeMap::new(),
}
}
pub fn load_from_file(path: &str) -> Result<Self> {
let bytes = std::fs::read(path)?;
let dag: SpendDag = rmp_serde::from_slice(&bytes)?;
Ok(dag)
}
pub fn dump_to_file(&self, path: &str) -> Result<()> {
let bytes = rmp_serde::to_vec(&self)?;
std::fs::write(path, bytes)?;
Ok(())
}
pub fn insert(&mut self, spend_addr: SpendAddress, spend: SignedSpend) {
let entries = self.spends.entry(spend_addr).or_default();
let existing_entry = entries.iter_mut().find(|(s, _idx)| {
match s {
Some(existing_spend) => existing_spend == &spend,
None => true,
}
});
let node_idx = match existing_entry {
Some(entry) => {
*entry = (Some(spend.clone()), entry.1);
NodeIndex::new(entry.1)
}
_ => {
let node_idx = self.dag.add_node(spend_addr);
entries.push((Some(spend.clone()), node_idx.index()));
node_idx
}
};
let spend_amount = spend.token();
for ancestor in spend.spend.parent_tx.inputs.iter() {
let ancestor_addr = SpendAddress::from_unique_pubkey(&ancestor.unique_pubkey);
let spends_at_addr = self.spends.entry(ancestor_addr).or_insert_with(|| {
let node_idx = self.dag.add_node(ancestor_addr);
vec![(None, node_idx.index())]
});
for (_, idx) in spends_at_addr {
let ancestor_idx = NodeIndex::new(*idx);
self.dag.update_edge(ancestor_idx, node_idx, *spend_amount);
}
}
for descendant in spend.spend.spent_tx.outputs.iter() {
let descendant_addr = SpendAddress::from_unique_pubkey(&descendant.unique_pubkey);
let spends_at_addr = self.spends.entry(descendant_addr).or_insert_with(|| {
let node_idx = self.dag.add_node(descendant_addr);
vec![(None, node_idx.index())]
});
for (_, idx) in spends_at_addr {
let descendant_idx = NodeIndex::new(*idx);
self.dag
.update_edge(node_idx, descendant_idx, descendant.amount);
}
}
}
pub fn check_and_insert(
&mut self,
spend_addr: SpendAddress,
spend: SignedSpend,
) -> Result<bool> {
if let Some(existing_spends) = self.spends.get(&spend_addr) {
match existing_spends.as_slice() {
[(Some(existing_spend), _)] if existing_spend == &spend => Ok(false),
[(Some(existing_spend), _)] if existing_spend != &spend => {
self.insert(spend_addr, spend.clone());
Err(Error::DoubleSpend(spend_addr))
}
[(None, _)] => {
self.insert(spend_addr, spend);
Ok(true)
}
_ => Err(Error::DoubleSpend(spend_addr)),
}
} else {
self.insert(spend_addr, spend);
Ok(true)
}
}
pub fn get_utxos(&self) -> Vec<SpendAddress> {
let mut leaves = Vec::new();
for node_index in self.dag.node_indices() {
if !self
.dag
.neighbors_directed(node_index, petgraph::Direction::Outgoing)
.any(|_| true)
{
let utxo_addr = self.dag[node_index];
leaves.push(utxo_addr);
}
}
leaves
}
pub fn dump_dot_format(&self) -> String {
format!("{:?}", Dot::with_config(&self.dag, &[]))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_spend_dag_serialisation() {
let dag = SpendDag::new();
let serialized_data = rmp_serde::to_vec(&dag).expect("Serialization failed");
let deserialized_instance: SpendDag =
rmp_serde::from_slice(&serialized_data).expect("Deserialization failed");
let reserialized_data =
rmp_serde::to_vec(&deserialized_instance).expect("Serialization failed");
assert_eq!(reserialized_data, serialized_data);
}
}