use crate::error::TskitError;
use crate::sys;
use crate::EdgeTable;
use crate::IndividualTable;
use crate::MigrationTable;
use crate::MutationTable;
use crate::NodeId;
use crate::NodeTable;
use crate::PopulationTable;
use crate::Position;
use crate::SimplificationOptions;
use crate::SiteTable;
use crate::SizeType;
use crate::TableCollection;
use crate::TableOutputOptions;
use crate::TreeFlags;
use crate::TreeSequenceFlags;
use crate::TskReturnValue;
use sys::bindings as ll_bindings;
#[cfg(feature = "provenance")]
use std::ffi::c_char;
use super::Tree;
pub struct TreeSequence {
pub(crate) inner: sys::TreeSequence,
tables: crate::TableCollection,
}
unsafe impl Send for TreeSequence {}
unsafe impl Sync for TreeSequence {}
impl TreeSequence {
pub fn new<F: Into<TreeSequenceFlags>>(
tables: TableCollection,
flags: F,
) -> Result<Self, TskitError> {
let raw_tables_ptr = tables.into_inner();
let mut inner = sys::TreeSequence::new(raw_tables_ptr, flags.into())?;
let tables = unsafe {
TableCollection::new_from_ll(sys::TableCollection::new_borrowed(
std::ptr::NonNull::new(inner.as_mut().tables).unwrap(),
))
}?;
Ok(Self { inner, tables })
}
fn as_ref(&self) -> &ll_bindings::tsk_treeseq_t {
self.inner.as_ref()
}
pub fn as_ptr(&self) -> *const ll_bindings::tsk_treeseq_t {
self.inner.as_ref()
}
pub fn as_mut_ptr(&mut self) -> *mut ll_bindings::tsk_treeseq_t {
self.inner.as_mut()
}
pub fn dump<O: Into<TableOutputOptions>>(&self, filename: &str, options: O) -> TskReturnValue {
let c_str = std::ffi::CString::new(filename).map_err(|_| {
TskitError::LibraryError("call to ffi::Cstring::new failed".to_string())
})?;
self.inner.dump(c_str, options.into().bits())
}
pub fn load(filename: impl AsRef<str>) -> Result<Self, TskitError> {
let tables = TableCollection::new_from_file(filename.as_ref())?;
Self::new(tables, TreeSequenceFlags::default())
}
pub fn dump_tables(self) -> Result<TableCollection, TskitError> {
assert!(!self.as_ptr().is_null());
let mut treeseq = self;
let tables = std::ptr::NonNull::new(unsafe { (*treeseq.as_ptr()).tables }).unwrap();
unsafe { (*treeseq.as_mut_ptr()).tables = std::ptr::null_mut() };
let tables = unsafe { crate::sys::TableCollection::new_owning_from_nonnull(tables) };
crate::TableCollection::new_from_ll(tables)
}
pub fn tree_iterator<'ts, F: Into<TreeFlags>>(
&'ts self,
flags: F,
) -> Result<Tree<'ts>, TskitError> {
let tree = Tree::new(&self.inner, flags)?;
Ok(tree)
}
pub fn tree_iterator_at_position<'ts, F: Into<TreeFlags>, P: Into<Position>>(
&'ts self,
flags: F,
at: P,
) -> Result<Tree<'ts>, TskitError> {
Tree::new_at_position(&self.inner, flags, at)
}
pub fn tree_iterator_at_index<'ts, F: Into<TreeFlags>>(
&'ts self,
flags: F,
at: i32,
) -> Result<Tree<'ts>, TskitError> {
Tree::new_at_index(&self.inner, flags, at)
}
pub fn sample_nodes(&self) -> &[NodeId] {
unsafe {
let num_samples = ll_bindings::tsk_treeseq_get_num_samples(self.as_ref());
sys::generate_slice(self.as_ref().samples, num_samples)
}
}
pub fn num_trees(&self) -> SizeType {
self.inner.num_trees()
}
pub fn kc_distance(&self, other: &TreeSequence, lambda: f64) -> Result<f64, TskitError> {
self.inner.kc_distance(&other.inner, lambda)
}
pub fn num_samples(&self) -> SizeType {
self.inner.num_samples()
}
pub fn simplify<O: Into<SimplificationOptions>>(
&self,
samples: &[NodeId],
options: O,
idmap: bool,
) -> Result<(Self, Option<Vec<NodeId>>), TskitError> {
let mut output_node_map: Vec<NodeId> = vec![];
if idmap {
output_node_map.resize(usize::try_from(self.nodes().num_rows())?, NodeId::NULL);
}
let mut inner = self.inner.simplify(
samples,
options.into(),
match idmap {
true => Some(&mut output_node_map),
false => None,
},
)?;
let tables = unsafe {
TableCollection::new_from_ll(sys::TableCollection::new_borrowed(
std::ptr::NonNull::new(inner.as_mut().tables).unwrap(),
))
}?;
Ok((
Self { inner, tables },
match idmap {
true => Some(output_node_map),
false => None,
},
))
}
pub fn keep_intervals<P>(
self,
intervals: impl Iterator<Item = (P, P)>,
) -> Result<Option<TableCollection>, TskitError>
where
P: Into<Position>,
{
self.dump_tables()?.keep_intervals(intervals)
}
#[cfg(feature = "provenance")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
pub fn add_provenance(&mut self, record: &str) -> Result<crate::ProvenanceId, TskitError> {
if record.is_empty() {
return Err(TskitError::ValueError {
got: "empty string".to_string(),
expected: "provenance record".to_string(),
});
}
let timestamp = chrono::Utc::now().to_string();
let rv = unsafe {
ll_bindings::tsk_provenance_table_add_row(
&mut (*self.inner.as_ref().tables).provenances,
timestamp.as_ptr() as *const c_char,
timestamp.len() as ll_bindings::tsk_size_t,
record.as_ptr() as *const c_char,
record.len() as ll_bindings::tsk_size_t,
)
};
handle_tsk_return_value!(rv, crate::ProvenanceId::from(rv))
}
pub fn edge_differences_iter<'ts>(
&'ts self,
) -> crate::edge_differences::EdgeDifferencesIterator<'ts> {
crate::edge_differences::EdgeDifferencesIterator::new(self)
}
pub fn tables(&self) -> &TableCollection {
&self.tables
}
pub fn edges(&self) -> &EdgeTable {
self.tables.edges()
}
pub fn nodes(&self) -> &NodeTable {
self.tables.nodes()
}
pub fn sites(&self) -> &SiteTable {
self.tables.sites()
}
pub fn migrations(&self) -> &MigrationTable {
self.tables.migrations()
}
pub fn mutations(&self) -> &MutationTable {
self.tables.mutations()
}
pub fn individuals(&self) -> &IndividualTable {
self.tables.individuals()
}
pub fn populations(&self) -> &PopulationTable {
self.tables.populations()
}
#[cfg(feature = "provenance")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
pub fn provenances(&self) -> &crate::provenance::ProvenanceTable {
self.tables.provenances()
}
pub fn individuals_iter(&self) -> impl Iterator<Item = crate::IndividualTableRow> + '_ {
self.individuals().iter()
}
pub fn nodes_iter(&self) -> impl Iterator<Item = crate::NodeTableRow> + '_ {
self.nodes().iter()
}
pub fn edges_iter(&self) -> impl Iterator<Item = crate::EdgeTableRow> + '_ {
self.edges().iter()
}
pub fn migrations_iter(&self) -> impl Iterator<Item = crate::MigrationTableRow> + '_ {
self.migrations().iter()
}
pub fn mutations_iter(&self) -> impl Iterator<Item = crate::MutationTableRow> + '_ {
self.mutations().iter()
}
pub fn populations_iter(&self) -> impl Iterator<Item = crate::PopulationTableRow> + '_ {
self.populations().iter()
}
pub fn sites_iter(&self) -> impl Iterator<Item = crate::SiteTableRow> + '_ {
self.sites().iter()
}
#[cfg(feature = "provenance")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
pub fn provenances_iter(
&self,
) -> impl Iterator<Item = crate::provenance::ProvenanceTableRow> + '_ {
self.provenances().iter()
}
pub fn samples_as_vector(&self) -> Vec<crate::NodeId> {
self.tables.samples_as_vector()
}
pub fn create_node_id_vector(
&self,
f: impl FnMut(&crate::NodeTableRow) -> bool,
) -> Vec<crate::NodeId> {
self.tables.create_node_id_vector(f)
}
#[doc(hidden)]
pub fn edge_insertion_order(&self) -> &[crate::EdgeId] {
self.inner.edge_insertion_order()
}
#[doc(hidden)]
pub fn edge_insertion_order_column(
&self,
) -> impl crate::TableColumn<crate::EdgeId, crate::EdgeId> + '_ {
crate::table_column::OpaqueTableColumn(self.edge_insertion_order())
}
#[doc(hidden)]
pub fn edge_removal_order(&self) -> &[crate::EdgeId] {
self.inner.edge_removal_order()
}
#[doc(hidden)]
pub fn edge_removal_order_column(
&self,
) -> impl crate::TableColumn<crate::EdgeId, crate::EdgeId> + '_ {
crate::table_column::OpaqueTableColumn(self.edge_removal_order())
}
}
impl TryFrom<TableCollection> for TreeSequence {
type Error = TskitError;
fn try_from(value: TableCollection) -> Result<Self, Self::Error> {
Self::new(value, TreeSequenceFlags::default())
}
}