|
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,189 @@ impl<A: Clone> CanonicalReason<A> { |
342 | 342 | } |
343 | 343 | } |
344 | 344 | } |
| 345 | + |
| 346 | +/// Iterator that yields transactions in topological order with proper sorting within levels. |
| 347 | +pub(crate) struct TopologicalIterator<'a, A> { |
| 348 | + /// Map of txid to its canonical transaction |
| 349 | + canonical_txs: HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>>, |
| 350 | + |
| 351 | + /// Current level of transactions to process |
| 352 | + current_level: Vec<Txid>, |
| 353 | + /// Next level of transactions to process |
| 354 | + next_level: Vec<Txid>, |
| 355 | + |
| 356 | + /// Adjacency list: parent txid -> list of children txids |
| 357 | + children_map: HashMap<Txid, Vec<Txid>>, |
| 358 | + /// Number of unprocessed parents for each transaction |
| 359 | + parent_count: HashMap<Txid, usize>, |
| 360 | + |
| 361 | + /// Current index in the current level |
| 362 | + current_index: usize, |
| 363 | +} |
| 364 | + |
| 365 | +impl<'a, A: Clone + Anchor> TopologicalIterator<'a, A> { |
| 366 | + pub(crate) fn new(canonical_txs: Vec<CanonicalTx<'a, Arc<Transaction>, A>>) -> Self { |
| 367 | + // Build a map from txid to canonical tx for quick lookup |
| 368 | + let mut tx_map: HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>> = HashMap::new(); |
| 369 | + let mut canonical_set: HashSet<Txid> = HashSet::new(); |
| 370 | + |
| 371 | + for canonical_tx in canonical_txs { |
| 372 | + let txid = canonical_tx.tx_node.txid; |
| 373 | + canonical_set.insert(txid); |
| 374 | + tx_map.insert(txid, canonical_tx); |
| 375 | + } |
| 376 | + |
| 377 | + // Build the dependency graph (txid -> parents it depends on) |
| 378 | + let mut dependencies: HashMap<Txid, Vec<Txid>> = HashMap::new(); |
| 379 | + let mut has_parents: HashSet<Txid> = HashSet::new(); |
| 380 | + |
| 381 | + for &txid in canonical_set.iter() { |
| 382 | + let canonical_tx = tx_map.get(&txid).expect("txid must exist in map"); |
| 383 | + let tx = &canonical_tx.tx_node.tx; |
| 384 | + |
| 385 | + // Find all parents (transactions this one depends on) |
| 386 | + let mut parents = Vec::new(); |
| 387 | + if !tx.is_coinbase() { |
| 388 | + for txin in &tx.input { |
| 389 | + let parent_txid = txin.previous_output.txid; |
| 390 | + // Only include if the parent is also canonical |
| 391 | + if canonical_set.contains(&parent_txid) { |
| 392 | + parents.push(parent_txid); |
| 393 | + has_parents.insert(txid); |
| 394 | + } |
| 395 | + } |
| 396 | + } |
| 397 | + |
| 398 | + if !parents.is_empty() { |
| 399 | + dependencies.insert(txid, parents); |
| 400 | + } |
| 401 | + } |
| 402 | + |
| 403 | + // Build adjacency list and parent counts for traversal |
| 404 | + let mut parent_count = HashMap::new(); |
| 405 | + let mut children_map: HashMap<Txid, Vec<Txid>> = HashMap::new(); |
| 406 | + |
| 407 | + for (txid, parents) in &dependencies { |
| 408 | + for parent_txid in parents { |
| 409 | + children_map.entry(*parent_txid).or_default().push(*txid); |
| 410 | + *parent_count.entry(*txid).or_insert(0) += 1; |
| 411 | + } |
| 412 | + } |
| 413 | + |
| 414 | + // Find root transactions (those with no parents in the canonical set) |
| 415 | + let roots: Vec<Txid> = canonical_set |
| 416 | + .iter() |
| 417 | + .filter(|&&txid| !has_parents.contains(&txid)) |
| 418 | + .copied() |
| 419 | + .collect(); |
| 420 | + |
| 421 | + // Sort the initial level |
| 422 | + let mut current_level = roots; |
| 423 | + Self::sort_level_by_chain_position(&mut current_level, &tx_map); |
| 424 | + |
| 425 | + Self { |
| 426 | + canonical_txs: tx_map, |
| 427 | + current_level, |
| 428 | + next_level: Vec::new(), |
| 429 | + children_map, |
| 430 | + parent_count, |
| 431 | + current_index: 0, |
| 432 | + } |
| 433 | + } |
| 434 | + |
| 435 | + /// Sort transactions within a level by their chain position |
| 436 | + /// Confirmed transactions come first (sorted by height), then unconfirmed (sorted by last_seen) |
| 437 | + fn sort_level_by_chain_position( |
| 438 | + level: &mut [Txid], |
| 439 | + canonical_txs: &HashMap<Txid, CanonicalTx<'a, Arc<Transaction>, A>>, |
| 440 | + ) { |
| 441 | + level.sort_by(|&a_txid, &b_txid| { |
| 442 | + let a_tx = canonical_txs.get(&a_txid).expect("txid must exist"); |
| 443 | + let b_tx = canonical_txs.get(&b_txid).expect("txid must exist"); |
| 444 | + |
| 445 | + use crate::ChainPosition; |
| 446 | + use core::cmp::Ordering; |
| 447 | + |
| 448 | + match (&a_tx.chain_position, &b_tx.chain_position) { |
| 449 | + // Both confirmed: sort by confirmation height |
| 450 | + ( |
| 451 | + ChainPosition::Confirmed { |
| 452 | + anchor: a_anchor, .. |
| 453 | + }, |
| 454 | + ChainPosition::Confirmed { |
| 455 | + anchor: b_anchor, .. |
| 456 | + }, |
| 457 | + ) => { |
| 458 | + let a_height = a_anchor.confirmation_height_upper_bound(); |
| 459 | + let b_height = b_anchor.confirmation_height_upper_bound(); |
| 460 | + a_height.cmp(&b_height) |
| 461 | + } |
| 462 | + // Confirmed comes before unconfirmed |
| 463 | + (ChainPosition::Confirmed { .. }, ChainPosition::Unconfirmed { .. }) => { |
| 464 | + Ordering::Less |
| 465 | + } |
| 466 | + // Unconfirmed comes after confirmed |
| 467 | + (ChainPosition::Unconfirmed { .. }, ChainPosition::Confirmed { .. }) => { |
| 468 | + Ordering::Greater |
| 469 | + } |
| 470 | + // Both unconfirmed: sort by first_seen (earlier timestamp first) |
| 471 | + ( |
| 472 | + ChainPosition::Unconfirmed { |
| 473 | + first_seen: a_first_seen, |
| 474 | + .. |
| 475 | + }, |
| 476 | + ChainPosition::Unconfirmed { |
| 477 | + first_seen: b_first_seen, |
| 478 | + .. |
| 479 | + }, |
| 480 | + ) => { |
| 481 | + // Earlier timestamps come first |
| 482 | + a_first_seen.cmp(b_first_seen) |
| 483 | + } |
| 484 | + } |
| 485 | + }); |
| 486 | + } |
| 487 | + |
| 488 | + fn advance_to_next_level(&mut self) { |
| 489 | + self.current_level = core::mem::take(&mut self.next_level); |
| 490 | + Self::sort_level_by_chain_position(&mut self.current_level, &self.canonical_txs); |
| 491 | + self.current_index = 0; |
| 492 | + } |
| 493 | +} |
| 494 | + |
| 495 | +impl<'a, A: Clone + Anchor> Iterator for TopologicalIterator<'a, A> { |
| 496 | + type Item = CanonicalTx<'a, Arc<Transaction>, A>; |
| 497 | + |
| 498 | + fn next(&mut self) -> Option<Self::Item> { |
| 499 | + // If we've exhausted the current level, move to next |
| 500 | + if self.current_index >= self.current_level.len() { |
| 501 | + if self.next_level.is_empty() { |
| 502 | + return None; |
| 503 | + } |
| 504 | + self.advance_to_next_level(); |
| 505 | + } |
| 506 | + |
| 507 | + let current_txid = self.current_level[self.current_index]; |
| 508 | + self.current_index += 1; |
| 509 | + |
| 510 | + // If this is the last item in current level, prepare dependents for next level |
| 511 | + if self.current_index == self.current_level.len() { |
| 512 | + // Process all dependents of all transactions in current level |
| 513 | + for &tx in &self.current_level { |
| 514 | + if let Some(children) = self.children_map.get(&tx) { |
| 515 | + for &child in children { |
| 516 | + if let Some(count) = self.parent_count.get_mut(&child) { |
| 517 | + *count -= 1; |
| 518 | + if *count == 0 { |
| 519 | + self.next_level.push(child); |
| 520 | + } |
| 521 | + } |
| 522 | + } |
| 523 | + } |
| 524 | + } |
| 525 | + } |
| 526 | + |
| 527 | + // Return the CanonicalTx for the current txid |
| 528 | + self.canonical_txs.get(¤t_txid).cloned() |
| 529 | + } |
| 530 | +} |
0 commit comments