tskit 0.16.0

//! Optional module for table and tree sequence provenance tables.
//!
//! This module is enabled via the `"provenance"` feature and provides
//! the following:
//!
//! * [`crate::TableCollection::add_provenance`]
//! * [`crate::TreeSequence::add_provenance`]
//! * [`ProvenanceTable`].
//! * [`crate::Provenance`], which is the value type returned by
//!   [`ProvenanceTable::iter`].
//!

use crate::sys;
use crate::ProvenanceId;
use crate::SizeType;
use sys::bindings as ll_bindings;

/// A provenance table.
///
/// # Notes
///
/// * The type is enabled by the `"provenance"` feature.
///
/// # Examples
///
/// ```rust
/// # #[cfg(feature = "provenance")]
/// # #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
/// {
/// use tskit::provenance::ProvenanceTable;
/// let mut provenances = ProvenanceTable::default();
/// let id = provenances.add_row("message").unwrap();
/// assert_eq!(id, 0);
/// assert_eq!(provenances.num_rows(), 1);
/// # }
/// ```
#[derive(Debug, Default)]
#[repr(transparent)]
pub struct ProvenanceTable {
    table_: sys::ProvenanceTable,
}

impl ProvenanceTable {
    // # Safety
    //
    // * this fn must NEVER by part of the public API
    // * all returned values must only be visible to the public API
    //   by REFERENCE (& or &mut)
    // * the input ptr must not be NULL
    // * the input ptr must point to an initialized table
    pub(crate) unsafe fn new_from_table(
        provenances: *mut ll_bindings::tsk_provenance_table_t,
    ) -> Result<Self, crate::TskitError> {
        let ptr = std::ptr::NonNull::new(provenances).unwrap();
        let table_ = unsafe { sys::ProvenanceTable::new_borrowed(ptr) };
        Ok(ProvenanceTable { table_ })
    }

    pub(crate) fn as_ref(&self) -> &ll_bindings::tsk_provenance_table_t {
        self.table_.as_ref()
    }

    /// Return the number of rows
    pub fn num_rows(&self) -> SizeType {
        self.as_ref().num_rows.into()
    }

    /// Get the ISO-formatted time stamp for row `row`.
    ///
    /// # Returns
    ///
    /// * `Some(String)` if `row` is valid.
    /// * `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    ///
    /// let mut tables = tskit::TableCollection::new(10.).unwrap();
    /// assert!(tables.add_provenance("foo").is_ok());
    /// if let Some(timestamp) = tables.provenances().timestamp(0) {
    ///  // then 0 is a valid row in the provenance table
    /// }
    /// # else {
    /// # panic!("Expected Some(timestamp)");
    /// # }
    /// ```
    pub fn timestamp<P: Into<ProvenanceId> + Copy>(&self, row: P) -> Option<&str> {
        self.table_.timestamp(row.into())
    }

    /// Get the provenance record for row `row`.
    ///
    /// # Returns
    ///
    /// * `Some(String)` if `row` is valid.
    /// * `None` otherwise.
    ///
    /// # Examples
    ///
    /// ```
    ///
    /// let mut tables = tskit::TableCollection::new(10.).unwrap();
    /// assert!(tables.add_provenance("foo").is_ok());
    /// if let Some(record) = tables.provenances().record(0) {
    ///  // then 0 is a valid row in the provenance table
    ///  # assert_eq!(record, "foo");
    /// }
    /// # else {
    /// # panic!("Expected Some(timestamp)");
    /// # }
    pub fn record<P: Into<ProvenanceId> + Copy>(&self, row: P) -> Option<&str> {
        self.table_.record(row.into())
    }

    /// Obtain a [`crate::Provenance`] for row `row`.
    ///
    /// # Returns
    ///
    /// * `Some(row)` if `r` is valid
    /// * `None` otherwise
    pub fn row<P: Into<ProvenanceId> + Copy>(&self, row: P) -> Option<crate::Provenance<'_>> {
        self.table_.row(row.into())
    }

    /// Return an iterator over rows of the table.
    /// The value of the iterator is [`crate::Provenance`].
    pub fn iter(&self) -> impl Iterator<Item = crate::Provenance<'_>> {
        self.table_.iter()
    }

    /// Clear all data from the table
    pub fn clear(&mut self) -> Result<i32, crate::TskitError> {
        handle_tsk_return_value!(self.table_.clear())
    }

    pub fn add_row(&mut self, record: &str) -> Result<ProvenanceId, crate::TskitError> {
        if record.is_empty() {
            return Err(crate::TskitError::ValueError {
                got: "empty string".to_owned(),
                expected: "provenance record".to_owned(),
            });
        }

        Ok(self.table_.add_row(record)?.into())
    }
}

#[cfg(test)]
mod test_provenances {
    #[test]
    fn test_empty_record_string() {
        // check for tables...
        let mut tables = crate::TableCollection::new(10.).unwrap();
        let s = String::from("");
        assert!(tables.add_provenance(&s).is_err());

        // and for tree sequences...
        tables.build_index().unwrap();
        let mut ts = tables
            .tree_sequence(crate::TreeSequenceFlags::default())
            .unwrap();
        assert!(ts.add_provenance(&s).is_err())
    }

    #[test]
    fn test_add_rows() {
        let records = ["banana".to_string(), "split".to_string()];
        let mut tables = crate::TableCollection::new(10.).unwrap();
        for (i, r) in records.iter().enumerate() {
            let row_id = tables.add_provenance(r).unwrap();
            assert!(row_id == i as crate::sys::bindings::tsk_id_t);
            assert_eq!(tables.provenances().record(row_id).unwrap(), *r);
        }
        assert_eq!(
            usize::try_from(tables.provenances().num_rows()).unwrap(),
            records.len()
        );
        for (i, row) in tables.provenance_iter().enumerate() {
            assert_eq!(records[i], row.record());
        }
        for (i, row) in tables.provenances().iter().enumerate() {
            assert_eq!(records[i], row.record());
        }

        assert!(tables.provenances().row(0).unwrap() == tables.provenances().row(0).unwrap());
        assert!(tables.provenances().row(0).unwrap() != tables.provenances().row(1).unwrap());

        for (row, i) in tables.provenances().iter().zip([0, 1].into_iter()) {
            assert_eq!(row.record(), &records[i]);
            let owned_row = tables.provenances().row(i as i32).unwrap();
            assert_eq!(row, owned_row);
            assert_eq!(owned_row, row);
        }
    }
}