Skip to content

Commit ceabe26

Browse files
molpopgenmolpopgen
committed
replace WrappedTskArray with idiomatic slices
1 parent ed8329b commit ceabe26

File tree

3 files changed

+153
-158
lines changed

3 files changed

+153
-158
lines changed

‎src/_macros.rs

Lines changed: 8 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -168,28 +168,6 @@ macro_rules! process_state_input {
168168
};
169169
}
170170

171-
macro_rules! index_for_wrapped_tsk_array_type {
172-
($name: ty, $index:ty, $output: ty) => {
173-
impl std::ops::Index<$index> for $name {
174-
type Output = $output;
175-
176-
fn index(&self, index: $index) -> &Self::Output {
177-
if index >= self.len() as $index {
178-
panic!("fatal: index out of range");
179-
}
180-
let rv = unsafe { self.array.offset(index as isize) };
181-
unsafe { &*rv }
182-
}
183-
}
184-
};
185-
}
186-
187-
macro_rules! wrapped_tsk_array_traits {
188-
($name: ty, $index:ty, $output: ty) => {
189-
index_for_wrapped_tsk_array_type!($name, $index, $output);
190-
};
191-
}
192-
193171
macro_rules! err_if_not_tracking_samples {
194172
($flags: expr, $rv: expr) => {
195173
match $flags.contains(crate::TreeFlags::SAMPLE_LISTS) {
@@ -228,6 +206,14 @@ macro_rules! iterator_for_nodeiterator {
228206
};
229207
}
230208

209+
macro_rules! tree_array_slice {
210+
($self: ident, $array: ident, $len: expr) => {
211+
unsafe {
212+
std::slice::from_raw_parts((*$self.as_ptr()).$array as *const tsk_id_t, $len as usize)
213+
}
214+
};
215+
}
216+
231217
#[cfg(test)]
232218
mod test {
233219
use crate::error::TskitError;

‎src/ffi.rs

Lines changed: 0 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -1,67 +1,5 @@
11
//! Define traits related to wrapping tskit stuff
22
3-
/// Indexable, iterable wrapper around C
4-
/// arrays.
5-
#[derive(Copy, Clone)]
6-
pub struct WrappedTskArray<T> {
7-
array: *const T,
8-
len_: crate::tsk_size_t,
9-
}
10-
11-
pub struct WrappedTskArrayIter<'a, T: Copy + 'a> {
12-
inner: &'a WrappedTskArray<T>,
13-
pos: crate::tsk_size_t,
14-
}
15-
16-
impl<'a, T: Copy> Iterator for WrappedTskArrayIter<'a, T> {
17-
type Item = T;
18-
19-
fn next(&mut self) -> Option<Self::Item> {
20-
if self.pos >= self.inner.len_ {
21-
None
22-
} else {
23-
let rv = Some(unsafe { *self.inner.array.offset(self.pos as isize) as T });
24-
self.pos += 1;
25-
rv
26-
}
27-
}
28-
}
29-
30-
impl<T: Copy> WrappedTskArray<T> {
31-
pub(crate) fn new(array: *const T, len: crate::tsk_size_t) -> Self {
32-
Self { array, len_: len }
33-
}
34-
35-
pub fn len(&self) -> crate::tsk_size_t {
36-
self.len_
37-
}
38-
39-
pub fn is_empty(&self) -> bool {
40-
self.len_ == 0
41-
}
42-
43-
/// # Safety
44-
///
45-
/// This function returns the raw C pointer,
46-
/// and is thus unsafe.
47-
pub fn as_ptr(&self) -> *const T {
48-
self.array
49-
}
50-
51-
pub fn iter(&self) -> WrappedTskArrayIter<T> {
52-
WrappedTskArrayIter {
53-
inner: self,
54-
pos: 0,
55-
}
56-
}
57-
}
58-
59-
pub(crate) type TskIdArray = WrappedTskArray<crate::tsk_id_t>;
60-
pub(crate) type Tskf64Array = WrappedTskArray<f64>;
61-
62-
wrapped_tsk_array_traits!(TskIdArray, crate::tsk_id_t, crate::tsk_id_t);
63-
wrapped_tsk_array_traits!(Tskf64Array, crate::tsk_id_t, f64);
64-
653
/// Wrap a tskit type
664
pub(crate) trait WrapTskitType<T> {
675
/// Encapsulate tsk_foo_t and return rust
@@ -83,7 +21,6 @@ pub(crate) trait WrapTskitConsumingType<T, C> {
8321
mod tests {
8422
use super::*;
8523
use crate::bindings as ll_bindings;
86-
use crate::tsk_size_t;
8724
use crate::TskitTypeAccess;
8825
use ll_bindings::tsk_table_collection_free;
8926

@@ -119,56 +56,4 @@ mod tests {
11956
let t = TableCollectionMock::new(10.);
12057
assert_eq!(t.sequence_length() as i64, 10);
12158
}
122-
123-
#[test]
124-
fn test_u32_array_wrapper() {
125-
let mut t = TableCollectionMock::new(10.);
126-
127-
let rv = unsafe {
128-
ll_bindings::tsk_edge_table_add_row(
129-
&mut (*t.as_mut_ptr()).edges,
130-
0.,
131-
10.,
132-
0,
133-
17,
134-
std::ptr::null(),
135-
0,
136-
)
137-
};
138-
panic_on_tskit_error!(rv);
139-
140-
let a = TskIdArray::new(unsafe { (*t.as_ptr()).edges.child }, 1);
141-
assert_eq!(a.len(), 1);
142-
assert_eq!(a[0], 17);
143-
144-
let mut v = vec![];
145-
for i in a.iter() {
146-
v.push(i);
147-
}
148-
assert_eq!(v.len() as tsk_size_t, a.len());
149-
assert_eq!(v[0], 17);
150-
}
151-
152-
#[should_panic]
153-
#[test]
154-
fn test_u32_array_wrapper_panic() {
155-
let mut t = TableCollectionMock::new(10.);
156-
157-
let rv = unsafe {
158-
ll_bindings::tsk_edge_table_add_row(
159-
&mut (*t.as_mut_ptr()).edges,
160-
0.,
161-
10.,
162-
0,
163-
17,
164-
std::ptr::null(),
165-
0,
166-
)
167-
};
168-
panic_on_tskit_error!(rv);
169-
170-
let a = TskIdArray::new(unsafe { (*t.as_ptr()).edges.child }, 1);
171-
assert_eq!(a.len(), 1);
172-
assert_eq!(a[1], 17);
173-
}
17459
}

‎src/trees.rs

Lines changed: 145 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -71,54 +71,159 @@ impl Tree {
7171
handle_tsk_return_value!(rv, tree)
7272
}
7373

74-
pub fn parent_array(&self) -> crate::ffi::TskIdArray {
75-
crate::ffi::TskIdArray::new(self.inner.parent, self.inner.num_nodes)
74+
/// ```compile_fail
75+
/// use streaming_iterator::StreamingIterator;
76+
/// let tables = tskit::TableCollection::new(1.).unwrap();
77+
/// let treeseq =
78+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
79+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap();
80+
/// while let Some(tree) = tree_iter.next() {
81+
/// let p = tree.parent_array();
82+
/// drop(tree_iter);
83+
/// for _ in p {}
84+
/// }
85+
/// ```
86+
pub fn parent_array(&self) -> &[tsk_id_t] {
87+
tree_array_slice!(self, parent, self.inner.num_nodes)
7688
}
7789

78-
pub fn samples_array(&self) -> Result<crate::ffi::TskIdArray, TskitError> {
90+
/// ```compile_fail
91+
/// use streaming_iterator::StreamingIterator;
92+
/// let tables = tskit::TableCollection::new(1.).unwrap();
93+
/// let treeseq =
94+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
95+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
96+
/// while let Some(tree) = tree_iter.next() {
97+
/// let p = tree.samples_array();
98+
/// drop(tree_iter);
99+
/// for _ in p {}
100+
/// }
101+
/// ```
102+
pub fn samples_array(&self) -> Result<&[tsk_id_t], TskitError> {
79103
let num_samples =
80104
unsafe { ll_bindings::tsk_treeseq_get_num_samples((*self.as_ptr()).tree_sequence) };
81-
err_if_not_tracking_samples!(
82-
self.flags,
83-
crate::ffi::TskIdArray::new(self.inner.samples, num_samples)
84-
)
105+
err_if_not_tracking_samples!(self.flags, tree_array_slice!(self, samples, num_samples))
85106
}
86107

87-
pub fn next_sample_array(&self) -> Result<crate::ffi::TskIdArray, TskitError> {
108+
/// ```compile_fail
109+
/// use streaming_iterator::StreamingIterator;
110+
/// let tables = tskit::TableCollection::new(1.).unwrap();
111+
/// let treeseq =
112+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
113+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
114+
/// while let Some(tree) = tree_iter.next() {
115+
/// let p = tree.next_sample_array();
116+
/// drop(tree_iter);
117+
/// for _ in p {}
118+
/// }
119+
/// ```
120+
pub fn next_sample_array(&self) -> Result<&[tsk_id_t], TskitError> {
88121
err_if_not_tracking_samples!(
89122
self.flags,
90-
crate::ffi::TskIdArray::new(self.inner.next_sample, self.inner.num_nodes)
123+
tree_array_slice!(self, next_sample, self.inner.num_nodes)
91124
)
92125
}
93126

94-
pub fn left_sample_array(&self) -> Result<crate::ffi::TskIdArray, TskitError> {
127+
/// ```compile_fail
128+
/// use streaming_iterator::StreamingIterator;
129+
/// let tables = tskit::TableCollection::new(1.).unwrap();
130+
/// let treeseq =
131+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
132+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
133+
/// while let Some(tree) = tree_iter.next() {
134+
/// let p = tree.left_sample_array();
135+
/// drop(tree_iter);
136+
/// for _ in p {}
137+
/// }
138+
/// ```
139+
pub fn left_sample_array(&self) -> Result<&[tsk_id_t], TskitError> {
95140
err_if_not_tracking_samples!(
96141
self.flags,
97-
crate::ffi::TskIdArray::new(self.inner.left_sample, self.inner.num_nodes)
142+
tree_array_slice!(self, left_sample, self.inner.num_nodes)
98143
)
99144
}
100145

101-
pub fn right_sample_array(&self) -> Result<crate::ffi::TskIdArray, TskitError> {
146+
/// ```compile_fail
147+
/// use streaming_iterator::StreamingIterator;
148+
/// let tables = tskit::TableCollection::new(1.).unwrap();
149+
/// let treeseq =
150+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
151+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::SAMPLE_LISTS).unwrap();
152+
/// while let Some(tree) = tree_iter.next() {
153+
/// let p = tree.right_sample_array();
154+
/// drop(tree_iter);
155+
/// for _ in p {}
156+
/// }
157+
/// ```
158+
pub fn right_sample_array(&self) -> Result<&[tsk_id_t], TskitError> {
102159
err_if_not_tracking_samples!(
103160
self.flags,
104-
crate::ffi::TskIdArray::new(self.inner.right_sample, self.inner.num_nodes)
161+
tree_array_slice!(self, right_sample, self.inner.num_nodes)
105162
)
106163
}
107164

108-
pub fn left_sib_array(&self) -> crate::ffi::TskIdArray {
109-
crate::ffi::TskIdArray::new(self.inner.left_sib, self.inner.num_nodes)
165+
/// ```compile_fail
166+
/// use streaming_iterator::StreamingIterator;
167+
/// let tables = tskit::TableCollection::new(1.).unwrap();
168+
/// let treeseq =
169+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
170+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap();
171+
/// while let Some(tree) = tree_iter.next() {
172+
/// let p = tree.left_sib_array();
173+
/// drop(tree_iter);
174+
/// for _ in p {}
175+
/// }
176+
/// ```
177+
pub fn left_sib_array(&self) -> &[tsk_id_t] {
178+
tree_array_slice!(self, left_sib, self.inner.num_nodes)
110179
}
111180

112-
pub fn right_sib_array(&self) -> crate::ffi::TskIdArray {
113-
crate::ffi::TskIdArray::new(self.inner.right_sib, self.inner.num_nodes)
181+
/// ```compile_fail
182+
/// use streaming_iterator::StreamingIterator;
183+
/// let tables = tskit::TableCollection::new(1.).unwrap();
184+
/// let treeseq =
185+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
186+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap();
187+
/// while let Some(tree) = tree_iter.next() {
188+
/// let p = tree.right_sib_array();
189+
/// drop(tree_iter);
190+
/// for _ in p {}
191+
/// }
192+
/// ```
193+
pub fn right_sib_array(&self) -> &[tsk_id_t] {
194+
tree_array_slice!(self, right_sib, self.inner.num_nodes)
114195
}
115196

116-
pub fn left_child_array(&self) -> crate::ffi::TskIdArray {
117-
crate::ffi::TskIdArray::new(self.inner.left_child, self.inner.num_nodes)
197+
/// ```compile_fail
198+
/// use streaming_iterator::StreamingIterator;
199+
/// let tables = tskit::TableCollection::new(1.).unwrap();
200+
/// let treeseq =
201+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
202+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap();
203+
/// while let Some(tree) = tree_iter.next() {
204+
/// let p = tree.left_child_array();
205+
/// drop(tree_iter);
206+
/// for _ in p {}
207+
/// }
208+
/// ```
209+
pub fn left_child_array(&self) -> &[tsk_id_t] {
210+
tree_array_slice!(self, left_child, self.inner.num_nodes)
118211
}
119212

120-
pub fn right_child_array(&self) -> crate::ffi::TskIdArray {
121-
crate::ffi::TskIdArray::new(self.inner.right_child, self.inner.num_nodes)
213+
/// ```compile_fail
214+
/// use streaming_iterator::StreamingIterator;
215+
/// let tables = tskit::TableCollection::new(1.).unwrap();
216+
/// let treeseq =
217+
/// tables.tree_sequence(tskit::TreeSequenceFlags::BUILD_INDEXES).unwrap();
218+
/// let mut tree_iter = treeseq.tree_iterator(tskit::TreeFlags::default()).unwrap();
219+
/// while let Some(tree) = tree_iter.next() {
220+
/// let p = tree.right_child_array();
221+
/// drop(tree_iter);
222+
/// for _ in p {}
223+
/// }
224+
/// ```
225+
pub fn right_child_array(&self) -> &[tsk_id_t] {
226+
tree_array_slice!(self, right_child, self.inner.num_nodes)
122227
}
123228

124229
fn left_sample(&self, u: tsk_id_t) -> Result<tsk_id_t, TskitError> {
@@ -1147,4 +1252,23 @@ pub(crate) mod test_trees {
11471252
stops_rev.reverse();
11481253
assert!(stops_fwd == stops_rev);
11491254
}
1255+
1256+
// FIXME: remove later
1257+
#[test]
1258+
fn test_array_lifetime() {
1259+
let treeseq = treeseq_from_small_table_collection_two_trees();
1260+
let mut tree_iter = treeseq.tree_iterator(TreeFlags::default()).unwrap();
1261+
if let Some(tree) = tree_iter.next() {
1262+
let pa = tree.parent_array();
1263+
let mut pc = vec![];
1264+
for i in pa.iter() {
1265+
pc.push(*i);
1266+
}
1267+
for (i, p) in pc.iter().enumerate() {
1268+
assert_eq!(pa[i], *p);
1269+
}
1270+
} else {
1271+
panic!("Expected a tree.");
1272+
}
1273+
}
11501274
}

0 commit comments

Comments
 (0)