generic_octree/
octree.rs

1use hashbrown::HashMap;
2use hashbrown::HashSet;
3
4use crate::{LocCode, OctreeNode, Orientation, AABB};
5use rayon::prelude::*;
6
7#[cfg(feature = "serialize")]
8use serde::{de::DeserializeOwned, Deserialize, Serialize};
9
10#[cfg(feature = "serialize")]
11use flate2::Compression;
12
13#[cfg(feature = "serialize")]
14use flate2::write::{ZlibDecoder, ZlibEncoder};
15
16#[cfg(feature = "serialize")]
17use std::{io::prelude::*, path::Path};
18
19#[derive(Debug)]
20#[cfg_attr(feature = "serialize", derive(Serialize, Deserialize))]
21pub struct Octree<L: LocCode, D: Send + Sync> {
22    pub content: HashMap<L, OctreeNode<D>>,
23    max_depth: u32,
24}
25
26#[cfg(feature = "dot_tree")]
27impl<L, D> Octree<L, D>
28where
29    L: LocCode + Serialize + DeserializeOwned,
30    D: Send + Sync + Serialize + DeserializeOwned,
31{
32    /// Load from voxel octree from files
33    /// TODO: Add better error management
34    pub fn load_from_file<P: AsRef<Path>>(path_ref: P) -> Result<Self, std::io::Error> {
35        let path = path_ref.as_ref();
36        match path.extension() {
37            Some(x) => match x.to_str() {
38                #[cfg(feature = "dot_tree")]
39                Some("tree") => {
40                    let file = std::fs::File::open(path)?;
41                    let contents: Vec<u8> = file.bytes().filter_map(Result::ok).collect();
42                    let mut decoder = ZlibDecoder::new(Vec::new());
43                    decoder.write_all(&contents)?;
44                    let contents = decoder.finish()?;
45                    let tree = bincode::deserialize_from::<&[u8], Self>(&contents).unwrap();
46                    Ok(tree)
47                }
48                _ => Err(std::io::Error::from(std::io::ErrorKind::InvalidData)),
49            },
50            None => Err(std::io::Error::from(std::io::ErrorKind::InvalidInput)),
51        }
52    }
53
54    /// Save octree to file
55    pub fn save_to_file<P: AsRef<Path>>(&self, path_ref: P) -> Result<(), std::io::Error> {
56        let path = path_ref.as_ref();
57        let binary = bincode::serialize(self).unwrap();
58        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
59        encoder.write_all(&binary)?;
60        std::fs::write(path, encoder.finish()?)
61    }
62}
63
64impl<'a, L, D> Octree<L, D>
65where
66    L: LocCode,
67    D: Send + Sync,
68{
69    /// Get the size of an octree
70    pub fn size(&self) -> usize {
71        self.content.len()
72    }
73}
74
75impl<'a, T, D> Octree<T, D>
76where
77    T: LocCode,
78    D: Copy + PartialEq + Send + Sync,
79{
80    /// Create a new Octree
81    pub fn new(max_depth: u32) -> Self {
82        let content = HashMap::default();
83        Self { content, max_depth }
84    }
85
86    /// Create an Octree with given pre-allocated space.
87    pub fn with_capacity(max_depth: u32, size: usize) -> Self {
88        let content = HashMap::with_capacity(size);
89        Self { content, max_depth }
90    }
91
92    pub fn depth(&self) -> u32 {
93        let keys = self.content.keys();
94        keys.max().unwrap_or(&T::root()).get_level()
95    }
96
97    /// Return a tree node a node.
98    pub fn lookup(&self, loc_code: T) -> Option<&OctreeNode<D>> {
99        self.content.get(&loc_code)
100    }
101
102    /// Insert a tree node.
103    pub fn insert(&mut self, location: T, node: OctreeNode<D>) -> T {
104        self.content.insert(location, node);
105        location
106    }
107
108    pub fn remove_node(&mut self, loc_code: T) {
109        self.content.remove(&loc_code);
110    }
111
112    /// Merge an AABB into the tree
113    pub fn merge(&mut self, aabb: AABB, data: D) {
114        let mut codes: Vec<T> = self
115            .merge_inner(aabb, data, (0.5, 0.5, 0.5), 1, T::root())
116            .into_iter()
117            .collect();
118        while !codes.is_empty() {
119            codes.sort();
120            codes.reverse();
121            codes = codes
122                .into_iter()
123                .filter_map(|code| self.assemble(code))
124                .filter(|code| *code > T::zero())
125                .collect::<HashSet<T>>()
126                .into_iter()
127                .collect();
128        }
129    }
130
131    fn assemble(&mut self, code: T) -> Option<T> {
132        let datas = (0_u8..8_u8)
133            .map(|number| (code << T::three()) | T::from(number))
134            .filter_map(|loc_code| self.lookup(loc_code))
135            .map(|node| node.data)
136            .collect::<Vec<D>>();
137        if datas.len() != 8 {
138            None
139        } else {
140            let elem = datas[0];
141            let is_same = datas
142                .into_iter()
143                .fold((true, elem), |acc, elem| (acc.0 && acc.1 == elem, acc.1))
144                .0;
145            if !is_same {
146                None
147            } else {
148                (0_u8..8_u8)
149                    .map(|number| (code << T::three()) | T::from(number))
150                    .for_each(|code| self.remove_node(code));
151                self.insert(code, OctreeNode::new(elem));
152                Some(code >> T::three())
153            }
154        }
155    }
156
157    /// Transform an Octree of data D into an Octree of data U, provided that
158    /// U implement From<D>
159    pub fn transform<U: From<D> + Send + Sync>(self) -> Octree<T, U> {
160        Octree {
161            content: self
162                .content
163                .into_par_iter()
164                .map(|(loc_code, data)| (loc_code, data.transform::<U>()))
165                .collect::<HashMap<T, OctreeNode<U>>>(),
166            max_depth: self.max_depth,
167        }
168    }
169
170    /// tree.transform_fn(Rgb::from_hex);
171    pub fn transform_fn<U: Send + Sync, F: Fn(D) -> U + Sync>(self, function: F) -> Octree<T, U> {
172        Octree {
173            content: self
174                .content
175                .into_par_iter()
176                .map(|(loc_code, data)| (loc_code, data.transform_fn(&function)))
177                .collect::<HashMap<T, OctreeNode<U>>>(),
178            max_depth: self.max_depth,
179        }
180    }
181
182    /// tree.transform_fn(Rgb::from_hex);
183    pub fn transform_nodes_fn<U: Send + Sync, F: Fn(T, OctreeNode<D>) -> OctreeNode<U> + Sync>(
184        self,
185        function: F,
186    ) -> Octree<T, U> {
187        Octree {
188            content: self
189                .content
190                .into_par_iter()
191                .map(|(loc_code, data)| (loc_code, function(loc_code, data)))
192                .collect::<HashMap<T, OctreeNode<U>>>(),
193            max_depth: self.max_depth,
194        }
195    }
196
197    /// Internal function for recursively merging AABB.
198    /// Returns a HashSet containing all the node that are affected by the merging, not all new nodes
199    /// These affected nodes can be scheduled to merge data outside here
200    fn merge_inner(
201        &mut self,
202        aabb: AABB,
203        data: D,
204        center: (f64, f64, f64),
205        depth: u32,
206        loc_code: T,
207    ) -> HashSet<T> {
208        let blocks = aabb.explode(center);
209        let max_depth = self.max_depth;
210
211        let (fitting, subdivisables): (Vec<AABB>, Vec<AABB>) = blocks
212            .into_iter()
213            .partition(|aabb| aabb.fit_in(depth, max_depth));
214
215        let mut codes: Vec<T> = fitting
216            .into_iter()
217            .map(|elem| {
218                self.insert(
219                    loc_code << T::three() | elem.orientation,
220                    OctreeNode::new(data),
221                )
222            })
223            .map(|loc_code| loc_code >> T::three())
224            .collect();
225
226        codes.extend(if depth != self.max_depth {
227            subdivisables
228                .into_iter()
229                .map(|aabb| {
230                    let new_loc_code = (loc_code << T::three()) | aabb.orientation;
231                    let new_center = aabb.orientation.make_new_center(new_loc_code, center);
232                    self.merge_inner(
233                        aabb.with_orientation(Orientation::N),
234                        data,
235                        new_center,
236                        depth + 1,
237                        new_loc_code,
238                    )
239                })
240                .flatten()
241                .collect()
242        } else {
243            Vec::<T>::default()
244        });
245
246        codes.into_iter().collect::<HashSet<T>>()
247    }
248}
249
250impl<'a, L, D> Octree<L, D>
251where
252    L: LocCode,
253    D: Send + Sync,
254{
255    #[cfg(feature = "vox")]
256    pub fn from_dotvox<U: AsRef<str>>(
257        path: U,
258        max_depth: u32,
259        optimal: crate::dot_vox::ConversionType,
260    ) -> Result<Vec<Octree<L, u32>>, &'static str> {
261        let vox = dot_vox::load(path.as_ref())?;
262        let octrees: Vec<Octree<L, u32>> = crate::dot_vox::vox_to_octrees(vox, max_depth, optimal);
263        Ok(octrees)
264    }
265}
generic_octree/octree.rs

generic_octree/
octree.rs
