polytrack_codes/v5/
mod.rs

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

polytrack_codes/v5/
mod.rs
