xp3/
index_set.rs

1/*
2 * Created on Tue Dec 15 2020
3 *
4 * Copyright (c) storycraft. Licensed under the Apache Licence 2.0.
5 */
6
7use std::{collections::{HashMap, hash_map::{Iter, Values}}, convert::TryFrom, io::{Cursor, Read, Write}};
8
9use byteorder::LittleEndian;
10use flate2::{Compression, read::{ZlibDecoder, ZlibEncoder}};
11
12use crate::XP3_INDEX_CONTINUE;
13
14use super::{XP3Error, XP3ErrorKind, XP3_INDEX_FILE_IDENTIFIER, index::{XP3Index, file::XP3FileIndex}};
15
16use byteorder::{WriteBytesExt, ReadBytesExt};
17
18#[derive(Debug, Copy, Clone)]
19#[repr(u8)]
20pub enum XP3IndexCompression {
21
22    UnCompressed = 0,
23
24    Compressed = 1
25
26}
27
28impl TryFrom<u8> for XP3IndexCompression {
29
30    type Error = XP3Error;
31
32    fn try_from(val: u8) -> Result<Self, Self::Error> {
33        match val {
34            val if val == Self::UnCompressed as u8 => Ok(Self::UnCompressed),
35            val if val == Self::Compressed as u8 => Ok(Self::Compressed),
36            _ => Err(XP3Error::new(XP3ErrorKind::InvalidFileIndexHeader, None)),
37        }
38    }
39}
40
41#[derive(Debug)]
42pub struct XP3IndexSet {
43
44    compression: XP3IndexCompression,
45
46    extra: Vec<XP3Index>,
47
48    file_map: HashMap<String, XP3FileIndex>
49
50}
51
52impl XP3IndexSet {
53
54    pub fn new(
55        compression: XP3IndexCompression,
56        extra: Vec<XP3Index>,
57        file_map: HashMap<String, XP3FileIndex>
58    ) -> Self {
59        Self {
60            compression,
61            extra,
62            file_map
63        }
64    }
65
66    pub fn compression(&self) -> XP3IndexCompression {
67        self.compression
68    }
69
70    pub fn set_compression(&mut self, compression: XP3IndexCompression) {
71        self.compression = compression;
72    }
73
74    pub fn indices(&self) -> Values<String, XP3FileIndex> {
75        self.file_map.values()
76    }
77
78    pub fn extra(&self) -> &Vec<XP3Index> {
79        &self.extra
80    }
81
82    pub fn extra_mut(&mut self) -> &mut Vec<XP3Index> {
83        &mut self.extra
84    }
85
86    pub fn file_map(&self) -> &HashMap<String, XP3FileIndex> {
87        &self.file_map
88    }
89
90    pub fn file_map_mut(&mut self) -> &mut HashMap<String, XP3FileIndex> {
91        &mut self.file_map
92    }
93
94    pub fn entries(&self) -> Iter<String, XP3FileIndex> {
95        self.file_map.iter()
96    }
97
98    pub fn get(&self, name: &String) -> Option<&XP3FileIndex> {
99        self.file_map.get(name)
100    }
101
102    /// Read xp3 file index set from stream.
103    pub fn from_bytes(stream: &mut impl Read) -> Result<(u64, Self), XP3Error> {
104        let mut raw_flag: u8 = stream.read_u8()?;
105        // Skip continue part
106        while raw_flag == XP3_INDEX_CONTINUE {
107            raw_flag = stream.read_u8()?;
108        }
109
110        let flag = XP3IndexCompression::try_from(raw_flag)?;
111
112        let index_data: Vec<u8>;
113        let index_size: u64;
114        let raw_index_read: u64;
115        match flag {
116            XP3IndexCompression::UnCompressed => {
117                index_size = stream.read_u64::<LittleEndian>()?;
118
119                let mut data = Vec::<u8>::with_capacity(index_size as usize);
120                stream.take(index_size).read_to_end(&mut data)?;
121
122                index_data = data;
123                raw_index_read = index_size;
124            },
125
126            XP3IndexCompression::Compressed => {
127                let compressed_size = stream.read_u64::<LittleEndian>()?;
128                index_size = stream.read_u64::<LittleEndian>()?;
129
130                let mut compressed_buffer = Vec::<u8>::with_capacity(compressed_size as usize);
131                stream.take(compressed_size).read_to_end(&mut compressed_buffer)?;
132
133                let decoder = ZlibDecoder::new(&compressed_buffer[..]);
134    
135                let mut data = Vec::new();
136                decoder.take(index_size).read_to_end(&mut data)?;
137    
138                index_data = data;
139                raw_index_read = compressed_size;
140            }
141        }
142
143        let mut index_buffer = Cursor::new(index_data);
144
145        let mut file_map: HashMap<String, XP3FileIndex> = HashMap::new();
146
147        let mut extra: Vec<XP3Index> = Vec::new();
148
149        let mut index_read: u64 = 0;
150        while index_read < index_size {
151            let (read, index) = XP3Index::from_bytes(&mut index_buffer)?;
152            
153            match index.identifier {
154
155                XP3_INDEX_FILE_IDENTIFIER => {
156                    let (_, file_index) = XP3FileIndex::from_bytes(index.data.len() as u64, &mut Cursor::new(&index.data))?;
157                    file_map.insert(file_index.info().name().clone(), file_index);
158                }
159
160                _ => {
161                    extra.push(index);
162                }
163            }
164
165            
166            index_read += read;
167        }
168
169        Ok((raw_index_read + index_size, Self::new(flag, extra, file_map)))
170    }
171
172    /// Write xp3 file index set to stream.
173    pub fn write_bytes(&mut self, stream: &mut impl Write) -> Result<u64, XP3Error> {
174        stream.write_u8(self.compression as u8)?;
175
176        let mut index_buffer = Vec::<u8>::new();
177
178        for extra_index in self.extra.iter_mut() {
179            extra_index.write_bytes(&mut index_buffer)?;
180        }
181
182        for index in self.file_map.values() {
183            let mut buffer = Vec::<u8>::new();
184            index.write_bytes(&mut buffer)?;
185            XP3Index { identifier: XP3_INDEX_FILE_IDENTIFIER, data: buffer }.write_bytes(&mut index_buffer)?;
186        }
187
188        match self.compression {
189            XP3IndexCompression::UnCompressed => {
190                let size = index_buffer.len() as u64;
191
192                stream.write_u64::<LittleEndian>(size)?;
193                
194                stream.write_all(&index_buffer)?;
195
196                Ok(9 + size)
197            },
198
199            XP3IndexCompression::Compressed => {
200                let mut compressed_data = Vec::<u8>::new();
201
202                // Reset read stream
203                let mut encoder = ZlibEncoder::new(Cursor::new(index_buffer), Compression::fast());
204
205                encoder.read_to_end(&mut compressed_data)?;
206
207                stream.write_u64::<LittleEndian>(compressed_data.len() as u64)?;
208                stream.write_u64::<LittleEndian>(encoder.get_ref().get_ref().len() as u64)?;
209
210                stream.write_all(&mut compressed_data)?;
211
212                Ok(17 + compressed_data.len() as u64)   
213            }
214        }
215    }
216
217}
xp3/index_set.rs

xp3/
index_set.rs
