polytrack_codes/tools/
mod.rs

1#![allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
2#[cfg(test)]
3mod tests;
4
5pub mod prelude {
6    pub use super::{read::*, write::*, Track};
7}
8
9use std::io::Read as _;
10
11use flate2::{
12    Compression,
13    read::{ZlibDecoder, ZlibEncoder},
14};
15
16const ENCODE_VALUES: [char; 62] = [
17    'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
18    'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l',
19    'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '0', '1', '2', '3', '4',
20    '5', '6', '7', '8', '9',
21];
22
23const DECODE_VALUES: [i32; 123] = [
24    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
25    -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
26    52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -1, -1, -1, -1, -1, -1, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8,
27    9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -1, -1, -1, -1, -1, -1, 26,
28    27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
29    51,
30];
31
32#[must_use]
33/// Encode the given byte buffer into base62 encoded text according to Polytrack's base62 implementation.
34/// Returns [`None`] if something failed in the process.
35pub fn encode(input: &[u8]) -> Option<String> {
36    let mut bit_pos = 0;
37    let mut res = String::new();
38
39    while bit_pos < 8 * input.len() {
40        let mut char_value = encode_chars(input, bit_pos)?;
41        // if char_num ends with 11110, shorten it to 5 bits
42        // (getting rid of value 62 and 63, which are too big for base62)
43        if (char_value & 30) == 30 {
44            char_value &= 31;
45            bit_pos += 5;
46        } else {
47            bit_pos += 6;
48        }
49        res.push(*ENCODE_VALUES.get(char_value)?);
50    }
51
52    Some(res)
53}
54
55#[must_use]
56/// Decode the given string as base62 text according to Polytrack's base62 implementation.
57/// Returns [`None`] if any character isn't valid for base62 encoded text.
58pub fn decode(input: &str) -> Option<Vec<u8>> {
59    let mut out_pos = 0;
60    let mut bytes_out: Vec<u8> = Vec::new();
61
62    for (i, ch) in input.chars().enumerate() {
63        let char_code = ch as usize;
64        let char_value = *DECODE_VALUES.get(char_code)?;
65        if char_value == -1 {
66            return None;
67        }
68        // 5 if char_value is 30 or 31, 6 otherwise (see encode for explanation)
69        let value_len = if (char_value & 30) == 30 { 5 } else { 6 };
70        decode_chars(
71            &mut bytes_out,
72            out_pos,
73            value_len,
74            char_value,
75            i == input.len() - 1,
76        );
77        out_pos += value_len;
78    }
79
80    Some(bytes_out)
81}
82
83fn encode_chars(bytes: &[u8], bit_index: usize) -> Option<usize> {
84    if bit_index >= 8 * bytes.len() {
85        return None;
86    }
87
88    let byte_index = bit_index / 8;
89    let current_byte = *bytes.get(byte_index)? as usize;
90    let offset = bit_index - 8 * byte_index;
91    if offset <= 2 || byte_index >= bytes.len() - 1 {
92        // move mask into right position, get only offset bits of current_byte, move back
93        Some((current_byte & (63 << offset)) >> offset)
94    } else {
95        let next_byte = *bytes.get(byte_index + 1)? as usize;
96        // same concept as above, move mask into right position,
97        // get correct bits of current and next byte, move back, combine the two
98        Some(
99            ((current_byte & (63 << offset)) >> offset)
100                | ((next_byte & (63 >> (8 - offset))) << (8 - offset)),
101        )
102    }
103}
104
105fn decode_chars(
106    bytes: &mut Vec<u8>,
107    bit_index: usize,
108    value_len: usize,
109    char_value: i32,
110    is_last: bool,
111) {
112    let byte_index = bit_index / 8;
113    while byte_index >= bytes.len() {
114        bytes.push(0);
115    }
116
117    // offset in current byte
118    let offset = bit_index - 8 * byte_index;
119
120    // writes value into byte (only part that fits)
121    bytes[byte_index] |= ((char_value << offset) & 0xFF) as u8;
122
123    // in case of value going into next byte add that part
124    if offset > 8 - value_len && !is_last {
125        let byte_index_next = byte_index + 1;
126        if byte_index_next >= bytes.len() {
127            bytes.push(0);
128        }
129
130        // write rest of value into next byte
131        bytes[byte_index_next] |= (char_value >> (8 - offset)) as u8;
132    }
133}
134
135#[must_use]
136pub fn decompress(data: &[u8]) -> Option<Vec<u8>> {
137    let mut decoder = ZlibDecoder::new(data);
138    let mut decompressed_data = Vec::new();
139    decoder.read_to_end(&mut decompressed_data).ok()?;
140    Some(decompressed_data)
141}
142
143#[must_use]
144pub fn compress(data: &[u8]) -> Option<Vec<u8>> {
145    let mut encoder = ZlibEncoder::new(data, Compression::best());
146    let mut compressed_data = Vec::new();
147    encoder.read_to_end(&mut compressed_data).ok()?;
148    Some(compressed_data)
149}
150
151#[must_use]
152pub fn hash_vec(track_data: Vec<u8>) -> String {
153    sha256::digest(track_data)
154}
155
156#[derive(Debug, PartialEq, Eq, Clone)]
157pub struct Track {
158    pub name: String,
159    pub author: Option<String>,
160    pub track_data: Vec<u8>,
161}
162
163pub(crate) mod read {
164    #[inline]
165    pub fn read_u8(buf: &[u8], offset: &mut usize) -> Option<u8> {
166        let res = buf.get(*offset).copied();
167        *offset += 1;
168        res
169    }
170    #[inline]
171    pub fn read_u16(buf: &[u8], offset: &mut usize) -> Option<u16> {
172        let res = Some(u16::from(*buf.get(*offset)?) | (u16::from(*buf.get(*offset + 1)?) << 8));
173        *offset += 2;
174        res
175    }
176    #[inline]
177    pub fn read_i24(buf: &[u8], offset: &mut usize) -> Option<i32> {
178        let res = Some(
179            i32::from(*buf.get(*offset)?)
180                | (i32::from(*buf.get(*offset + 1)?) << 8)
181                | (i32::from(*buf.get(*offset + 2)?) << 16),
182        );
183        *offset += 3;
184        res
185    }
186    #[inline]
187    pub fn read_u32(buf: &[u8], offset: &mut usize) -> Option<u32> {
188        let res = Some(
189            u32::from(*buf.get(*offset)?)
190                | (u32::from(*buf.get(*offset + 1)?) << 8)
191                | (u32::from(*buf.get(*offset + 2)?) << 16)
192                | (u32::from(*buf.get(*offset + 3)?) << 24),
193        );
194        *offset += 4;
195        res
196    }
197}
198
199pub(crate) mod write {
200    #[inline]
201    pub fn write_u8(data: &mut Vec<u8>, value: u32) {
202        data.push((value & 0xFF) as u8);
203    }
204    #[inline]
205    pub fn write_u16(data: &mut Vec<u8>, value: u32) {
206        data.push((value & 0xFF) as u8);
207        data.push((value >> 8 & 0xFF) as u8);
208    }
209    #[inline]
210    pub fn write_u24(data: &mut Vec<u8>, value: u32) {
211        data.push((value & 0xFF) as u8);
212        data.push((value >> 8 & 0xFF) as u8);
213        data.push((value >> 16 & 0xFF) as u8);
214    }
215    #[inline]
216    pub fn write_u32(data: &mut Vec<u8>, value: u32) {
217        data.push((value & 0xFF) as u8);
218        data.push((value >> 8 & 0xFF) as u8);
219        data.push((value >> 16 & 0xFF) as u8);
220        data.push((value >> 24 & 0xFF) as u8);
221    }
222}
polytrack_codes/tools/mod.rs

polytrack_codes/tools/
mod.rs
