emote_psb\types/
binary_tree.rs

1/*
2 * Created on Sun Dec 27 2020
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7use std::{collections::{HashMap, hash_map}, io::{Read, Seek, Write}, slice::Iter};
8
9use crate::{PsbError, PsbErrorKind, internal::SafeIndexVec};
10
11use super::{PsbValue, collection::PsbUintArray};
12
13/// Binary tree
14pub struct PsbBinaryTree {
15
16    pub list: Vec<Vec<u8>>
17
18}
19
20impl PsbBinaryTree {
21
22    pub fn new() -> Self {
23        Self {
24            list: Vec::new()
25        }
26    }
27
28    pub fn list(&self) -> &Vec<Vec<u8>> {
29        &self.list
30    }
31
32    pub fn list_mut(&mut self) -> &mut Vec<Vec<u8>> {
33        &mut self.list
34    }
35
36    pub fn len(&self) -> usize {
37        self.list.len()
38    }
39
40    pub fn iter(&self) -> Iter<'_, Vec<u8>> {
41        self.list.iter()
42    }
43
44    pub fn unwrap(self) -> Vec<Vec<u8>> {
45        self.list
46    }
47
48    pub fn from_bytes<T: Read + Seek>(stream: &mut T) -> Result<(u64, Self), PsbError> {
49        let (offsets_read, offsets) = match PsbValue::from_bytes(stream)? {
50    
51            (read, PsbValue::IntArray(array)) => Ok((read, array)),
52
53            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
54
55        }?;
56        let (tree_read, tree) = match PsbValue::from_bytes(stream)? {
57
58            (read, PsbValue::IntArray(array)) => Ok((read, array)),
59
60            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
61
62        }?;
63        let (indexes_read, indexes) = match PsbValue::from_bytes(stream)? {
64
65            (read, PsbValue::IntArray(array)) => Ok((read, array)),
66
67            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
68
69        }?;
70
71        // Unwrap all to vec
72        let offsets = offsets.unwrap();
73        let tree = tree.unwrap();
74        let indexes = indexes.unwrap();
75
76        let mut list = Vec::<Vec<u8>>::with_capacity(indexes.len());
77
78        for index in indexes {
79            let mut buffer = Vec::<u8>::new();
80            
81            let mut id = tree[index as usize];
82
83            while id != 0 {
84                // travel to child tree
85                let next = tree[id as usize];
86
87                // get values from offsets
88                let decoded = id - offsets[next as usize];
89                
90                id = next;
91
92                buffer.push(decoded as u8);
93            }
94
95            buffer.reverse();
96            list.push(buffer);
97        }
98
99        Ok((offsets_read + tree_read + indexes_read, Self::from(list)))
100    }
101
102    pub fn build_tree(&self) -> TreeNode {
103        let mut root = TreeNode::new();
104        
105        for data in &self.list {
106            let mut last_node = &mut root;
107
108            for byte in data {
109                last_node = last_node.get_or_insert_mut(*byte);
110            }
111
112            last_node.get_or_insert(0);
113        }
114        
115        root
116    }
117
118    pub fn write_bytes(&self, stream: &mut impl Write) -> Result<u64, PsbError> {
119        let mut root = self.build_tree();
120
121        let mut offsets = SafeIndexVec::new();
122        let mut tree = SafeIndexVec::new();
123        let mut indexes = SafeIndexVec::new();
124
125        offsets.push(1);
126        self.make_sub_tree(&mut root, Vec::new(), &mut offsets, &mut tree, &mut indexes);
127
128        let offsets_written = PsbValue::IntArray(PsbUintArray::from(offsets.into_inner())).write_bytes(stream)?;
129        let tree_written = PsbValue::IntArray(PsbUintArray::from(tree.into_inner())).write_bytes(stream)?;
130        let indexes_written = PsbValue::IntArray(PsbUintArray::from(indexes.into_inner())).write_bytes(stream)?;
131
132        Ok(offsets_written + tree_written + indexes_written)
133    }
134
135    // Returns last node 
136    fn make_sub_tree(
137        &self,
138        current_node: &mut TreeNode,
139        value: Vec<u8>,
140        offsets: &mut SafeIndexVec<u64>,
141        tree: &mut SafeIndexVec<u64>,
142        indexes: &mut SafeIndexVec<u64>
143    ) {
144        let min_value = *current_node.min_value().unwrap_or(&0);
145        let begin_pos = current_node.begin_pos;
146        let current_id = current_node.id;
147
148        // make_tree
149        for (child_value, child) in current_node.iter_mut() {
150            let id = if current_id == 0 || min_value < 1 {
151                *child_value as u64 + offsets.get(current_id as usize).unwrap()
152            } else {
153                (*child_value - min_value) as u64 + begin_pos
154            };
155
156            tree.set(id as usize, current_id);
157            child.id = id;
158        }
159
160        for (child_value, child) in current_node.iter_mut() {
161            let child_max = *child.max_value().unwrap_or(&0) as usize;
162            let child_min = *child.min_value().unwrap_or(&0) as usize;
163
164            let pos = {
165                let len = tree.len();
166                if len > child_max {
167                    len
168                } else {
169                    tree.set(child_max, 0);
170
171                    tree.len()
172                }
173            };
174
175            let count = child_max - child_min;
176            let end = pos + count;
177            
178            tree.set(end, 0);
179
180            if *child_value == 0 {
181                let index = self.list.iter().position(|val| val.eq(&value)).unwrap() as u64;
182                offsets.set(child.id as usize, index);
183                indexes.set(index as usize, child.id);
184            } else {
185                let offset = (pos - child_min) as u64;
186                offsets.set(child.id as usize, offset);
187                child.begin_pos = pos as u64;
188            }
189        }
190
191        for (child_value, child) in current_node.iter_mut() {
192            let mut value = value.clone();
193            value.push(*child_value);
194            self.make_sub_tree(child, value, offsets, tree, indexes);
195        }
196    }
197
198}
199
200impl From<Vec<Vec<u8>>> for PsbBinaryTree {
201
202    fn from(list: Vec<Vec<u8>>) -> Self {
203        Self {
204            list
205        }
206    }
207
208}
209
210#[derive(Debug)]
211pub struct TreeNode {
212
213    /// Children value, node
214    children: HashMap<u8, TreeNode>,
215    
216    pub begin_pos: u64,
217    pub id: u64
218
219}
220
221impl TreeNode {
222
223    pub fn new() -> Self {
224        Self {
225            children: HashMap::new(),
226            id: 0,
227            begin_pos: 0
228        }
229    }
230
231    pub fn min_value(&self) -> Option<&u8> {
232        self.children.keys().min()
233    }
234
235    pub fn max_value(&self) -> Option<&u8> {
236        self.children.keys().max()
237    }
238
239    pub fn iter(&self) -> hash_map::Iter<u8, Self> {
240        self.children.iter()
241    }
242
243    pub fn iter_mut(&mut self) -> hash_map::IterMut<u8, Self> {
244        self.children.iter_mut()
245    }
246
247    pub fn len(&self) -> usize {
248        self.children.len()
249    }
250
251    pub fn get(&self, value: u8) -> Option<&Self> {
252        self.children.get(&value)
253    }
254
255    pub fn get_mut(&mut self, value: u8) -> Option<&mut Self> {
256        self.children.get_mut(&value)
257    }
258
259    pub fn get_or_insert(&mut self, value: u8) -> &Self {
260        if !self.children.contains_key(&value) {
261            let new_node = Self::new();
262
263            self.children.insert(value, new_node);
264        }
265
266        self.children.get(&value).unwrap()
267    }
268
269    pub fn get_or_insert_mut(&mut self, value: u8) -> &mut Self {
270        if !self.children.contains_key(&value) {
271            let new_node = Self::new();
272
273            self.children.insert(value, new_node);
274        }
275
276        self.children.get_mut(&value).unwrap()
277    }
278
279}
emote_psb\types/binary_tree.rs

emote_psb\types/
binary_tree.rs
