From 04ccd96523f60e0c4af15112b921c5d74fd01945 Mon Sep 17 00:00:00 2001
From: "Kevin R. Thornton" <krthornt@uci.edu>
Date: Thu, 26 Jan 2023 14:06:46 -0800
Subject: [PATCH] feat: TableViews can now lend &mut for all table types.

---
 src/table_collection.rs | 37 +++++++++++++++++++++++++++++++++++++
 src/table_views.rs      | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/src/table_collection.rs b/src/table_collection.rs
index 4ff97d75b..e26e71463 100644
--- a/src/table_collection.rs
+++ b/src/table_collection.rs
@@ -1244,6 +1244,29 @@ impl TableCollection {
         to self.views {
             /// Get mutable reference to the [``NodeTable``](crate::NodeTable).
             pub fn nodes_mut(&mut self) -> &mut crate::NodeTable;
+
+            /// Get mutable reference to the [``EdgeTable``](crate::EdgeTable).
+            pub fn edges_mut(&mut self) -> &mut crate::EdgeTable;
+
+            /// Get mutable reference to the [``SiteTable``](crate::SiteTable).
+            pub fn sites_mut(&mut self) -> &mut crate::SiteTable;
+
+            /// Get mutable reference to the [``MutationTable``](crate::MutationTable).
+            pub fn mutations_mut(&mut self) -> &mut crate::MutationTable;
+
+            /// Get mutable reference to the [``IndividualTable``](crate::IndividualTable).
+            pub fn individuals_mut(&mut self) -> &mut crate::IndividualTable;
+
+            /// Get mutable reference to the [``PopulationTable``](crate::PopulationTable).
+            pub fn populations_mut(&mut self) -> &mut crate::PopulationTable;
+
+            /// Get mutable reference to the [``MigrationTable``](crate::PopulationTable).
+            pub fn migrations_mut(&mut self) -> &mut crate::MigrationTable;
+
+            #[cfg(feature="provenance")]
+            #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
+            /// Get mutable reference to the [``ProvenanceTable``](crate::ProvenanceTable).
+            pub fn provenances_mut(&mut self) -> &mut crate::provenance::ProvenanceTable;
         }
     }
 
@@ -1259,3 +1282,17 @@ impl TableCollection {
         &mut *self.inner
     }
 }
+
+#[test]
+fn test_mutable_node_references() {
+    let mut tables = crate::TableCollection::new(10.0).unwrap();
+    let _nodes: &mut crate::NodeTable = tables.nodes_mut();
+    let _edges: &mut crate::EdgeTable = tables.edges_mut();
+    let _sites: &mut crate::SiteTable = tables.sites_mut();
+    let _mutations: &mut crate::MutationTable = tables.mutations_mut();
+    let _individuals: &mut crate::IndividualTable = tables.individuals_mut();
+    let _populations: &mut crate::PopulationTable = tables.populations_mut();
+    let _migrations: &mut crate::MigrationTable = tables.migrations_mut();
+    #[cfg(feature = "provenance")]
+    let _provenances: &mut crate::provenance::ProvenanceTable = tables.provenances_mut();
+}
diff --git a/src/table_views.rs b/src/table_views.rs
index ea682496b..a99c96dc8 100644
--- a/src/table_views.rs
+++ b/src/table_views.rs
@@ -90,6 +90,11 @@ impl TableViews {
         &self.edges
     }
 
+    /// Get mutable reference to the [``EdgeTable``](crate::EdgeTable).
+    pub fn edges_mut(&mut self) -> &mut EdgeTable {
+        &mut self.edges
+    }
+
     /// Get reference to the [``NodeTable``](crate::NodeTable).
     pub fn nodes(&self) -> &NodeTable {
         &self.nodes
@@ -105,26 +110,51 @@ impl TableViews {
         &self.sites
     }
 
+    /// Get mutable reference to the [``SiteTable``](crate::SiteTable).
+    pub fn sites_mut(&mut self) -> &mut SiteTable {
+        &mut self.sites
+    }
+
     /// Get reference to the [``MutationTable``](crate::MutationTable).
     pub fn mutations(&self) -> &MutationTable {
         &self.mutations
     }
 
+    /// Get mutable reference to the [``MutationTable``](crate::MutationTable).
+    pub fn mutations_mut(&mut self) -> &mut MutationTable {
+        &mut self.mutations
+    }
+
     /// Get reference to the [``IndividualTable``](crate::IndividualTable).
     pub fn individuals(&self) -> &IndividualTable {
         &self.individuals
     }
 
+    /// Get mutable reference to the [``IndividualTable``](crate::IndividualTable).
+    pub fn individuals_mut(&mut self) -> &mut IndividualTable {
+        &mut self.individuals
+    }
+
     /// Get reference to the [``PopulationTable``](crate::PopulationTable).
     pub fn populations(&self) -> &PopulationTable {
         &self.populations
     }
 
+    /// Get mutable reference to the [``PopulationTable``](crate::PopulationTable).
+    pub fn populations_mut(&mut self) -> &mut PopulationTable {
+        &mut self.populations
+    }
+
     /// Get reference to the [``MigrationTable``](crate::MigrationTable).
     pub fn migrations(&self) -> &MigrationTable {
         &self.migrations
     }
 
+    /// Get mutable reference to the [``MigrationTable``](crate::MigrationTable).
+    pub fn migrations_mut(&mut self) -> &mut MigrationTable {
+        &mut self.migrations
+    }
+
     #[cfg(feature = "provenance")]
     #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
     /// Get reference to the [``ProvenanceTable``](crate::provenance::ProvenanceTable)
@@ -132,6 +162,13 @@ impl TableViews {
         &self.provenances
     }
 
+    #[cfg(feature = "provenance")]
+    #[cfg_attr(doc_cfg, doc(cfg(feature = "provenance")))]
+    /// Get mutable reference to the [``ProvenanceTable``](crate::provenance::ProvenanceTable)
+    pub fn provenances_mut(&mut self) -> &mut ProvenanceTable {
+        &mut self.provenances
+    }
+
     /// Return an iterator over the edges.
     pub fn edges_iter(&self) -> impl Iterator<Item = crate::EdgeTableRow> + '_ {
         self.edges.iter()