From 52364f576c5ba3d4ddcf83f21fe9dd8cd893e19b Mon Sep 17 00:00:00 2001
From: "Kevin R. Thornton" <krthornt@uci.edu>
Date: Thu, 21 Jul 2022 05:42:28 -0700
Subject: [PATCH] perf: implement slice access to individual data

BREAKING CHANGE: changes return value of
IndividualTable::parents and ::location
---
 src/_macros.rs          | 13 +++++++------
 src/individual_table.rs |  8 ++++----
 src/table_collection.rs | 31 ++++++++++++++++++++++++++++---
 3 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/src/_macros.rs b/src/_macros.rs
index f352e6586..74ed25b32 100644
--- a/src/_macros.rs
+++ b/src/_macros.rs
@@ -70,7 +70,7 @@ macro_rules! unsafe_tsk_ragged_column_access {
         }
     }};
 
-    ($i: expr, $lo: expr, $hi: expr, $array: expr, $offset_array: expr, $offset_array_len: expr, $output_id_type: expr) => {{
+    ($i: expr, $lo: expr, $hi: expr, $array: expr, $offset_array: expr, $offset_array_len: expr, $output_id_type: ty) => {{
         let i = $crate::SizeType::try_from($i)?;
         if $i < $lo || i >= $hi {
             Err(TskitError::IndexError {})
@@ -86,11 +86,12 @@ macro_rules! unsafe_tsk_ragged_column_access {
             if start == stop {
                 Ok(None)
             } else {
-                let mut buffer = vec![];
-                for i in start..stop {
-                    buffer.push($output_id_type(unsafe { *$array.offset(i as isize) }));
-                }
-                Ok(Some(buffer))
+                Ok(Some(unsafe {
+                    std::slice::from_raw_parts(
+                        $array.offset(start as isize) as *const $output_id_type,
+                        stop as usize - start as usize,
+                    )
+                }))
             }
         }
     }};
diff --git a/src/individual_table.rs b/src/individual_table.rs
index c15a4a98f..4b235ebfe 100644
--- a/src/individual_table.rs
+++ b/src/individual_table.rs
@@ -61,8 +61,8 @@ fn make_individual_table_row(table: &IndividualTable, pos: tsk_id_t) -> Option<I
         let rv = IndividualTableRow {
             id: pos.into(),
             flags: table.flags(pos).unwrap(),
-            location: table.location(pos).unwrap(),
-            parents: table.parents(pos).unwrap(),
+            location: table.location(pos).unwrap().map(|s| s.to_vec()),
+            parents: table.parents(pos).unwrap().map(|s| s.to_vec()),
             metadata: table_row_decode_metadata!(table_ref, pos),
         };
         Some(rv)
@@ -131,7 +131,7 @@ impl<'a> IndividualTable<'a> {
     pub fn location<I: Into<IndividualId> + Copy>(
         &self,
         row: I,
-    ) -> Result<Option<Vec<Location>>, TskitError> {
+    ) -> Result<Option<&[Location]>, TskitError> {
         unsafe_tsk_ragged_column_access!(
             row.into().0,
             0,
@@ -151,7 +151,7 @@ impl<'a> IndividualTable<'a> {
     pub fn parents<I: Into<IndividualId> + Copy>(
         &self,
         row: I,
-    ) -> Result<Option<Vec<IndividualId>>, TskitError> {
+    ) -> Result<Option<&[IndividualId]>, TskitError> {
         unsafe_tsk_ragged_column_access!(
             row.into().0,
             0,
diff --git a/src/table_collection.rs b/src/table_collection.rs
index 1a1d32afb..c4435a214 100644
--- a/src/table_collection.rs
+++ b/src/table_collection.rs
@@ -1567,15 +1567,40 @@ mod test_adding_individual {
         );
 
         // Empty slices are a thing, causing None to be in the rows.
-
-        let mut tables = crate::test_fixtures::make_empty_table_collection(10.);
         let row_id = tables
             .add_individual(0, &[] as &[f64], &[] as &[IndividualId])
             .unwrap();
         let row = tables.individuals().row(row_id).unwrap();
-        assert_eq!(row.id, IndividualId::from(0));
+        assert_eq!(row.id, IndividualId::from(1));
         assert!(row.location.is_none());
         assert!(row.parents.is_none());
+        let row_id = tables
+            .add_individual(0, &[0.2, 0.4, 0.6], &[1, 2, 3, 4])
+            .unwrap();
+        assert_eq!(row_id, 2);
+        assert_eq!(
+            tables.individuals().location(row_id).unwrap(),
+            Some(
+                [
+                    crate::Location::from(0.2),
+                    crate::Location::from(0.4),
+                    crate::Location::from(0.6)
+                ]
+                .as_slice()
+            )
+        );
+        assert_eq!(
+            tables.individuals().parents(row_id).unwrap(),
+            Some(
+                [
+                    crate::IndividualId::from(1),
+                    crate::IndividualId::from(2),
+                    crate::IndividualId::from(3),
+                    crate::IndividualId::from(4)
+                ]
+                .as_slice()
+            )
+        );
     }
 
     #[test]