diff --git a/fwdpp/ts/types/tree_sequence.hpp b/fwdpp/ts/types/tree_sequence.hpp index 0c41c14b3..475fc9548 100644 --- a/fwdpp/ts/types/tree_sequence.hpp +++ b/fwdpp/ts/types/tree_sequence.hpp @@ -23,14 +23,13 @@ namespace fwdpp template class tree_iterator { private: - const tree_sequence* treeseq_; + std::reference_wrapper> treeseq_; marginal_tree tree_; public: - tree_iterator(const tree_sequence* ts, - const std::vector& samples, + tree_iterator(const tree_sequence& ts, std::uint32_t flags) - : treeseq_{ts}, tree_{ts->nodes.size(), samples, + : treeseq_{ts}, tree_{ts.num_nodes(), ts.samples(), static_cast(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& + samples() const + { + return samples_; + } + + tree_iterator + trees(std::uint32_t flags) + { + return tree_iterator{*this, flags}; + } }; #if __cplusplus < 201703L diff --git a/testsuite/tree_sequences/test_tree_sequence.cc b/testsuite/tree_sequences/test_tree_sequence.cc index 0af150d81..bdb25e00b 100644 --- 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; //explicit instantiations template class fwdpp::ts::types::tree_sequence; template class fwdpp::ts::types::tree_sequence; +template class fwdpp::ts::types::tree_iterator; +template class fwdpp::ts::types::tree_iterator;