1 #include "multitree_iterator.hpp"
2
3 namespace terraces {
4
multitree_iterator(const multitree_node * root)5 multitree_iterator::multitree_iterator(const multitree_node* root)
6 : m_tree(2 * root->num_leaves - 1), m_choices(m_tree.size()),
7 m_unconstrained_choices(m_tree.size()) {
8 m_choices[0] = {root};
9 init_subtree(0);
10 }
11
init_subtree(index i,index single_leaf)12 void multitree_iterator::init_subtree(index i, index single_leaf) {
13 m_tree[i].lchild() = none;
14 m_tree[i].rchild() = none;
15 m_tree[i].taxon() = single_leaf;
16 }
17
init_subtree(index i,multitree_nodes::two_leaves two_leaves)18 void multitree_iterator::init_subtree(index i, multitree_nodes::two_leaves two_leaves) {
19 const auto l = i + 1;
20 const auto r = i + 2;
21 m_tree[i].lchild() = l;
22 m_tree[i].rchild() = r;
23 m_tree[i].taxon() = none;
24 m_tree[l] = {i, none, none, two_leaves.left_leaf};
25 m_tree[r] = {i, none, none, two_leaves.right_leaf};
26 }
27
init_subtree(index i,multitree_nodes::unconstrained unconstrained)28 void multitree_iterator::init_subtree(index i, multitree_nodes::unconstrained unconstrained) {
29 m_unconstrained_choices[i] = small_bipartition::full_set(unconstrained.num_leaves());
30 init_subtree_unconstrained(i, unconstrained);
31 }
32
init_subtree_unconstrained(index i,multitree_nodes::unconstrained data)33 void multitree_iterator::init_subtree_unconstrained(index i, multitree_nodes::unconstrained data) {
34 const auto& bip = m_unconstrained_choices[i];
35 auto& node = m_tree[i];
36 if (bip.num_leaves() <= 2) {
37 if (bip.num_leaves() == 1) {
38 node.lchild() = none;
39 node.rchild() = none;
40 node.taxon() = data.begin[bip.leftmost_leaf()];
41 } else {
42 node.lchild() = i + 1;
43 node.rchild() = i + 2;
44 node.taxon() = none;
45 m_tree[i + 1] = {i, none, none, data.begin[bip.leftmost_leaf()]};
46 m_tree[i + 2] = {i, none, none, data.begin[bip.rightmost_leaf()]};
47 }
48 } else {
49 const auto lbip = small_bipartition{bip.left_mask()};
50 const auto rbip = small_bipartition{bip.right_mask()};
51 const auto left = i + 1;
52 const auto right = i + 1 + 2 * lbip.num_leaves() - 1;
53 node.lchild() = left;
54 node.rchild() = right;
55 node.taxon() = none;
56 m_unconstrained_choices[left] = lbip;
57 m_unconstrained_choices[right] = rbip;
58 m_tree[node.lchild()].parent() = i;
59 m_tree[node.rchild()].parent() = i;
60
61 init_subtree_unconstrained(right, data);
62 init_subtree_unconstrained(left, data);
63 init_subtree_unconstrained(right, data);
64 }
65 }
66
init_subtree(index i,multitree_nodes::inner_node inner)67 void multitree_iterator::init_subtree(index i, multitree_nodes::inner_node inner) {
68 const auto left = inner.left;
69 const auto right = inner.right;
70 const auto lindex = i + 1;
71 const auto rindex = lindex + (2 * left->num_leaves - 1);
72 m_tree[i].lchild() = lindex;
73 m_tree[i].rchild() = rindex;
74 m_tree[i].taxon() = none;
75 m_tree[lindex].parent() = i;
76 m_tree[rindex].parent() = i;
77 m_choices[lindex] = {left};
78 m_choices[rindex] = {right};
79 init_subtree(lindex);
80 init_subtree(rindex);
81 }
82
init_subtree(index i)83 void multitree_iterator::init_subtree(index i) {
84 const auto mt_node = m_choices[i].current;
85 switch (mt_node->type) {
86 case multitree_node_type::base_single_leaf:
87 return init_subtree(i, mt_node->single_leaf);
88 case multitree_node_type::base_two_leaves:
89 return init_subtree(i, mt_node->two_leaves);
90 case multitree_node_type::base_unconstrained:
91 return init_subtree(i, mt_node->unconstrained);
92 case multitree_node_type::inner_node:
93 return init_subtree(i, mt_node->inner_node);
94 case multitree_node_type::alternative_array:
95 assert(false && "Malformed multitree: Nested alternative_arrays");
96 break;
97 case multitree_node_type::unexplored:
98 assert(false && "Must not use multitree_iterator with unexplored nodes");
99 break;
100 }
101 }
102
next(index root)103 bool multitree_iterator::next(index root) {
104 auto node = m_tree[root];
105 auto left = node.lchild();
106 auto right = node.rchild();
107 auto& choice = m_choices[root];
108 switch (choice.current->type) {
109 case multitree_node_type::base_single_leaf:
110 case multitree_node_type::base_two_leaves:
111 return false;
112 case multitree_node_type::base_unconstrained:
113 return next_unconstrained(root, choice.current->unconstrained);
114 case multitree_node_type::inner_node:
115 case multitree_node_type::alternative_array:
116 return next(left) || (next(right) && reset(left)) ||
117 (choice.has_choices() && choice.next() && (init_subtree(root), true));
118 case multitree_node_type::unexplored: {
119 assert(false && "Must not use multitree_iterator with unexplored nodes");
120 return false;
121 }
122 default:
123 assert(false && "Unknown node type in multitree");
124 return false;
125 }
126 }
127
next_unconstrained(index root,multitree_nodes::unconstrained data)128 bool multitree_iterator::next_unconstrained(index root, multitree_nodes::unconstrained data) {
129 auto node = m_tree[root];
130 auto left = node.lchild();
131 auto right = node.rchild();
132 auto& choice = m_unconstrained_choices[root];
133 if (!choice.has_choices()) {
134 return false;
135 }
136 return next_unconstrained(left, data) ||
137 (next_unconstrained(right, data) && reset_unconstrained(left, data)) ||
138 (choice.next() && (init_subtree_unconstrained(root, data), true));
139 }
140
reset(index root)141 bool multitree_iterator::reset(index root) {
142 auto& choice = m_choices[root];
143 if (choice.has_choices()) {
144 choice.reset();
145 }
146 switch (choice.current->type) {
147 case multitree_node_type::base_single_leaf:
148 case multitree_node_type::base_two_leaves:
149 break;
150 case multitree_node_type::base_unconstrained:
151 reset_unconstrained(root, choice.current->unconstrained);
152 break;
153 case multitree_node_type::inner_node:
154 case multitree_node_type::alternative_array:
155 init_subtree(root);
156 break;
157 case multitree_node_type::unexplored: {
158 assert(false && "Must not use multitree_iterator with unexplored nodes");
159 break;
160 }
161 default:
162 assert(false && "Unknown node type in multitree");
163 break;
164 }
165 return true;
166 }
167
reset_unconstrained(index root,multitree_nodes::unconstrained data)168 bool multitree_iterator::reset_unconstrained(index root, multitree_nodes::unconstrained data) {
169 auto& choice = m_unconstrained_choices[root];
170 if (choice.has_choices()) {
171 choice.reset();
172 }
173 init_subtree_unconstrained(root, data);
174 return true;
175 }
176
next()177 bool multitree_iterator::next() { return next(0); }
178
179 } // namespace terraces
180