pulsejet_rs/
decode.rs

1use std::{f32::consts::PI, mem::size_of};
2
3use crate::{cmath::CMath, common::*};
4
5trait SkipForward {
6    fn skip_forward(&mut self, offset: usize);
7}
8
9impl<T> SkipForward for &[T] {
10    fn skip_forward(&mut self, offset: usize) {
11        *self = &self[offset..];
12    }
13}
14
15/// Decodes an encoded pulsejet sample into a newly-allocated buffer.
16///
17/// This function is optimized for size and designed to be compiled in a
18/// size-constrained environment. In such environments, it's common not
19/// to have access to all of the required math functions, and instead
20/// implement them by hand. For this reason, this decoder does not
21/// depend on any such functions directly, and instead expects that
22/// `cos`, `exp2`, `sin`, and `sqrt` functions are defined using a [`CMath`]
23/// trait. Default implemention is provided through [`Std`](`crate::Std`) type.
24/// pulsejet expects that these functions behave similarly
25/// to the corresponding similarly-named cmath functions. This shim
26/// mechanism can also be used to provide less accurate, speed-optimized
27/// versions of these functions if desired.
28///
29/// Additionally, this function will not perform any error checking or
30/// handling. The included metadata API can be used for high-level error
31/// checking before decoding takes place if required (albeit not in a
32/// non-size-constrained environment).
33///
34/// # Returns
35///
36/// Decoded samples in the [-1, 1] range (normalized).
37pub fn decode<M: CMath>(mut input_stream: &[u8]) -> Vec<f32> {
38    // Skip tag and codec version
39    input_stream.skip_forward(8);
40
41    // Read frame count, determine number of samples, and allocate
42    // output sample buffer
43    let mut num_frames = u16::from_le_bytes(
44        input_stream[..size_of::<u16>()].try_into().unwrap(),
45    ) as u32;
46    input_stream.skip_forward(size_of::<u16>());
47    let num_samples = num_frames * FRAME_SIZE;
48    let mut samples = vec![0f32; num_samples as usize];
49
50    // We're going to decode one more frame than we output, so adjust
51    // the frame count
52    num_frames += 1;
53
54    // Set up and skip window mode stream
55    let mut window_mode_stream = input_stream;
56    input_stream.skip_forward(num_frames as usize);
57
58    // Set up and skip quantized band bin stream
59    let mut quantized_band_bin_stream = input_stream;
60    input_stream.skip_forward((num_frames * NUM_TOTAL_BINS) as usize);
61
62    // Allocate padded sample buffer, and fill with silence
63    let num_padded_samples = num_samples + FRAME_SIZE * 2;
64    let mut padded_samples = vec![0f32; num_padded_samples as usize];
65
66    // Initialize LCG
67    let mut lcg_state: u32 = 0;
68
69    // Clear quantized band energy predictions
70    let mut quantized_band_energy_predictions = vec![0u8; NUM_BANDS];
71
72    // Decode frames
73    for frame_index in 0..num_frames {
74        // Read window mode for this frame
75        let window_mode = match window_mode_stream[0] {
76            0 => WindowMode::Long,
77            1 => WindowMode::Short,
78            2 => WindowMode::Start,
79            3 => WindowMode::Stop,
80            _ => unreachable!(),
81        };
82        window_mode_stream.skip_forward(1);
83
84        // Determine subframe configuration from window mode
85        let mut num_subframes: u32 = 1;
86        let mut subframe_window_offset: u32 = 0;
87        let mut subframe_window_size: u32 = LONG_WINDOW_SIZE;
88        if window_mode == WindowMode::Short {
89            num_subframes = NUM_SHORT_WINDOWS_PER_FRAME;
90            subframe_window_offset =
91                LONG_WINDOW_SIZE / 4 - SHORT_WINDOW_SIZE / 4;
92            subframe_window_size = SHORT_WINDOW_SIZE;
93        }
94
95        // Decode subframe(s)
96        for subframe_index in 0..num_subframes {
97            // Decode bands
98            let mut window_bins = vec![0f32; FRAME_SIZE as usize];
99            let mut band_bins = &mut window_bins[..];
100            for band_index in 0..NUM_BANDS {
101                // Decode band bins
102                let num_bins =
103                    BAND_TO_NUM_BINS[band_index] as u32 / num_subframes;
104                let mut num_nonzero_bins: u32 = 0;
105                for bin_index in 0..num_bins {
106                    let bin_q = quantized_band_bin_stream[0] as i8;
107                    quantized_band_bin_stream.skip_forward(1);
108                    if bin_q != 0 {
109                        num_nonzero_bins += 1;
110                    }
111                    let bin = bin_q as f32;
112                    band_bins[bin_index as usize] = bin;
113                }
114
115                // If this band is significantly sparse, fill in (nearly)
116                // spectrally flat noise
117                let bin_fill = (num_nonzero_bins as f32) / (num_bins as f32);
118                let noise_fill_threshold = 0.1f32;
119                if bin_fill < noise_fill_threshold {
120                    let bin_sparsity = (noise_fill_threshold - bin_fill)
121                        / noise_fill_threshold;
122                    let noise_fill_gain = bin_sparsity * bin_sparsity;
123                    for bin_index in 0..num_bins {
124                        let noise_sample =
125                            ((lcg_state >> 16) as i8) as f32 / 127.0f32;
126                        band_bins[bin_index as usize] +=
127                            noise_sample * noise_fill_gain;
128
129                        // Transition LCG state using
130                        // Numerical Recipes parameters
131                        lcg_state = lcg_state
132                            .wrapping_mul(1664525)
133                            .wrapping_add(1013904223);
134                    }
135                }
136
137                // Decode band energy
138                let quantized_band_energy_residual = input_stream[0];
139                input_stream.skip_forward(1);
140                let quantized_band_energy: u8 =
141                    quantized_band_energy_predictions[band_index]
142                        .wrapping_add(quantized_band_energy_residual);
143                quantized_band_energy_predictions[band_index] =
144                    quantized_band_energy;
145                let band_energy = M::exp2(
146                    (quantized_band_energy) as f32 / 64.0f32 * 40.0f32
147                        - 20.0f32,
148                ) * (num_bins) as f32;
149
150                // Normalize band bins and scale by band energy
151                let epsilon: f32 = 1e-27f32;
152                let mut band_bin_energy = epsilon;
153                for bin_index in 0..num_bins {
154                    let bin = band_bins[bin_index as usize];
155                    band_bin_energy += bin * bin;
156                }
157                band_bin_energy = M::sqrt(band_bin_energy);
158                let bin_scale = band_energy / band_bin_energy;
159                for bin_index in 0..num_bins {
160                    band_bins[bin_index as usize] *= bin_scale;
161                }
162
163                band_bins = &mut band_bins[num_bins as usize..];
164            }
165
166            // Apply the IMDCT to the subframe bins, then apply the appropriate
167            // window to the resulting samples, and finally accumulate them into
168            // the padded output buffer
169            let frame_offset = frame_index * FRAME_SIZE;
170            let window_offset = subframe_window_offset
171                + subframe_index * subframe_window_size / 2;
172            for n in 0..subframe_window_size {
173                let n_plus_half = n as f32 + 0.5f32;
174
175                let mut sample = 0.0f32;
176                for k in 0..subframe_window_size / 2 {
177                    sample += (2.0f32 / (subframe_window_size / 2) as f32)
178                        * window_bins[k as usize]
179                        * M::cos(
180                            PI / (subframe_window_size / 2) as f32
181                                * (n_plus_half
182                                    + (subframe_window_size / 4) as f32)
183                                * ((k) as f32 + 0.5f32),
184                        );
185                }
186
187                let window = mdct_window(n, subframe_window_size, window_mode);
188                padded_samples[(frame_offset + window_offset + n) as usize] +=
189                    sample * window;
190            }
191        }
192    }
193
194    let size = num_samples as usize;
195    samples[..size]
196        .copy_from_slice(&padded_samples[FRAME_SIZE as usize..][..size]);
197
198    samples
199}