use crate::{
pack,
pack::tree::{Item, Tree},
};
impl<T> Tree<T> {
#[allow(unsafe_code)]
unsafe fn put_data(&self, index: usize, data: T) {
let items_mut: &mut Vec<Item<T>> = &mut *(self.items.get());
items_mut.get_unchecked_mut(index).data = data;
}
#[allow(unsafe_code)]
unsafe fn offset(&self, index: usize) -> u64 {
let items: &Vec<Item<T>> = &*(self.items.get());
items.get_unchecked(index).offset
}
#[allow(unsafe_code)]
unsafe fn entry_slice(&self, index: usize) -> std::ops::Range<u64> {
let items: &Vec<Item<T>> = &*(self.items.get());
let start = items.get_unchecked(index).offset;
let end = items
.get(index + 1)
.map(|e| e.offset)
.or(self.pack_entries_end)
.expect("traversal(…) to have set this value (BUG)");
start..end
}
#[allow(unsafe_code)]
unsafe fn take_entry(&self, index: usize) -> (T, Vec<usize>)
where
T: Default,
{
let items_mut: &mut Vec<Item<T>> = &mut *(self.items.get());
let item = items_mut.get_unchecked_mut(index);
let children = std::mem::take(&mut item.children);
let data = std::mem::take(&mut item.data);
(data, children)
}
#[allow(unsafe_code)]
unsafe fn from_iter_take_entry_if_root(&self, index: usize) -> Option<(T, Vec<usize>)>
where
T: Default,
{
let items_mut: &mut Vec<Item<T>> = &mut *(self.items.get());
let item = items_mut.get_unchecked_mut(index);
if item.is_root {
let children = std::mem::take(&mut item.children);
let data = std::mem::take(&mut item.data);
Some((data, children))
} else {
None
}
}
}
impl<T> Tree<T> {
pub fn iter_root_chunks(&mut self, chunk_size: usize) -> Chunks<'_, T> {
Chunks {
tree: self,
chunk_size,
cursor: 0,
}
}
}
pub struct Node<'a, T> {
tree: &'a Tree<T>,
index: usize,
children: Vec<usize>,
pub data: T,
}
impl<'a, T> Node<'a, T>
where
T: Default,
{
pub fn offset(&self) -> u64 {
#[allow(unsafe_code)]
unsafe {
self.tree.offset(self.index)
}
}
pub fn entry_slice(&self) -> pack::data::EntryRange {
#[allow(unsafe_code)]
unsafe {
self.tree.entry_slice(self.index)
}
}
pub fn store_changes_then_into_child_iter(self) -> impl Iterator<Item = Node<'a, T>> {
#[allow(unsafe_code)]
unsafe {
self.tree.put_data(self.index, self.data)
};
let Self { tree, children, .. } = self;
children.into_iter().map(move |index| {
#[allow(unsafe_code)]
let (data, children) = unsafe { tree.take_entry(index) };
Node {
tree,
data,
children,
index,
}
})
}
}
pub struct Chunks<'a, T> {
tree: &'a Tree<T>,
chunk_size: usize,
cursor: usize,
}
impl<'a, T> Iterator for Chunks<'a, T>
where
T: Default,
{
type Item = Vec<Node<'a, T>>;
fn next(&mut self) -> Option<Self::Item> {
if self.cursor == self.tree.one_past_last_seen_root {
return None;
}
let mut items_remaining = self.chunk_size;
let mut res = Vec::with_capacity(self.chunk_size);
while items_remaining > 0 && self.cursor < self.tree.one_past_last_seen_root {
#[allow(unsafe_code)]
if let Some((data, children)) = unsafe { self.tree.from_iter_take_entry_if_root(self.cursor) } {
res.push(Node {
tree: self.tree,
index: self.cursor,
children,
data,
});
items_remaining -= 1;
}
self.cursor += 1;
}
if res.is_empty() {
None
} else {
Some(res)
}
}
}