diff --git a/Cargo.lock b/Cargo.lock index 832646e6..ae9e6dd1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -383,6 +383,15 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "ordered-float" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7d950ca161dc355eaf28f82b11345ed76c6e1f6eb1f4f4479e0323b9e2fbd0e" +dependencies = [ + "num-traits", +] + [[package]] name = "pkg-config" version = "0.3.30" @@ -600,6 +609,7 @@ dependencies = [ "chrono", "clap", "libc", + "ordered-float", "pkg-config", "rand", "rand_distr", diff --git a/Cargo.toml b/Cargo.toml index 811c9043..b853fb9e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ serde_json = {version = "1.0.120"} bincode = "1.3.1" rand = "0.8.6" rand_distr = "0.4.0" +ordered-float = "5.3.0" [build-dependencies] bindgen = "0.69.4" diff --git a/src/sys/treeseq.rs b/src/sys/treeseq.rs index 0095d59f..499584a1 100644 --- a/src/sys/treeseq.rs +++ b/src/sys/treeseq.rs @@ -220,6 +220,22 @@ impl TreeSequence { current_row: 0, } } + + pub fn breakpoints(&self) -> &[super::newtypes::Position] { + assert!(!self.as_ref().breakpoints.is_null()); + let size = 1 + usize::try_from(self.num_trees()).unwrap(); + + // SAFETY: pointer is not NULL, the size is correct, + // and Position is a simple newtype for f64 + unsafe { + std::slice::from_raw_parts( + self.as_ref() + .breakpoints + .cast::(), + size, + ) + } + } } #[test] diff --git a/src/trees/treeseq.rs b/src/trees/treeseq.rs index dcbbca49..47f7500a 100644 --- a/src/trees/treeseq.rs +++ b/src/trees/treeseq.rs @@ -769,6 +769,11 @@ impl TreeSequence { pub fn into_mut_ptr(self) -> Option> { std::ptr::NonNull::new(self.inner.into_raw()) } + + /// Breakpoint positions, including 0.0 and the sequence length + pub fn breakpoints(&self) -> &[Position] { + self.inner.breakpoints() + } } impl TryFrom for TreeSequence { diff --git a/tests/test_trees.rs b/tests/test_trees.rs index 4cfa03ff..bd049434 100644 --- a/tests/test_trees.rs +++ b/tests/test_trees.rs @@ -1190,3 +1190,26 @@ fn flatten_nodes_from_individuals() { .collect::>(); assert_eq!(nodes, &[0, 1, 2, 3, 4]); } + +#[test] +fn test_breakpoints() { + let ts = treeseq_from_small_table_collection_two_trees(); + let bp = ts.breakpoints(); + assert_eq!(bp.len(), usize::try_from(ts.num_trees()).unwrap() + 1); + let mut tbp = vec![]; + let mut tree_iterator = ts.tree_iterator(0).unwrap(); + while let Some(tree) = tree_iterator.next() { + let l: f64 = tree.interval().0.into(); + let r: f64 = tree.interval().1.into(); + tbp.push(ordered_float::NotNan::new(l).unwrap()); + tbp.push(ordered_float::NotNan::new(r).unwrap()); + } + tbp.sort_unstable(); + tbp.dedup(); + assert_eq!(bp.len(), tbp.len()); + for (i, j) in bp.iter().zip(tbp.iter()) { + let l = f64::from(i); + let r = f64::from(*j); + assert_eq!(l, r) + } +}