polytrack_codes/v5/
mod.rs

1#[cfg(test)]
2mod tests;
3
4use std::fmt::Display;
5
6use num_enum::TryFromPrimitive;
7
8use crate::tools::{self, prelude::*};
9use crate::{Block, Part, Track};
10
11pub const CP_IDS: [u8; 4] = [52, 65, 75, 77];
12pub const START_IDS: [u8; 4] = [5, 91, 92, 93];
13
14#[derive(Debug, PartialEq, Eq, Clone)]
15pub struct V5TrackInfo {
16    author: String,
17}
18#[derive(Debug, PartialEq, Eq, Clone)]
19pub struct V5Track {
20    pub metadata: V5TrackMetadata,
21    pub parts: Vec<V5Part>,
22}
23#[derive(Debug, PartialEq, Eq, Clone, Copy)]
24pub struct V5TrackMetadata {
25    pub env: V5Environment,
26    pub sun_dir: u8,
27
28    pub min_x: i32,
29    pub min_y: i32,
30    pub min_z: i32,
31
32    pub data_bytes: u8,
33}
34
35#[derive(TryFromPrimitive, Debug, PartialEq, Eq, Clone, Copy)]
36#[repr(u8)]
37pub enum V5Environment {
38    Summer,
39    Winter,
40    Desert,
41}
42
43#[derive(Debug, PartialEq, Eq, Clone)]
44pub struct V5Part {
45    pub id: u8,
46    pub amount: u32,
47    pub blocks: Vec<V5Block>,
48}
49
50#[derive(Debug, PartialEq, Eq, Clone)]
51pub struct V5Block {
52    pub x: u32,
53    pub y: u32,
54    pub z: u32,
55
56    pub rotation: u8,
57    pub dir: V5Direction,
58
59    pub color: u8,
60    pub cp_order: Option<u16>,
61    pub start_order: Option<u32>,
62}
63
64#[derive(TryFromPrimitive, Debug, PartialEq, Eq, Clone, Copy)]
65#[repr(u8)]
66pub enum V5Direction {
67    YPos,
68    YNeg,
69    XPos,
70    XNeg,
71    ZPos,
72    ZNeg,
73}
74
75impl Track for V5Track {
76    type Part = V5Part;
77
78    type Metadata = V5TrackMetadata;
79    fn meta(&self) -> Self::Metadata {
80        self.metadata
81    }
82    fn parts(&self) -> Vec<Self::Part> {
83        self.parts.clone()
84    }
85
86    fn decode_meta(data: &[u8], offset: &mut usize) -> Option<Self::Metadata> {
87        let env = V5Environment::try_from(read_u8(data, offset)?).ok()?;
88        let sun_dir = read_u8(data, offset)?;
89
90        let min_x = read_u32(data, offset)?.cast_signed();
91        let min_y = read_u32(data, offset)?.cast_signed();
92        let min_z = read_u32(data, offset)?.cast_signed();
93
94        let data_bytes = read_u8(data, offset)?;
95
96        Some(V5TrackMetadata {
97            env,
98            sun_dir,
99
100            min_x,
101            min_y,
102            min_z,
103
104            data_bytes,
105        })
106    }
107    fn encode_meta(&self, data: &mut Vec<u8>) {
108        data.push(self.metadata.env as u8);
109        data.push(self.metadata.sun_dir);
110        write_u32(data, self.metadata.min_x.cast_unsigned());
111        write_u32(data, self.metadata.min_y.cast_unsigned());
112        write_u32(data, self.metadata.min_z.cast_unsigned());
113        data.push(self.metadata.data_bytes);
114    }
115    fn from_data(metadata: Self::Metadata, parts: Vec<Self::Part>) -> Self {
116        Self { metadata, parts }
117    }
118
119    type TrackInfo = V5TrackInfo;
120    fn decode_track_code(track_code: &str) -> Option<(String, Self::TrackInfo, Vec<u8>)> {
121        // only use the actual data, skipping the "PolyTrack2"
122        let track_code = track_code.get(10..)?;
123        // ZLIB header 0x78DA is always encoded to `4p` and other stuff
124        let td_start = track_code.find("4p")?;
125        let track_data = track_code.get(td_start..)?;
126
127        // (base64-decode and then decompress using zlib) x2
128        let step1 = tools::decode(track_data)?;
129        let step2 = tools::decompress(&step1)?;
130        let step2_str = String::from_utf8(step2).ok()?;
131        let step3 = tools::decode(&step2_str)?;
132        let step4 = tools::decompress(&step3)?;
133
134        let name_len = *step4.first()? as usize;
135        let name = String::from_utf8(step4.get(1..=name_len)?.to_vec()).ok()?;
136
137        let author_len = *step4.get(1 + name_len)? as usize;
138        let author = String::from_utf8(
139            step4
140                .get((name_len + 2)..(name_len + author_len + 2))?
141                .to_vec(),
142        )
143        .ok()?;
144        let track_data = step4.get((name_len + author_len + 2)..)?.to_vec();
145
146        Some((name, V5TrackInfo { author }, track_data))
147    }
148    fn encode_track_code(name: String, info: Self::TrackInfo, track_data: &[u8]) -> Option<String> {
149        let mut data: Vec<u8> = Vec::new();
150
151        let mut name = name.as_bytes().to_vec();
152        data.push(name.len().try_into().ok()?);
153        data.append(&mut name);
154
155        let mut author = info.author.as_bytes().to_vec();
156        data.push(author.len().try_into().ok()?);
157        data.append(&mut author);
158
159        data.append(&mut track_data.into());
160
161        // (compress using zlib and then base62-encode) x2
162        let step1 = tools::compress_first(&data)?;
163        let step2_str = tools::encode(&step1)?;
164        let step2 = step2_str.as_bytes();
165        let step3 = tools::compress_final(step2)?;
166        let step4 = tools::encode(&step3)?;
167
168        // prepend the "PolyTrack1"
169        let track_code = String::from("PolyTrack1") + &step4;
170        Some(track_code)
171    }
172}
173impl Part for V5Part {
174    type Block = V5Block;
175    fn id(&self) -> u8 {
176        self.id
177    }
178    fn amount(&self) -> u32 {
179        self.amount
180    }
181    fn blocks(&self) -> Vec<Self::Block> {
182        self.blocks.clone()
183    }
184    fn from_data(id: u8, amount: u32, blocks: Vec<Self::Block>) -> Self {
185        Self { id, amount, blocks }
186    }
187
188    fn decode_header(data: &[u8], offset: &mut usize) -> Option<(u8, u32)> {
189        Some((read_u8(data, offset)?, read_u32(data, offset)?))
190    }
191    fn encode_header(&self, data: &mut Vec<u8>) {
192        data.push(self.id);
193        write_u32(data, self.amount);
194    }
195}
196impl Block for V5Block {
197    type Track = V5Track;
198
199    type Extra = (V5Direction, u8, Option<u16>, Option<u32>);
200    fn extra_data(&self) -> Self::Extra {
201        (self.dir, self.color, self.cp_order, self.start_order)
202    }
203
204    type Coord = u32;
205    fn pos(&self) -> (Self::Coord, Self::Coord, Self::Coord) {
206        (self.x, self.y, self.z)
207    }
208    fn rot(&self) -> u8 {
209        self.rotation
210    }
211
212    fn decode(
213        data: &[u8],
214        offset: &mut usize,
215        id: u8,
216        meta: <Self::Track as Track>::Metadata,
217    ) -> Option<Self> {
218        let data_bytes = meta.data_bytes;
219        let x_bytes = data_bytes & 3;
220        let y_bytes = (data_bytes >> 2) & 3;
221        let z_bytes = (data_bytes >> 4) & 3;
222
223        let mut x = 0;
224        for i in 0..x_bytes {
225            x |= u32::from(*data.get(*offset + (i as usize))?) << (8 * i);
226        }
227        *offset += x_bytes as usize;
228
229        let mut y = 0;
230        for i in 0..y_bytes {
231            y |= u32::from(*data.get(*offset + (i as usize))?) << (8 * i);
232        }
233        *offset += y_bytes as usize;
234
235        let mut z = 0;
236        for i in 0..z_bytes {
237            z |= u32::from(*data.get(*offset + (i as usize))?) << (8 * i);
238        }
239        *offset += z_bytes as usize;
240
241        let rotation = read_u8(data, offset)?;
242        if rotation > 3 {
243            return None;
244        }
245        let dir = V5Direction::try_from(read_u8(data, offset)?).ok()?;
246        let color = read_u8(data, offset)?;
247        // no custom color support for now
248        if color > 3 && color < 32 && color > 40 {
249            return None;
250        }
251
252        let cp_order = if CP_IDS.contains(&id) {
253            Some(read_u16(data, offset)?)
254        } else {
255            None
256        };
257        let start_order = if START_IDS.contains(&id) {
258            Some(read_u32(data, offset)?)
259        } else {
260            None
261        };
262
263        Some(Self {
264            x,
265            y,
266            z,
267
268            rotation,
269            dir,
270
271            color,
272            cp_order,
273            start_order,
274        })
275    }
276    fn encode(&self, data: &mut Vec<u8>, meta: <Self::Track as Track>::Metadata) {
277        let data_bytes = meta.data_bytes;
278        let x_bytes = data_bytes & 3;
279        let y_bytes = (data_bytes >> 2) & 3;
280        let z_bytes = (data_bytes >> 4) & 3;
281
282        match x_bytes {
283            1 => write_u8(data, self.x),
284            2 => write_u16(data, self.x),
285            3 => write_u24(data, self.x),
286            4 => write_u32(data, self.x),
287            _ => {}
288        }
289        match y_bytes {
290            1 => write_u8(data, self.y),
291            2 => write_u16(data, self.y),
292            3 => write_u24(data, self.y),
293            4 => write_u32(data, self.y),
294            _ => {}
295        }
296        match z_bytes {
297            1 => write_u8(data, self.z),
298            2 => write_u16(data, self.z),
299            3 => write_u24(data, self.z),
300            4 => write_u32(data, self.z),
301            _ => {}
302        }
303        data.push(self.rotation);
304        data.push(self.dir as u8);
305        data.push(self.color);
306        if let Some(cp_order) = self.cp_order {
307            write_u16(data, cp_order.into());
308        }
309        if let Some(start_order) = self.start_order {
310            write_u32(data, start_order);
311        }
312    }
313}
314
315impl Display for V5Environment {
316    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
317        match self {
318            Self::Summer => write!(f, "Summer"),
319            Self::Winter => write!(f, "Winter"),
320            Self::Desert => write!(f, "Desert"),
321        }
322    }
323}
polytrack_codes/v5/mod.rs

polytrack_codes/v5/
mod.rs
