Skip to content

Commit 85731da

Browse files
committed
feat: access nodes in Tree as slice of NodeId
1 parent e5a8c0b commit 85731da

File tree

1 file changed

+40
-0
lines changed

1 file changed

+40
-0
lines changed

src/sys/tree.rs

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,46 @@ impl<'treeseq> LLTree<'treeseq> {
156156
}
157157
}
158158

159+
pub fn nodes(&self, order: NodeTraversalOrder) -> Result<Box<[NodeId]>, TskitError> {
160+
let mut nodes: Vec<NodeId> = vec![
161+
NodeId::NULL;
162+
unsafe { super::bindings::tsk_tree_get_size_bound(self.as_ll_ref()) }
163+
as usize
164+
];
165+
166+
let mut num_nodes: super::bindings::tsk_size_t = 0;
167+
let ptr = std::ptr::addr_of_mut!(num_nodes);
168+
unsafe {
169+
super::bindings::tsk_tree_preorder(
170+
self.as_ll_ref(),
171+
nodes.as_mut_ptr() as *mut super::bindings::tsk_id_t,
172+
ptr,
173+
);
174+
}
175+
176+
let code = match order {
177+
NodeTraversalOrder::Preorder => unsafe {
178+
super::bindings::tsk_tree_preorder(
179+
self.as_ll_ref(),
180+
nodes.as_mut_ptr() as *mut super::bindings::tsk_id_t,
181+
ptr,
182+
)
183+
},
184+
NodeTraversalOrder::Postorder => unsafe {
185+
super::bindings::tsk_tree_preorder(
186+
self.as_ll_ref(),
187+
nodes.as_mut_ptr() as *mut super::bindings::tsk_id_t,
188+
ptr,
189+
)
190+
},
191+
};
192+
if code == 0 {
193+
nodes.resize(num_nodes as usize, NodeId::NULL);
194+
}
195+
196+
handle_tsk_return_value!(code, nodes.into_boxed_slice())
197+
}
198+
159199
pub fn children(&self, u: NodeId) -> impl Iterator<Item = NodeId> + '_ {
160200
NodeIteratorAdapter(ChildIterator::new(self, u))
161201
}

0 commit comments

Comments
 (0)