|
1 | 1 | use crate::collections::{HashMap, HashSet, VecDeque}; |
2 | | -use crate::tx_graph::{TxAncestors, TxDescendants}; |
| 2 | +use crate::tx_graph::{CanonicalTx, TxAncestors, TxDescendants}; |
3 | 3 | use crate::{Anchor, ChainOracle, TxGraph}; |
4 | 4 | use alloc::boxed::Box; |
5 | 5 | use alloc::collections::BTreeSet; |
@@ -342,3 +342,157 @@ impl<A: Clone> CanonicalReason<A> { |
342 | 342 | } |
343 | 343 | } |
344 | 344 | } |
| 345 | + |
| 346 | +/// Iterator based on the Kahn's Algorithm, that yields transactions in topological spending order |
| 347 | +/// in depth, and properly sorted with level. |
| 348 | +/// |
| 349 | +/// NOTE: Please refer to the Kahn's Algorithm reference: https://dl.acm.org/doi/pdf/10.1145/368996.369025 |
| 350 | +pub(crate) struct TopologicalIterator<'a, A> { |
| 351 | + /// Map of txid to its canonical transaction |
| 352 | + canonical_txs: HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>>, |
| 353 | + |
| 354 | + /// Current level of transactions to process |
| 355 | + current_level: Vec<Txid>, |
| 356 | + /// Next level of transactions to process |
| 357 | + next_level: Vec<Txid>, |
| 358 | + |
| 359 | + /// Adjacency list: parent txid -> list of children txids |
| 360 | + children_map: HashMap<Txid, Vec<Txid>>, |
| 361 | + /// Number of unprocessed parents for each transaction |
| 362 | + parent_count: HashMap<Txid, usize>, |
| 363 | + |
| 364 | + /// Current index in the current level |
| 365 | + current_index: usize, |
| 366 | +} |
| 367 | + |
| 368 | +impl<'a, A: Clone + Anchor> TopologicalIterator<'a, A> { |
| 369 | + /// Constructs [`TopologicalIterator`] from a list of `canonical_txs` (e.g [`CanonicalIter`]), |
| 370 | + /// in order to handle all the graph building internally. |
| 371 | + pub(crate) fn new( |
| 372 | + canonical_txs: impl Iterator<Item = CanonicalTx<'a, Arc<Transaction>, A>>, |
| 373 | + ) -> Self { |
| 374 | + // Build a map from txid to canonical tx for quick lookup |
| 375 | + let mut tx_map: HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>> = HashMap::new(); |
| 376 | + let mut canonical_set: HashSet<Txid> = HashSet::new(); |
| 377 | + |
| 378 | + for canonical_tx in canonical_txs { |
| 379 | + let txid = canonical_tx.tx_node.txid; |
| 380 | + canonical_set.insert(txid); |
| 381 | + tx_map.insert(txid, canonical_tx); |
| 382 | + } |
| 383 | + |
| 384 | + // Build the dependency graph (txid -> parents it depends on) |
| 385 | + let mut dependencies: HashMap<Txid, Vec<Txid>> = HashMap::new(); |
| 386 | + let mut has_parents: HashSet<Txid> = HashSet::new(); |
| 387 | + |
| 388 | + for &txid in canonical_set.iter() { |
| 389 | + let canonical_tx = tx_map.get(&txid).expect("txid must exist in map"); |
| 390 | + let tx = &canonical_tx.tx_node.tx; |
| 391 | + |
| 392 | + // Find all parents (transactions this one depends on) |
| 393 | + let mut parents = Vec::new(); |
| 394 | + if !tx.is_coinbase() { |
| 395 | + for txin in &tx.input { |
| 396 | + let parent_txid = txin.previous_output.txid; |
| 397 | + // Only include if the parent is also canonical |
| 398 | + if canonical_set.contains(&parent_txid) { |
| 399 | + parents.push(parent_txid); |
| 400 | + has_parents.insert(txid); |
| 401 | + } |
| 402 | + } |
| 403 | + } |
| 404 | + |
| 405 | + if !parents.is_empty() { |
| 406 | + dependencies.insert(txid, parents); |
| 407 | + } |
| 408 | + } |
| 409 | + |
| 410 | + // Build adjacency list and parent counts for traversal |
| 411 | + let mut parent_count = HashMap::new(); |
| 412 | + let mut children_map: HashMap<Txid, Vec<Txid>> = HashMap::new(); |
| 413 | + |
| 414 | + for (txid, parents) in &dependencies { |
| 415 | + for parent_txid in parents { |
| 416 | + children_map.entry(*parent_txid).or_default().push(*txid); |
| 417 | + *parent_count.entry(*txid).or_insert(0) += 1; |
| 418 | + } |
| 419 | + } |
| 420 | + |
| 421 | + // Find root transactions (those with no parents in the canonical set) |
| 422 | + let roots: Vec<Txid> = canonical_set |
| 423 | + .iter() |
| 424 | + .filter(|&&txid| !has_parents.contains(&txid)) |
| 425 | + .copied() |
| 426 | + .collect(); |
| 427 | + |
| 428 | + // Sort the initial level |
| 429 | + let mut current_level = roots; |
| 430 | + Self::sort_level_by_chain_position(&mut current_level, &tx_map); |
| 431 | + |
| 432 | + Self { |
| 433 | + canonical_txs: tx_map, |
| 434 | + current_level, |
| 435 | + next_level: Vec::new(), |
| 436 | + children_map, |
| 437 | + parent_count, |
| 438 | + current_index: 0, |
| 439 | + } |
| 440 | + } |
| 441 | + |
| 442 | + /// Sort transactions within a level by their chain position |
| 443 | + /// Confirmed transactions come first (sorted by height), then unconfirmed (sorted by last_seen) |
| 444 | + fn sort_level_by_chain_position( |
| 445 | + level: &mut [Txid], |
| 446 | + canonical_txs: &HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>>, |
| 447 | + ) { |
| 448 | + level.sort_by(|&a_txid, &b_txid| { |
| 449 | + let a_tx = canonical_txs.get(&a_txid).expect("txid must exist"); |
| 450 | + let b_tx = canonical_txs.get(&b_txid).expect("txid must exist"); |
| 451 | + |
| 452 | + a_tx.cmp(b_tx) |
| 453 | + }); |
| 454 | + } |
| 455 | + |
| 456 | + fn advance_to_next_level(&mut self) { |
| 457 | + self.current_level = core::mem::take(&mut self.next_level); |
| 458 | + Self::sort_level_by_chain_position(&mut self.current_level, &self.canonical_txs); |
| 459 | + self.current_index = 0; |
| 460 | + } |
| 461 | +} |
| 462 | + |
| 463 | +impl<'a, A: Clone + Anchor> Iterator for TopologicalIterator<'a, A> { |
| 464 | + type Item = CanonicalTx<'a, Arc<Transaction>, A>; |
| 465 | + |
| 466 | + fn next(&mut self) -> Option<Self::Item> { |
| 467 | + // If we've exhausted the current level, move to next |
| 468 | + if self.current_index >= self.current_level.len() { |
| 469 | + if self.next_level.is_empty() { |
| 470 | + return None; |
| 471 | + } |
| 472 | + self.advance_to_next_level(); |
| 473 | + } |
| 474 | + |
| 475 | + let current_txid = self.current_level[self.current_index]; |
| 476 | + self.current_index += 1; |
| 477 | + |
| 478 | + // If this is the last item in current level, prepare dependents for next level |
| 479 | + if self.current_index == self.current_level.len() { |
| 480 | + // Process all dependents of all transactions in current level |
| 481 | + for &tx in &self.current_level { |
| 482 | + if let Some(children) = self.children_map.get(&tx) { |
| 483 | + for &child in children { |
| 484 | + if let Some(count) = self.parent_count.get_mut(&child) { |
| 485 | + *count -= 1; |
| 486 | + if *count == 0 { |
| 487 | + self.next_level.push(child); |
| 488 | + } |
| 489 | + } |
| 490 | + } |
| 491 | + } |
| 492 | + } |
| 493 | + } |
| 494 | + |
| 495 | + // Return the CanonicalTx for the current txid |
| 496 | + self.canonical_txs.get(¤t_txid).cloned() |
| 497 | + } |
| 498 | +} |
0 commit comments