Merge pull request #20 from molpopgen/metadata

Add metadata support.

Author: molpopgen

.github/workflows/tests.yml

@@ -33,7 +33,9 @@ jobs:
           override: true
       - run: cargo check
       - name: run tests
-        run: cargo test
+        run: |
+            cargo test
+            cargo test --examples
 
   build-osx:
     runs-on: macos-latest
@@ -53,7 +55,9 @@ jobs:
           override: true
       - run: cargo check
       - name: run tests
-        run: cargo test 
+        run: |
+            cargo test
+            cargo test --examples
 
   fmt:
     name: rust fmt

Cargo.toml

@@ -13,6 +13,11 @@ repository = "https://github.com/molpopgen/tskit_rust"
 
 [dependencies]
 thiserror = "1.0"
+libc = "0.2.81"
+
+[dev-dependencies]
+serde = {version = "1.0.118", features = ["derive"]}
+bincode = "1.3.1"
 
 [build-dependencies]
 bindgen = "0.56.0"

examples/mutation.rs

@@ -0,0 +1,42 @@
+#![macro_use]
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Mutation {
+    pub effect_size: f64,
+    pub dominance: f64,
+    pub origin_time: i32,
+}
+
+macro_rules! make_mutation_metadata_run {
+    () => {
+        pub fn run() {
+            let mut tables = TableCollection::new(1000.).unwrap();
+            // The simulation generates a mutation:
+            let m = Mutation {
+                effect_size: -0.235423,
+                dominance: 0.5,
+                origin_time: 1,
+            };
+
+            // The mutation's data are included as metadata:
+            tables
+                .add_mutation_with_metadata(0, 0, 0, 0.0, None, Some(&m))
+                .unwrap();
+
+            // Decoding requres 2 unwraps:
+            // 1. The first is to handle errors.
+            // 2. The second is b/c metadata are optional,
+            //    so a row may return None.
+            let decoded = tables.mutations().metadata::<Mutation>(0).unwrap().unwrap();
+
+            // Check that we've made the round trip:
+            assert_eq!(decoded.origin_time, 1);
+            assert!((m.effect_size - decoded.effect_size).abs() < f64::EPSILON);
+            assert!((m.dominance - decoded.dominance).abs() < f64::EPSILON);
+        }
+    };
+}
+
+#[allow(dead_code)]
+fn main() {}

examples/mutation_metadata_bincode.rs

@@ -0,0 +1,38 @@
+use tskit_rust::metadata;
+use tskit_rust::*;
+mod mutation;
+
+use mutation::Mutation;
+
+impl metadata::MetadataRoundtrip for Mutation {
+    fn encode(&self) -> Result<Vec<u8>, metadata::MetadataError> {
+        let mut rv = vec![];
+        rv.extend(self.effect_size.to_le_bytes().iter().copied());
+        rv.extend(self.dominance.to_le_bytes().iter().copied());
+        rv.extend(self.origin_time.to_le_bytes().iter().copied());
+        Ok(rv)
+    }
+
+    fn decode(md: &[u8]) -> Result<Self, metadata::MetadataError> {
+        use std::convert::TryInto;
+        let (effect_size_bytes, rest) = md.split_at(std::mem::size_of::<f64>());
+        let (dominance_bytes, rest) = rest.split_at(std::mem::size_of::<f64>());
+        let (origin_time_bytes, _) = rest.split_at(std::mem::size_of::<i32>());
+        Ok(Self {
+            effect_size: f64::from_le_bytes(effect_size_bytes.try_into().unwrap()),
+            dominance: f64::from_le_bytes(dominance_bytes.try_into().unwrap()),
+            origin_time: i32::from_le_bytes(origin_time_bytes.try_into().unwrap()),
+        })
+    }
+}
+
+make_mutation_metadata_run!();
+
+#[test]
+fn run_bincode() {
+    run();
+}
+
+fn main() {
+    run();
+}

examples/mutation_metadata_std.rs

@@ -0,0 +1,29 @@
+use tskit_rust::metadata;
+use tskit_rust::*;
+mod mutation;
+
+use mutation::Mutation;
+
+// Implement the metadata trait for our mutation
+// type.  Will will use the standard library for the implementation
+// details.
+impl metadata::MetadataRoundtrip for Mutation {
+    fn encode(&self) -> Result<Vec<u8>, metadata::MetadataError> {
+        Ok(bincode::serialize(&self).unwrap())
+    }
+
+    fn decode(md: &[u8]) -> Result<Self, metadata::MetadataError> {
+        Ok(bincode::deserialize(md).unwrap())
+    }
+}
+
+make_mutation_metadata_run!();
+
+#[test]
+fn run_std() {
+    run();
+}
+
+fn main() {
+    run();
+}

src/_macros.rs

@@ -58,6 +58,42 @@ macro_rules! build_tskit_type {
     };
 }
 
+macro_rules! metadata_to_vector {
+    ($T: ty, $self: expr, $row: expr) => {
+        crate::metadata::char_column_to_vector(
+            $self.table_.metadata,
+            $self.table_.metadata_offset,
+            $row,
+            $self.table_.num_rows,
+            $self.table_.metadata_length,
+        )?
+    };
+}
+
+macro_rules! decode_metadata_row {
+    ($T: ty, $buffer: expr) => {
+        match $buffer {
+            None => Ok(None),
+            Some(v) => Ok(Some(<$T as crate::metadata::MetadataRoundtrip>::decode(
+                &v,
+            )?)),
+        }
+    };
+}
+
+macro_rules! process_state_input {
+    ($state: expr) => {
+        match $state {
+            Some(x) => (
+                x.as_ptr() as *const libc::c_char,
+                x.len() as crate::bindings::tsk_size_t,
+                $state,
+            ),
+            None => (std::ptr::null(), 0, $state),
+        }
+    };
+}
+
 #[cfg(test)]
 mod test {
     use crate::error::TskitRustError;

src/edge_table.rs

@@ -1,4 +1,5 @@
 use crate::bindings as ll_bindings;
+use crate::metadata;
 use crate::{tsk_id_t, tsk_size_t, TskitRustError};
 
 /// An immutable view of an edge table.
@@ -59,4 +60,12 @@ impl<'a> EdgeTable<'a> {
     pub fn right(&'a self, row: tsk_id_t) -> Result<f64, TskitRustError> {
         unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.right);
     }
+
+    pub fn metadata<T: metadata::MetadataRoundtrip>(
+        &'a self,
+        row: tsk_id_t,
+    ) -> Result<Option<T>, TskitRustError> {
+        let buffer = metadata_to_vector!(T, self, row);
+        decode_metadata_row!(T, buffer)
+    }
 }

src/error.rs

@@ -16,6 +16,13 @@ pub enum TskitRustError {
     /// Wrapper around tskit C API error codes.
     #[error("{}", get_tskit_error_message(*code))]
     ErrorCode { code: i32 },
+    /// A redirection of [``crate::metadata::MetadataError``]
+    #[error("{value:?}")]
+    MetadataError {
+        /// The redirected error
+        #[from]
+        value: crate::metadata::MetadataError,
+    },
 }
 
 /// Takes the return code from a tskit

src/lib.rs

@@ -10,6 +10,7 @@ mod _macros; // Starts w/_ to be sorted at front by rustfmt!
 mod edge_table;
 pub mod error;
 pub mod ffi;
+pub mod metadata;
 mod mutation_table;
 mod node_table;
 mod population_table;

src/metadata.rs

@@ -0,0 +1,149 @@
+use crate::bindings::{tsk_id_t, tsk_size_t};
+use thiserror::Error;
+
+pub trait MetadataRoundtrip {
+    fn encode(&self) -> Result<Vec<u8>, MetadataError>;
+    fn decode(md: &[u8]) -> Result<Self, MetadataError>
+    where
+        Self: Sized;
+}
+
+pub(crate) struct EncodedMetadata {
+    encoded: Option<Vec<u8>>,
+}
+
+impl EncodedMetadata {
+    pub(crate) fn new(md: Option<&dyn MetadataRoundtrip>) -> Result<Self, MetadataError> {
+        match md {
+            Some(x) => {
+                let e = x.encode()?;
+                Ok(Self { encoded: Some(e) })
+            }
+            None => Ok(Self { encoded: None }),
+        }
+    }
+
+    pub(crate) fn as_ptr(&self) -> *const libc::c_char {
+        match &self.encoded {
+            Some(x) => x.as_ptr() as *const libc::c_char,
+            None => std::ptr::null(),
+        }
+    }
+
+    pub(crate) fn len(&self) -> tsk_size_t {
+        match &self.encoded {
+            Some(x) => x.len() as tsk_size_t,
+            None => 0,
+        }
+    }
+}
+
+#[derive(Error, Debug, PartialEq)]
+pub enum MetadataError {
+    /// Error related to types implementing
+    /// [``MetadataRoundtrip``]
+    #[error("{}", *msg)]
+    RoundtripError { msg: String },
+}
+
+pub(crate) fn char_column_to_vector(
+    column: *const libc::c_char,
+    column_offset: *const tsk_size_t,
+    row: tsk_id_t,
+    num_rows: tsk_size_t,
+    column_length: tsk_size_t,
+) -> Result<Option<Vec<u8>>, crate::TskitRustError> {
+    if row < 0 || (row as tsk_size_t) >= num_rows {
+        return Err(crate::TskitRustError::IndexError {});
+    }
+    if column_length == 0 {
+        return Ok(None);
+    }
+    let start = unsafe { *column_offset.offset(row as isize) };
+    let stop = if (row as tsk_size_t) < num_rows {
+        unsafe { *column_offset.offset((row + 1) as isize) }
+    } else {
+        column_length
+    };
+    if start >= stop {
+        return Ok(None);
+    }
+    if column_length == 0 {
+        return Ok(None);
+    }
+    let mut buffer = vec![];
+    for i in start..stop {
+        buffer.push(unsafe { *column.offset(i as isize) } as u8);
+    }
+    Ok(Some(buffer))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_vec8_cast_to_c_string() {
+        let v: Vec<u8> = vec![0, 1, b'\0', 2, 3];
+        let c = v.as_ptr() as *const libc::c_char;
+        for (i, vi) in v.iter().enumerate() {
+            assert_eq!(*vi as i8, unsafe { *c.add(i) });
+        }
+
+        let _ = match Some(&v) {
+            Some(x) => x.as_ptr() as *const libc::c_char,
+            None => std::ptr::null(),
+        };
+    }
+
+    struct F {
+        x: i32,
+        y: u32,
+    }
+
+    impl MetadataRoundtrip for F {
+        fn encode(&self) -> Result<Vec<u8>, MetadataError> {
+            let mut rv = vec![];
+            rv.extend(self.x.to_le_bytes().iter().copied());
+            rv.extend(self.y.to_le_bytes().iter().copied());
+            Ok(rv)
+        }
+        fn decode(md: &[u8]) -> Result<Self, MetadataError> {
+            use std::convert::TryInto;
+            let (x_int_bytes, rest) = md.split_at(std::mem::size_of::<i32>());
+            let (y_int_bytes, _) = rest.split_at(std::mem::size_of::<u32>());
+            Ok(Self {
+                x: i32::from_le_bytes(x_int_bytes.try_into().unwrap()),
+                y: u32::from_le_bytes(y_int_bytes.try_into().unwrap()),
+            })
+        }
+    }
+
+    #[test]
+    fn test_metadata_round_trip() {
+        let f = F { x: -3, y: 42 };
+        let v = f.encode().unwrap();
+        let c = v.as_ptr() as *const libc::c_char;
+        let mut d = vec![];
+        for i in 0..v.len() {
+            d.push(unsafe { *c.add(i as usize) as u8 });
+        }
+        let df = F::decode(&d).unwrap();
+        assert_eq!(f.x, df.x);
+        assert_eq!(f.y, df.y);
+    }
+
+    #[test]
+    fn test_encoded_metadata_roundtrip() {
+        let f = F { x: -3, y: 42 };
+        let enc = EncodedMetadata::new(Some(&f)).unwrap();
+        let p = enc.as_ptr();
+        let mut d = vec![];
+        for i in 0..enc.len() {
+            d.push(unsafe { *p.add(i as usize) as u8 });
+        }
+        let df = F::decode(&d).unwrap();
+        assert_eq!(f.x, df.x);
+        assert_eq!(f.y, df.y);
+    }
+}