Skip to content

Commit

Permalink
Add EdgeId (#138)
Browse files Browse the repository at this point in the history
  • Loading branch information
molpopgen authored Jul 21, 2021
1 parent 9baff07 commit 6bae100
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 24 deletions.
33 changes: 17 additions & 16 deletions src/edge_table.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
use crate::bindings as ll_bindings;
use crate::metadata;
use crate::{tsk_id_t, tsk_size_t, TskitError};
use crate::{EdgeId, NodeId};

/// Row of an [`EdgeTable`]
pub struct EdgeTableRow {
pub id: tsk_id_t,
pub id: EdgeId,
pub left: f64,
pub right: f64,
pub parent: tsk_id_t,
pub child: tsk_id_t,
pub parent: NodeId,
pub child: NodeId,
pub metadata: Option<Vec<u8>>,
}

Expand All @@ -26,7 +27,7 @@ impl PartialEq for EdgeTableRow {
fn make_edge_table_row(table: &EdgeTable, pos: tsk_id_t) -> Option<EdgeTableRow> {
if pos < table.num_rows() as tsk_id_t {
let rv = EdgeTableRow {
id: pos,
id: pos.into(),
left: table.left(pos).unwrap(),
right: table.right(pos).unwrap(),
parent: table.parent(pos).unwrap(),
Expand Down Expand Up @@ -87,8 +88,8 @@ impl<'a> EdgeTable<'a> {
///
/// Will return [``IndexError``](crate::TskitError::IndexError)
/// if ``row`` is out of range.
pub fn parent(&'a self, row: tsk_id_t) -> Result<tsk_id_t, TskitError> {
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.parent)
pub fn parent<E: Into<EdgeId> + Copy>(&'a self, row: E) -> Result<NodeId, TskitError> {
unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.parent, NodeId)
}

/// Return the ``child`` value from row ``row`` of the table.
Expand All @@ -97,8 +98,8 @@ impl<'a> EdgeTable<'a> {
///
/// Will return [``IndexError``](crate::TskitError::IndexError)
/// if ``row`` is out of range.
pub fn child(&'a self, row: tsk_id_t) -> Result<tsk_id_t, TskitError> {
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.child)
pub fn child<E: Into<EdgeId> + Copy>(&'a self, row: E) -> Result<NodeId, TskitError> {
unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.child, NodeId)
}

/// Return the ``left`` value from row ``row`` of the table.
Expand All @@ -107,8 +108,8 @@ impl<'a> EdgeTable<'a> {
///
/// Will return [``IndexError``](crate::TskitError::IndexError)
/// if ``row`` is out of range.
pub fn left(&'a self, row: tsk_id_t) -> Result<f64, TskitError> {
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.left)
pub fn left<E: Into<EdgeId> + Copy>(&'a self, row: E) -> Result<f64, TskitError> {
unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.left)
}

/// Return the ``right`` value from row ``row`` of the table.
Expand All @@ -117,15 +118,15 @@ impl<'a> EdgeTable<'a> {
///
/// Will return [``IndexError``](crate::TskitError::IndexError)
/// if ``row`` is out of range.
pub fn right(&'a self, row: tsk_id_t) -> Result<f64, TskitError> {
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.right)
pub fn right<E: Into<EdgeId> + Copy>(&'a self, row: E) -> Result<f64, TskitError> {
unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.right)
}

pub fn metadata<T: metadata::MetadataRoundtrip>(
&'a self,
row: tsk_id_t,
row: EdgeId,
) -> Result<Option<T>, TskitError> {
let buffer = metadata_to_vector!(self, row)?;
let buffer = metadata_to_vector!(self, row.0)?;
decode_metadata_row!(T, buffer)
}

Expand All @@ -145,7 +146,7 @@ impl<'a> EdgeTable<'a> {
/// # Errors
///
/// [`TskitError::IndexError`] if `r` is out of range.
pub fn row(&self, r: tsk_id_t) -> Result<EdgeTableRow, TskitError> {
table_row_access!(r, self, make_edge_table_row)
pub fn row<E: Into<EdgeId> + Copy>(&self, r: E) -> Result<EdgeTableRow, TskitError> {
table_row_access!(r.into().0, self, make_edge_table_row)
}
}
10 changes: 10 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,12 +187,22 @@ pub struct MutationId(tsk_id_t);
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
pub struct MigrationId(tsk_id_t);

/// An edge ID
///
/// This is an integer referring to a row of an [``EdgeTable``].
///
/// The features for this type follow the same pattern as for [``NodeId``]
#[repr(transparent)]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
pub struct EdgeId(tsk_id_t);

impl_id_traits!(NodeId);
impl_id_traits!(IndividualId);
impl_id_traits!(PopulationId);
impl_id_traits!(SiteId);
impl_id_traits!(MutationId);
impl_id_traits!(MigrationId);
impl_id_traits!(EdgeId);

// tskit defines this via a type cast
// in a macro. bindgen thus misses it.
Expand Down
51 changes: 43 additions & 8 deletions src/table_collection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use crate::TreeSequenceFlags;
use crate::TskReturnValue;
use crate::TskitTypeAccess;
use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TSK_NULL};
use crate::{IndividualId, MutationId, NodeId, PopulationId, SiteId};
use crate::{EdgeId, IndividualId, MutationId, NodeId, PopulationId, SiteId};
use ll_bindings::tsk_table_collection_free;

/// A table collection.
Expand Down Expand Up @@ -172,7 +172,7 @@ impl TableCollection {
right: f64,
parent: P,
child: C,
) -> TskReturnValue {
) -> Result<EdgeId, TskitError> {
self.add_edge_with_metadata(left, right, parent, child, None)
}

Expand All @@ -184,7 +184,7 @@ impl TableCollection {
parent: P,
child: C,
metadata: Option<&dyn MetadataRoundtrip>,
) -> TskReturnValue {
) -> Result<EdgeId, TskitError> {
let md = EncodedMetadata::new(metadata)?;
let rv = unsafe {
ll_bindings::tsk_edge_table_add_row(
Expand All @@ -198,7 +198,7 @@ impl TableCollection {
)
};

handle_tsk_return_value!(rv)
handle_tsk_return_value!(rv, EdgeId::from(rv))
}

/// Add a row to the individual table
Expand Down Expand Up @@ -440,11 +440,11 @@ impl TableCollection {
/// If `self.is_indexed()` is `true`, return a non-owning
/// slice containing the edge insertion order.
/// Otherwise, return `None`.
pub fn edge_insertion_order(&self) -> Option<&[tsk_id_t]> {
pub fn edge_insertion_order(&self) -> Option<&[EdgeId]> {
if self.is_indexed() {
Some(unsafe {
std::slice::from_raw_parts(
(*self.as_ptr()).indexes.edge_insertion_order,
(*self.as_ptr()).indexes.edge_insertion_order as *const EdgeId,
(*self.as_ptr()).indexes.num_edges as usize,
)
})
Expand All @@ -456,11 +456,11 @@ impl TableCollection {
/// If `self.is_indexed()` is `true`, return a non-owning
/// slice containing the edge removal order.
/// Otherwise, return `None`.
pub fn edge_removal_order(&self) -> Option<&[tsk_id_t]> {
pub fn edge_removal_order(&self) -> Option<&[EdgeId]> {
if self.is_indexed() {
Some(unsafe {
std::slice::from_raw_parts(
(*self.as_ptr()).indexes.edge_removal_order,
(*self.as_ptr()).indexes.edge_removal_order as *const EdgeId,
(*self.as_ptr()).indexes.num_edges as usize,
)
})
Expand Down Expand Up @@ -803,6 +803,41 @@ mod test {
assert!(*i >= 0);
assert!(*i < tables.edges().num_rows() as tsk_id_t);
}

// The "transparent" casts are such black magic that we
// should probably test against what C thinks is going on :)
let input = unsafe {
std::slice::from_raw_parts(
(*tables.as_ptr()).indexes.edge_insertion_order,
(*tables.as_ptr()).indexes.num_edges as usize,
)
};

assert!(!input.is_empty());

let tables_input = tables.edge_insertion_order().unwrap();

assert_eq!(input.len(), tables_input.len());

for i in 0..input.len() {
assert_eq!(EdgeId::from(input[i]), tables_input[i]);
}

let output = unsafe {
std::slice::from_raw_parts(
(*tables.as_ptr()).indexes.edge_removal_order,
(*tables.as_ptr()).indexes.num_edges as usize,
)
};
assert!(!output.is_empty());

let tables_output = tables.edge_removal_order().unwrap();

assert_eq!(output.len(), tables_output.len());

for i in 0..output.len() {
assert_eq!(EdgeId::from(output[i]), tables_output[i]);
}
}

#[test]
Expand Down

0 comments on commit 6bae100

Please sign in to comment.