diff --git a/src/thin/runs.rs b/src/thin/runs.rs index 451e7c7..be8e287 100644 --- a/src/thin/runs.rs +++ b/src/thin/runs.rs @@ -19,6 +19,7 @@ pub struct Gatherer { heads: BTreeSet, tails: BTreeSet, entries: BTreeMap, + back: BTreeMap, } impl Gatherer { @@ -28,6 +29,7 @@ impl Gatherer { heads: BTreeSet::new(), tails: BTreeSet::new(), entries: BTreeMap::new(), + back: BTreeMap::new(), } } @@ -59,6 +61,14 @@ impl Gatherer { if let Some(prev) = self.prev { let e = self.entries.get_mut(&prev).unwrap(); e.neighbours.insert(b); + + if let Some(back) = self.back.get(&b) { + if *back != prev { + self.mark_head(b); + } + } else { + self.back.insert(b, prev); + } } else { self.mark_head(b); } @@ -153,6 +163,8 @@ mod tests { #[test] fn gather() { + // the first value defines the input runs, + // and the second defines expected sub sequences struct Test(Vec>, Vec>); let tests = vec![ @@ -177,6 +189,10 @@ mod tests { ], vec![vec![1], vec![2], vec![3, 4], vec![5, 6]], ), + Test( + vec![vec![2, 3, 4, 5], vec![2, 3, 4, 6], vec![1, 3, 4, 5]], + vec![vec![1], vec![2], vec![3, 4], vec![5], vec![6]], + ), ]; for t in tests {