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, Track, hash_vec};
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)]
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)]
28#[repr(u8)]
29pub enum Environment {
30    Summer,
31    Winter,
32    Desert,
33}
34
35#[derive(Debug, PartialEq, Eq)]
36pub struct Part {
37    pub id: u8,
38    pub amount: u32,
39    pub blocks: Vec<Block>,
40}
41
42#[derive(Debug, PartialEq, Eq)]
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)]
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/// Decodes the (raw binary) track data into a struct
106/// representing everything that is in the data.
107///
108/// Fields of all involved structs correspond exactly to how
109/// the data is stored in Polytrack itself.
110/// Returns [`None`] if the data is not valid track data.
111pub fn decode_track_data(data: &[u8]) -> Option<TrackInfo> {
112    #[inline]
113    fn read_u8(buf: &[u8], offset: &mut usize) -> Option<u8> {
114        let res = buf.get(*offset).copied();
115        *offset += 1;
116        res
117    }
118    #[inline]
119    fn read_u16(buf: &[u8], offset: &mut usize) -> Option<u16> {
120        let res = Some(u16::from(*buf.get(*offset)?) | (u16::from(*buf.get(*offset + 1)?) << 8));
121        *offset += 2;
122        res
123    }
124    #[inline]
125    fn read_u32(buf: &[u8], offset: &mut usize) -> Option<u32> {
126        let res = Some(
127            u32::from(*buf.get(*offset)?)
128                | (u32::from(*buf.get(*offset + 1)?) << 8)
129                | (u32::from(*buf.get(*offset + 2)?) << 16)
130                | (u32::from(*buf.get(*offset + 3)?) << 24),
131        );
132        *offset += 4;
133        res
134    }
135
136    let mut offset = 0;
137
138    let env = Environment::try_from(read_u8(data, &mut offset)?).ok()?;
139    let sun_dir = read_u8(data, &mut offset)?;
140
141    let min_x = read_u32(data, &mut offset)? as i32;
142    let min_y = read_u32(data, &mut offset)? as i32;
143    let min_z = read_u32(data, &mut offset)? as i32;
144
145    let data_bytes = read_u8(data, &mut offset)?;
146    let x_bytes = data_bytes & 3;
147    let y_bytes = (data_bytes >> 2) & 3;
148    let z_bytes = (data_bytes >> 4) & 3;
149
150    let mut parts = Vec::new();
151    while offset < data.len() {
152        let id = read_u8(data, &mut offset)?;
153        let amount = read_u32(data, &mut offset)?;
154
155        let mut blocks = Vec::new();
156        for _ in 0..amount {
157            let mut x = 0;
158            for i in 0..x_bytes {
159                x |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
160            }
161            offset += x_bytes as usize;
162
163            let mut y = 0;
164            for i in 0..x_bytes {
165                y |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
166            }
167            offset += y_bytes as usize;
168
169            let mut z = 0;
170            for i in 0..x_bytes {
171                z |= u32::from(*data.get(offset + (i as usize))?) << (8 * i);
172            }
173            offset += z_bytes as usize;
174
175            let rotation = read_u8(data, &mut offset)?;
176            if rotation > 3 {
177                return None;
178            }
179            let dir = Direction::try_from(read_u8(data, &mut offset)?).ok()?;
180            let color = read_u8(data, &mut offset)?;
181            // no custom color support for now
182            if color > 3 && color < 32 && color > 40 {
183                return None;
184            }
185
186            let cp_order = if CP_IDS.contains(&id) {
187                Some(read_u16(data, &mut offset)?)
188            } else {
189                None
190            };
191            let start_order = if START_IDS.contains(&id) {
192                Some(read_u32(data, &mut offset)?)
193            } else {
194                None
195            };
196
197            blocks.push(Block {
198                x,
199                y,
200                z,
201
202                rotation,
203                dir,
204
205                color,
206                cp_order,
207                start_order,
208            });
209        }
210        parts.push(Part { id, amount, blocks });
211    }
212
213    Some(TrackInfo {
214        env,
215        sun_dir,
216
217        min_x,
218        min_y,
219        min_z,
220
221        data_bytes,
222        parts,
223    })
224}
225
226#[must_use]
227/// Computes the track ID for a given track code. Returns [`None`] if something failed in the process.
228pub fn export_to_id(track_code: &str) -> Option<String> {
229    let track_data = decode_track_code(track_code)?;
230    let id = hash_vec(track_data.track_data);
231    Some(id)
232}
233
234impl Display for Environment {
235    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236        match self {
237            Self::Summer => write!(f, "Summer"),
238            Self::Winter => write!(f, "Winter"),
239            Self::Desert => write!(f, "Desert"),
240        }
241    }
242}
polytrack_codes/v5/mod.rs

polytrack_codes/v5/
mod.rs
