diff --git a/rust/hg-core/src/discovery.rs b/rust/hg-core/src/discovery.rs --- a/rust/hg-core/src/discovery.rs +++ b/rust/hg-core/src/discovery.rs @@ -17,6 +17,8 @@ use super::{Graph, GraphError, Revision, NULL_REVISION}; use crate::ancestors::MissingAncestors; use crate::dagops; +use std::cmp::Reverse; +use std::collections::BinaryHeap; use std::collections::{HashMap, HashSet, VecDeque}; type Rng = self::rand_pcg::Pcg32; @@ -212,15 +214,49 @@ missing: impl IntoIterator, ) -> Result<(), GraphError> { self.ensure_undecided()?; - let range = dagops::range( - &self.graph, - missing, - self.undecided.as_ref().unwrap().iter().cloned(), - )?; - let undecided_mut = self.undecided.as_mut().unwrap(); - for missrev in range { - self.missing.insert(missrev); - undecided_mut.remove(&missrev); + if let Some(ref children) = self.children_cache { + let mut seen: HashSet = HashSet::new(); + let mut tovisit: VecDeque = + missing.into_iter().map(|r| -r).collect(); + let undecided_mut = self.undecided.as_mut().unwrap(); + loop { + match tovisit.pop_front() { + None => { + break; + } + Some(opp_rev) => { + let rev = -opp_rev; + self.missing.insert(rev); + undecided_mut.remove(&rev); + match children.get(&rev) { + None => { + continue; + } + Some(this_children) => { + for child in this_children.iter().cloned() { + if child != NULL_REVISION + && !seen.contains(&child) + { + tovisit.push_back(-child); + seen.insert(child); + } + } + } + } + } + } + } + } else { + let range = dagops::range( + &self.graph, + missing, + self.undecided.as_ref().unwrap().iter().cloned(), + )?; + let undecided_mut = self.undecided.as_mut().unwrap(); + for missrev in range { + self.missing.insert(missrev); + undecided_mut.remove(&missrev); + } } Ok(()) } @@ -482,6 +518,24 @@ } #[test] + fn test_discovery_with_children_cache() -> Result<(), GraphError> { + let mut disco = full_disco(); + disco.add_common_revisions(vec![11, 12])?; + disco.ensure_children_cache(); + disco.add_missing_revisions(vec![8, 10])?; + assert_eq!(sorted_undecided(&disco), vec![5]); + assert_eq!(sorted_missing(&disco), vec![8, 10, 13]); + assert!(!disco.is_complete()); + + disco.add_common_revisions(vec![5])?; + assert_eq!(sorted_undecided(&disco), vec![]); + assert_eq!(sorted_missing(&disco), vec![8, 10, 13]); + assert!(disco.is_complete()); + assert_eq!(sorted_common_heads(&disco)?, vec![5, 11, 12]); + Ok(()) + } + + #[test] fn test_limit_sample_no_need_to() { let sample = vec![1, 2, 3, 4]; assert_eq!(full_disco().limit_sample(sample, 10), vec![1, 2, 3, 4]);