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        last_modified: None,
101        track_data,
102    })
103}
104
105#[must_use]
106/// Encodes the given track struct into a track code.
107/// Returns [`None`] if something failed in the process.
108///
109/// Output might differ slightly from Polytrack's output
110/// because of Zlib shenanigans, but is still compatible.
111pub fn encode_track_code(track: &Track) -> Option<String> {
112    let mut data: Vec<u8> = Vec::new();
113
114    let mut name = track.name.as_bytes().to_vec();
115    data.push(name.len().try_into().ok()?);
116    data.append(&mut name);
117
118    if let Some(author) = &track.author {
119        let mut author = author.as_bytes().to_vec();
120        data.push(author.len().try_into().ok()?);
121        data.append(&mut author);
122    } else {
123        data.push(0);
124    }
125
126    data.append(&mut track.track_data.clone());
127
128    // (compress using zlib and then base64-encode) x2
129    let step1 = tools::compress_first(&data)?;
130    let step2_str = tools::encode(&step1)?;
131    let step2 = step2_str.as_bytes();
132    let step3 = tools::compress_final(step2)?;
133    let step4 = tools::encode(&step3)?;
134
135    // prepend the "PolyTrack1"
136    let track_code = String::from("PolyTrack1") + &step4;
137    Some(track_code)
138}
139
140#[must_use]
141/// Decodes the (raw binary) track data into a struct
142/// representing everything that is in the data.
143///
144/// Fields of all involved structs correspond exactly to how
145/// the data is stored in Polytrack itself.
146/// Returns [`None`] if the data is not valid track data.
147pub fn decode_track_data(data: &[u8]) -> Option<TrackInfo> {
148    let mut offset = 0;
149
150    let env = Environment::try_from(read_u8(data, &mut offset)?).ok()?;
151    let sun_dir = read_u8(data, &mut offset)?;
152
153    let min_x = read_u32(data, &mut offset)?.cast_signed();
154    let min_y = read_u32(data, &mut offset)?.cast_signed();
155    let min_z = read_u32(data, &mut offset)?.cast_signed();
156
157    let data_bytes = read_u8(data, &mut offset)?;
158    let x_bytes = data_bytes & 3;
159    let y_bytes = (data_bytes >> 2) & 3;
160    let z_bytes = (data_bytes >> 4) & 3;
161
162    let mut parts = Vec::new();
163    while offset < data.len() {
164        let id = read_u8(data, &mut offset)?;
165        let amount = read_u32(data, &mut offset)?;
166
167        let mut blocks = Vec::new();
168        for _ in 0..amount {
169            let mut x = 0;
170            for i in 0..x_bytes {
171                x |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
172            }
173            offset += x_bytes as usize;
174
175            let mut y = 0;
176            for i in 0..y_bytes {
177                y |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
178            }
179            offset += y_bytes as usize;
180
181            let mut z = 0;
182            for i in 0..z_bytes {
183                z |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
184            }
185            offset += z_bytes as usize;
186
187            let rotation = read_u8(data, &mut offset)?;
188            if rotation > 3 {
189                return None;
190            }
191            let dir = Direction::try_from(read_u8(data, &mut offset)?).ok()?;
192            let color = read_u8(data, &mut offset)?;
193            // no custom color support for now
194            if color > 3 && color < 32 && color > 40 {
195                return None;
196            }
197
198            let cp_order = if CP_IDS.contains(&id) {
199                Some(read_u16(data, &mut offset)?)
200            } else {
201                None
202            };
203            let start_order = if START_IDS.contains(&id) {
204                Some(read_u32(data, &mut offset)?)
205            } else {
206                None
207            };
208
209            blocks.push(Block {
210                x,
211                y,
212                z,
213
214                rotation,
215                dir,
216
217                color,
218                cp_order,
219                start_order,
220            });
221        }
222        parts.push(Part { id, amount, blocks });
223    }
224
225    Some(TrackInfo {
226        env,
227        sun_dir,
228
229        min_x,
230        min_y,
231        min_z,
232
233        data_bytes,
234        parts,
235    })
236}
237
238#[must_use]
239/// Encodes the `TrackInfo` struct into raw binary data.
240pub fn encode_track_data(track_info: &TrackInfo) -> Option<Vec<u8>> {
241    let mut data = Vec::new();
242
243    data.push(track_info.env as u8);
244    data.push(track_info.sun_dir);
245    write_u32(&mut data, track_info.min_x.cast_unsigned());
246    write_u32(&mut data, track_info.min_y.cast_unsigned());
247    write_u32(&mut data, track_info.min_z.cast_unsigned());
248    data.push(track_info.data_bytes);
249    let x_bytes = track_info.data_bytes & 3;
250    let y_bytes = (track_info.data_bytes >> 2) & 3;
251    let z_bytes = (track_info.data_bytes >> 4) & 3;
252    for part in &track_info.parts {
253        data.push(part.id);
254        write_u32(&mut data, part.amount);
255        for block in &part.blocks {
256            match x_bytes {
257                1 => write_u8(&mut data, block.x),
258                2 => write_u16(&mut data, block.x),
259                3 => write_u24(&mut data, block.x),
260                4 => write_u32(&mut data, block.x),
261                _ => {}
262            }
263            match y_bytes {
264                1 => write_u8(&mut data, block.y),
265                2 => write_u16(&mut data, block.y),
266                3 => write_u24(&mut data, block.y),
267                4 => write_u32(&mut data, block.y),
268                _ => {}
269            }
270            match z_bytes {
271                1 => write_u8(&mut data, block.z),
272                2 => write_u16(&mut data, block.z),
273                3 => write_u24(&mut data, block.z),
274                4 => write_u32(&mut data, block.z),
275                _ => {}
276            }
277            data.push(block.rotation);
278            data.push(block.dir as u8);
279            data.push(block.color);
280            if let Some(cp_order) = block.cp_order {
281                write_u16(&mut data, cp_order.into());
282            }
283            if let Some(start_order) = block.start_order {
284                write_u32(&mut data, start_order);
285            }
286        }
287    }
288
289    Some(data)
290}
291
292#[must_use]
293/// Computes the track ID for a given track code. Returns [`None`] if something failed in the process.
294pub fn export_to_id(track_code: &str) -> Option<String> {
295    let track_data = decode_track_code(track_code)?;
296    let id = hash_vec(track_data.track_data);
297    Some(id)
298}
299
300impl Display for Environment {
301    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
302        match self {
303            Self::Summer => write!(f, "Summer"),
304            Self::Winter => write!(f, "Winter"),
305            Self::Desert => write!(f, "Desert"),
306        }
307    }
308}
polytrack_codes/v5/mod.rs

polytrack_codes/v5/
mod.rs
