clean up ownership semantics

molpopgen · Jul 7, 2021 · cf96361 · cf96361
1 parent 29ab1f2
commit cf96361
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 4 deletions.
diff --git a/fwdpp/ts/types/tree_sequence.hpp b/fwdpp/ts/types/tree_sequence.hpp
@@ -23,14 +23,13 @@ namespace fwdpp
             template <typename SignedInteger> class tree_iterator
             {
               private:
-                const tree_sequence<SignedInteger>* treeseq_;
+                std::reference_wrapper<const tree_sequence<SignedInteger>> treeseq_;
                 marginal_tree<SignedInteger> tree_;
 
               public:
-                tree_iterator(const tree_sequence<SignedInteger>* ts,
-                              const std::vector<SignedInteger>& samples,
+                tree_iterator(const tree_sequence<SignedInteger>& ts,
                               std::uint32_t flags)
-                    : treeseq_{ts}, tree_{ts->nodes.size(), samples,
+                    : treeseq_{ts}, tree_{ts.num_nodes(), ts.samples(),
                                           static_cast<bool>(flags
                                                             & tree_flags::TRACK_SAMPLES)}
                 {
@@ -139,6 +138,24 @@ namespace fwdpp
                       num_trees_{init_num_trees()}
                 {
                 }
+
+                std::size_t
+                num_nodes() const
+                {
+                    return tables_->nodes.size();
+                }
+
+                const std::vector<SignedInteger>&
+                samples() const
+                {
+                    return samples_;
+                }
+
+                tree_iterator<SignedInteger>
+                trees(std::uint32_t flags)
+                {
+                    return tree_iterator<SignedInteger>{*this, flags};
+                }
             };
 
 #if __cplusplus < 201703L

diff --git a/testsuite/tree_sequences/test_tree_sequence.cc b/testsuite/tree_sequences/test_tree_sequence.cc
@@ -9,4 +9,6 @@ using tree_sequence64 = fwdpp::ts::types::tree_sequence<std::int64_t>;
 //explicit instantiations
 template class fwdpp::ts::types::tree_sequence<std::int32_t>;
 template class fwdpp::ts::types::tree_sequence<std::int64_t>;
+template class fwdpp::ts::types::tree_iterator<std::int32_t>;
+template class fwdpp::ts::types::tree_iterator<std::int64_t>;