emote_psb\types/
collection.rs

1/*
2 * Created on Fri Dec 25 2020
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7use std::{collections::{HashMap, hash_map}, io::{Read, Seek, SeekFrom, Write}, ops::Index, slice::Iter};
8
9use crate::{PsbError, PsbErrorKind, PsbRefs};
10
11use byteorder::{ReadBytesExt, WriteBytesExt};
12use itertools::Itertools;
13
14use super::{PSB_TYPE_INTEGER_ARRAY_N, PsbValue, number::PsbNumber};
15
16#[cfg(feature = "serde")]
17use serde::{Serialize, Deserialize};
18
19#[derive(Debug, Clone, PartialEq, Eq, Hash)]
20#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
21#[cfg_attr(feature = "serde", serde(transparent))]
22pub struct PsbUintArray {
23
24    vec: Vec<u64>
25
26}
27
28impl PsbUintArray {
29
30    pub fn new() -> Self {
31        Self {
32            vec: Vec::new()
33        }
34    }
35
36    pub fn len(&self) -> usize {
37        self.vec.len()
38    }
39
40    pub fn iter(&self) -> Iter<'_, u64> {
41        self.vec.iter()
42    }
43
44    pub fn vec(&self) -> &Vec<u64> {
45        &self.vec
46    }
47
48    pub fn vec_mut(&mut self) -> &mut Vec<u64> {
49        &mut self.vec
50    }
51
52    pub fn unwrap(self) -> Vec<u64> {
53        self.vec
54    }
55
56    /// Item byte size
57    pub fn get_item_n(&self) -> u8 {
58        PsbNumber::get_uint_n(*self.vec.iter().max().unwrap_or(&0))
59    }
60
61    pub fn get_n(&self) -> u8 {
62        PsbNumber::get_uint_n(self.vec.len() as u64)
63    }
64
65    pub fn from_bytes(n: u8, stream: &mut impl Read) -> Result<(u64, PsbUintArray), PsbError> {
66        let (count_read, item_count) = PsbNumber::read_uint(n, stream)?;
67
68        let item_byte_size = stream.read_u8()? - PSB_TYPE_INTEGER_ARRAY_N;
69
70        let mut list = Vec::<u64>::new();
71
72        let mut item_total_read = 0_u64;
73        for _ in 0..item_count {
74            let (item_read, item) = PsbNumber::read_uint(item_byte_size, stream)?;
75            list.push(item as u64);
76
77            item_total_read += item_read;
78        }
79
80        Ok((count_read + item_total_read + 1, PsbUintArray::from(list)))
81    }
82
83    pub fn write_bytes(&self, stream: &mut impl Write) -> Result<u64, PsbError> {
84        let len = self.vec.len() as u64;
85
86        let count_written = PsbNumber::write_uint(self.get_n(), len, stream)? as u64;
87
88        let n = self.get_item_n();
89
90        stream.write_u8(n + PSB_TYPE_INTEGER_ARRAY_N)?;
91
92        for num in self.vec.iter() {
93            PsbNumber::write_uint(n, *num, stream)?;
94        }
95
96        Ok(1 + count_written + n as u64 * len as u64)
97    }
98
99}
100
101impl From<Vec<u64>> for PsbUintArray {
102
103    fn from(vec: Vec<u64>) -> Self {
104        Self {
105            vec
106        }
107    }
108
109}
110
111impl Index<usize> for PsbUintArray {
112    
113    type Output = u64;
114
115    fn index(&self, index: usize) -> &Self::Output {
116        &self.vec[index]
117    }
118}
119
120#[derive(Debug, PartialEq)]
121#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
122#[cfg_attr(feature = "serde", serde(transparent))]
123pub struct PsbList {
124
125    values: Vec<PsbValue>
126
127}
128
129impl PsbList {
130
131    pub fn new() -> Self {
132        Self {
133            values: Vec::new()
134        }
135    }
136
137    pub fn values(&self) -> &Vec<PsbValue> {
138        &self.values
139    }
140
141    pub fn len(&self) -> usize {
142        self.values.len()
143    }
144
145    pub fn iter(&self) -> Iter<'_, PsbValue> {
146        self.values.iter()
147    }
148
149    pub fn unwrap(self) -> Vec<PsbValue> {
150        self.values
151    }
152
153    pub fn from_bytes<T: Read + Seek>(stream: &mut T, table: &PsbRefs) -> Result<(u64, PsbList), PsbError> {
154        let (offsets_read, ref_offsets) = match PsbValue::from_bytes(stream)? {
155    
156            (read, PsbValue::IntArray(array)) => Ok((read, array)),
157
158            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
159
160        }?;
161
162        if ref_offsets.len() < 1 {
163            return Ok((offsets_read, Self::new()));
164        }
165
166        let max_offset = ref_offsets.iter().max().unwrap();
167
168        let mut values = Vec::<PsbValue>::with_capacity(ref_offsets.len());
169
170        let start = stream.seek(SeekFrom::Current(0)).unwrap();
171        let mut total_read = 0_u64;
172
173        for offset in ref_offsets.iter() {
174            stream.seek(SeekFrom::Start(start + *offset as u64))?;
175            let (read, val) = PsbValue::from_bytes_refs(stream, table)?;
176
177            values.push(val);
178
179            if max_offset == offset {
180                total_read = read + *offset as u64;
181            }
182        }
183
184        stream.seek(SeekFrom::Start(start + total_read))?;
185
186        Ok((offsets_read + total_read, Self::from(values)))
187    }
188
189    pub fn write_bytes(&self, stream: &mut impl Write, table: &PsbRefs) -> Result<u64, PsbError> {
190        let mut value_offset_cache = HashMap::<u64, &PsbValue>::new();
191
192        let mut offsets = Vec::<u64>::new();
193        let mut data_buffer = Vec::<u8>::new();
194
195        let mut total_data_written = 0_u64;
196        for value in &self.values {
197            let mut cached = false;
198            for (offset, cache_value) in &value_offset_cache {
199                if value == *cache_value {
200                    offsets.push(*offset);
201                    cached = true;
202                    break;
203                }
204            }
205
206            if !cached {
207                value_offset_cache.insert(total_data_written, &value);
208                offsets.push(total_data_written);
209
210                total_data_written += value.write_bytes_refs(&mut data_buffer, table)?;
211            }
212        }
213
214        let offset_written = PsbValue::IntArray(PsbUintArray::from(offsets)).write_bytes(stream)?;
215        stream.write_all(&data_buffer)?;
216
217        Ok(offset_written + total_data_written)
218    }
219
220    pub fn collect_strings(&self, vec: &mut Vec<String>) {
221        for child in self.values.iter() {
222            match child {
223
224                PsbValue::Object(child_obj) => {
225                    child_obj.collect_strings(vec);
226                }
227
228                PsbValue::List(child_list) => {
229                    child_list.collect_strings(vec);
230                }
231
232                PsbValue::String(string) => {
233                    if !vec.contains(string.string()) {
234                        vec.push(string.string().clone());
235                    }
236                }
237
238                _ => {}
239            }
240            
241        }
242    }
243
244    pub fn collect_names(&self, vec: &mut Vec<String>) {
245        for child in self.values.iter() {
246            match child {
247
248                PsbValue::Object(child_obj) => {
249                    child_obj.collect_names(vec);
250                }
251
252                PsbValue::List(child_list) => {
253                    child_list.collect_names(vec);
254                }
255
256                _ => {}
257            }
258            
259        }
260    }
261
262}
263
264impl From<Vec<PsbValue>> for PsbList {
265
266    fn from(values: Vec<PsbValue>) -> Self {
267        Self {
268            values
269        }
270    }
271
272}
273
274#[derive(Debug, PartialEq)]
275#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
276#[cfg_attr(feature = "serde", serde(transparent))]
277pub struct PsbObject {
278
279    // key, PsbValue Map
280    map: HashMap<String, PsbValue>
281
282}
283
284impl PsbObject {
285
286    pub fn new() -> Self {
287        Self {
288            map: HashMap::new()
289        }
290    }
291
292    pub fn len(&self) -> usize {
293        self.map.len()
294    }
295
296    pub fn get_value(&self, key: String) -> Option<&PsbValue> {
297        self.map.get(&key)
298    }
299
300    pub fn map(&self) -> &HashMap<String, PsbValue> {
301        &self.map
302    }
303
304    pub fn iter(&self) -> hash_map::Iter<'_, String, PsbValue>{
305        self.map.iter()
306    }
307
308    pub fn unwrap(self) -> HashMap<String, PsbValue> {
309        self.map
310    }
311
312    pub fn from_bytes<T: Read + Seek>(stream: &mut T, table: &PsbRefs) -> Result<(u64, PsbObject), PsbError> {
313        let (names_read, name_refs) = match PsbValue::from_bytes(stream)? {
314    
315            (read, PsbValue::IntArray(array)) => Ok((read, array)),
316
317            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
318
319        }?;
320
321        let (offsets_read, ref_offsets) = match PsbValue::from_bytes(stream)? {
322    
323            (read, PsbValue::IntArray(array)) => Ok((read, array)),
324
325            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
326
327        }?;
328
329        if name_refs.len() < 1 {
330            return Ok((names_read + offsets_read, Self::new()));
331        }
332
333        let max_offset = ref_offsets.iter().max().unwrap();
334
335        let mut map = HashMap::<String, PsbValue>::new();
336
337        let start = stream.seek(SeekFrom::Current(0)).unwrap();
338        let mut total_read = 0_u64;
339
340        for (name_ref, offset) in name_refs.iter().zip(ref_offsets.iter()) {
341            stream.seek(SeekFrom::Start(start + *offset as u64))?;
342            let (read, val) = PsbValue::from_bytes_refs(stream, table)?;
343
344            let key = table.names().get(*name_ref as usize);
345           
346            if key.is_none() {
347                return Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None));
348            }
349
350            map.insert(key.unwrap().clone(), val);
351
352            if *max_offset == *offset {
353                total_read = read + *offset as u64;
354            }
355        }
356
357        stream.seek(SeekFrom::Start(start + total_read))?;
358
359        Ok((names_read + offsets_read + total_read, Self::from(map)))
360    }
361
362    pub fn write_bytes(&self, stream: &mut impl Write, ref_table: &PsbRefs) -> Result<u64, PsbError> {
363        let mut value_offset_cache = HashMap::<u64, &PsbValue>::new();
364
365        let mut ref_cache = HashMap::<&String, u64>::new();
366
367        let mut name_refs = Vec::<u64>::new();
368        let mut offsets = Vec::<u64>::new();
369        let mut data_buffer = Vec::<u8>::new();
370
371        let mut total_data_written = 0_u64;
372
373        for name in self.map.keys().into_iter().sorted() {
374            let value = self.map.get(name).unwrap();
375
376            let name_ref = if ref_cache.contains_key(name) {
377                *ref_cache.get(name).unwrap()
378            } else {
379                match ref_table.find_name_index(name) {
380                    Some(index) => {
381                        ref_cache.insert(name, index);
382
383                        Ok(index)
384                    },
385
386                    None => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
387                }?
388            };
389
390            name_refs.push(name_ref);
391
392            let mut cached = false;
393            for (offset, cache_value) in value_offset_cache.iter() {
394                if value == *cache_value {
395                    offsets.push(*offset);
396                    cached = true;
397                    break;
398                }
399            }
400
401            if !cached {
402                value_offset_cache.insert(total_data_written, &value);
403                offsets.push(total_data_written);
404
405                total_data_written += value.write_bytes_refs(&mut data_buffer, ref_table)?;
406            }
407        }
408
409        let names_written = PsbValue::IntArray(PsbUintArray::from(name_refs)).write_bytes(stream)?;
410        let offset_written = PsbValue::IntArray(PsbUintArray::from(offsets)).write_bytes(stream)?;
411
412        stream.write_all(&data_buffer)?;
413
414        Ok(names_written + offset_written + total_data_written as u64)
415    }
416
417    pub fn collect_names(&self, vec: &mut Vec<String>) {
418        for (name, child) in self.map.iter() {
419            match child {
420
421                PsbValue::Object(child_obj) => {
422                    child_obj.collect_names(vec);
423                }
424
425                PsbValue::List(child_list) => {
426                    child_list.collect_names(vec);
427                }
428
429                _ => {}
430            }
431
432            if !vec.contains(&name) {
433                vec.push(name.clone());
434            }
435        }
436    }
437
438    pub fn collect_strings(&self, vec: &mut Vec<String>) {
439        for (_, child) in self.map.iter() {
440            match child {
441
442                PsbValue::Object(child_obj) => {
443                    child_obj.collect_strings(vec);
444                }
445
446                PsbValue::List(child_list) => {
447                    child_list.collect_strings(vec);
448                }
449
450                PsbValue::String(string) => {
451                    if !vec.contains(string.string()) {
452                        vec.push(string.string().clone());
453                    }
454                }
455
456                _ => {}
457            }
458            
459        }
460    }
461
462}
463
464impl From<HashMap<String, PsbValue>> for PsbObject {
465
466    fn from(map: HashMap<String, PsbValue>) -> Self {
467        Self {
468            map
469        }
470    }
471
472}
emote_psb\types/collection.rs

emote_psb\types/
collection.rs
