ogg_opus/
encode.rs

1use std::cmp::min;
2use std::process;
3
4use crate::Error;
5use crate::common::*;
6
7use byteorder::{LittleEndian, ByteOrder};
8use ogg::PacketWriter;
9use audiopus::{Bitrate, coder::{Encoder as OpusEnc, GenericCtl}};
10use rand::Rng;
11
12//--- Final range  things ------------------------------------------------------
13
14#[cfg(test)]
15use std::cell::RefCell;
16
17#[cfg(test)]
18thread_local! {
19    static LAST_FINAL_RANGE: RefCell<u32> = RefCell::new(0);
20}
21
22#[cfg(test)]
23fn set_final_range(r:u32) {
24    LAST_FINAL_RANGE.with(|f|*f.borrow_mut() = r);
25}
26
27// Just here so that it can be used in the function
28#[cfg(not(test))]
29fn set_final_range(_:u32) {}
30
31#[cfg(test)]
32pub(crate) fn get_final_range() -> u32 {
33    LAST_FINAL_RANGE.with(|f|*f.borrow())
34}
35
36//--- Code ---------------------------------------------------------------------
37
38const VER: &str = std::env!("CARGO_PKG_VERSION");
39
40const fn to_samples<const S_PS: u32>(ms: u32) -> usize {
41    ((S_PS * ms) / 1000) as usize
42}
43
44
45// In microseconds
46const fn calc_fr_size(us: u32, channels:u8, sps:u32) -> usize {
47    let samps_ms = (sps * us) as u32;
48    const US_TO_MS: u32 = 10;
49    ((samps_ms * channels as u32 ) / (1000 * US_TO_MS )) as usize
50}
51
52
53const fn opus_channels(val: u8) -> audiopus::Channels{
54    if val == 0 {
55        // Never should be 0
56        audiopus::Channels::Mono
57    }
58    else if val == 1 {
59       audiopus::Channels::Mono
60    }
61    else {
62       audiopus::Channels::Stereo
63    }
64}
65
66pub fn encode<const S_PS: u32, const NUM_CHANNELS: u8>(audio: &[i16]) -> Result<Vec<u8>, Error> {
67    //NOTE: In the future the S_PS const generic will let us use const on a lot 
68    // of things, until then we need to use variables
69
70    // This should have a bitrate of 24 Kb/s, exactly what IBM recommends
71
72    // More frame time, sligtly less overhead more problematic packet loses,
73    // a frame time of 20ms is considered good enough for most applications
74    
75    // Data
76    const FRAME_TIME_MS: u32 = 20;
77    const MAX_PACKET: usize = 4000; // Maximum theorical recommended by Opus
78    const MIN_FRAME_MICROS: u32 = 25;
79
80    let frame_samples: usize = to_samples::<S_PS>(FRAME_TIME_MS);
81    let frame_size: usize = frame_samples * (NUM_CHANNELS as usize);
82
83    // Generate the serial which is nothing but a value to identify a stream, we
84    // will also use the process id so that two programs don't use 
85    // the same serial even if getting one at the same time
86    let mut rnd = rand::thread_rng();
87    let serial = rnd.gen::<u32>() ^ process::id();
88    let mut buffer: Vec<u8> = Vec::new();
89    
90    let mut packet_writer = PacketWriter::new(&mut buffer);
91
92    let opus_sr = s_ps_to_audiopus(S_PS)?;
93
94    let mut opus_encoder = OpusEnc::new(opus_sr, opus_channels(NUM_CHANNELS), audiopus::Application::Audio)?;
95    opus_encoder.set_bitrate(Bitrate::BitsPerSecond(24000))?;
96
97    let skip = opus_encoder.lookahead().unwrap() as u16;
98    let skip_us = skip as usize;
99    let tot_samples = audio.len() + skip_us;
100    let skip_48 = calc_sr(
101        skip,
102        S_PS,
103        OGG_OPUS_SPS
104    );
105
106    let max = (tot_samples as f32 / frame_size as f32).floor() as u32;
107
108    let calc = |counter: u32| -> usize {
109        (counter as usize) * frame_size
110    };
111
112    let calc_samples = |counter:u32| -> usize {
113        (counter as usize) * frame_samples
114    };
115
116    const fn granule<const S_PS: u32>(val: usize) -> u64 {
117        calc_sr_u64(val as u64, S_PS, OGG_OPUS_SPS)
118    }
119
120    let opus_head: [u8; 19] = [
121        b'O', b'p', b'u', b's', b'H', b'e', b'a', b'd', // Magic header
122        1, // Version number, always 1
123        NUM_CHANNELS, // Channels
124        0, 0,//Pre-skip
125        0, 0, 0, 0, // Original Hz (informational)
126        0, 0, // Output gain
127        0, // Channel map family
128        // If Channel map != 0, here should go channel mapping table
129    ];
130
131    fn encode_vec(opus_encoder: &mut OpusEnc, audio: &[i16]) -> Result<Box<[u8]>, Error> {
132        let mut output: Vec<u8> = vec![0; MAX_PACKET];
133		let result = opus_encoder.encode(audio, output.as_mut_slice())?;
134		output.truncate(result);
135		Ok(output.into_boxed_slice())
136    }
137
138    fn encode_with_skip(opus_encoder: &mut OpusEnc, audio: &[i16], pos_a: usize, pos_b: usize, skip_us: usize) -> Result<Box<[u8]>, Error> {
139        if pos_a > skip_us {
140            encode_vec(opus_encoder, &audio[pos_a-skip_us..pos_b-skip_us])
141        }
142        else {
143            let mut buf = vec![0; pos_b-pos_a];
144            if pos_b > skip_us {
145                buf[skip_us - pos_a..].copy_from_slice(&audio[.. pos_b - skip_us]);
146            }
147            encode_vec(opus_encoder, &buf)
148        }
149    }
150
151    fn is_end_of_stream(pos: usize, max: usize) -> ogg::PacketWriteEndInfo {
152        if pos == max {
153            ogg::PacketWriteEndInfo::EndStream
154        }
155        else {
156            ogg::PacketWriteEndInfo::NormalPacket
157        }
158    }
159
160    let mut head = opus_head;
161    LittleEndian::write_u16(&mut head[10..12], skip_48 as u16); // Write pre-skip
162    LittleEndian::write_u32(&mut head[12..16], S_PS); // Write Samples per second
163
164    let mut opus_tags : Vec<u8> = Vec::with_capacity(60);
165    let vendor_str = format!("ogg-opus {}", VER);
166    opus_tags.extend(b"OpusTags");
167    let mut len_bf = [0u8;4];
168    LittleEndian::write_u32(&mut len_bf, vendor_str.len() as u32);
169    opus_tags.extend(&len_bf);
170    opus_tags.extend(vendor_str.bytes());
171    opus_tags.extend(&[0]); // No user comments
172
173    packet_writer.write_packet(Box::new(head), serial, ogg::PacketWriteEndInfo::EndPage, 0)?;
174    packet_writer.write_packet(opus_tags.into_boxed_slice(), serial, ogg::PacketWriteEndInfo::EndPage, 0)?;
175
176    // Do all frames
177    for counter in 0..max{ // Last value of counter is max - 1 
178        let pos_a: usize = calc(counter);
179        let pos_b: usize = calc(counter + 1);
180        
181        assert!((pos_b - pos_a) <= frame_size);
182        
183        let new_buffer = encode_with_skip(&mut opus_encoder, audio, pos_a, pos_b, skip_us)?;
184
185        packet_writer.write_packet(
186            new_buffer,
187            serial,
188            is_end_of_stream(pos_b, tot_samples),
189            granule::<S_PS>(skip_us + calc_samples(counter + 1)
190        ))?;
191    }
192
193    // Calc the biggest frame buffer that still is either smaller or the
194    // same size as the input
195    fn calc_biggest_spills<T:PartialOrd + Copy>(val: T, possibles: &[T]) -> Option<T> {
196        for container in possibles.iter().rev()  {
197            if *container <= val {
198                return Some(*container)
199            }
200        }
201        None
202    }
203
204    fn encode_no_skip(opus_encoder: &mut OpusEnc, audio: &[i16], start: usize, frame_size : usize) -> Result<Box<[u8]>, Error> {
205        encode_vec(opus_encoder, &audio[start .. start + frame_size])
206    }
207
208    // Try to add as less of empty audio as possible, first everything into
209    // small frames, and on the last one, if needed fill with 0, since the
210    // last one is going to be smaller this should be much less of a problem
211    let mut last_sample = calc(max);
212    assert!(last_sample <= audio.len() + skip_us);
213    let frame_sizes: [usize; 4] = [
214            calc_fr_size(MIN_FRAME_MICROS, NUM_CHANNELS, S_PS),
215            calc_fr_size(50, NUM_CHANNELS, S_PS),
216            calc_fr_size(100, NUM_CHANNELS, S_PS),
217            calc_fr_size(200, NUM_CHANNELS, S_PS)
218    ];
219
220    while last_sample < tot_samples {
221
222            let rem_samples = tot_samples - last_sample;
223            let last_audio_s = last_sample - min(last_sample,skip_us);
224
225            match calc_biggest_spills(rem_samples, &frame_sizes) {
226                Some(frame_size) => {
227                    let enc = if last_sample >= skip_us {
228                        encode_no_skip(&mut opus_encoder, audio, last_audio_s, frame_size)?
229                    }
230                    else {
231                        encode_with_skip(&mut opus_encoder, audio, last_sample, last_sample + frame_size, skip_us)?
232                    };
233                    last_sample += frame_size;
234                    packet_writer.write_packet(
235                        enc,
236                        serial, 
237                        is_end_of_stream(last_sample, tot_samples),
238                        granule::<S_PS>(last_sample/(NUM_CHANNELS as usize))
239                    )?;
240                }
241                None => {
242                    // Maximum size for a 2.5 ms frame
243                    const MAX_25_SIZE: usize = calc_fr_size(MIN_FRAME_MICROS, MAX_NUM_CHANNELS, OGG_OPUS_SPS);
244                    let mut in_buffer = [0i16;MAX_25_SIZE];
245                    let rem_skip = skip_us - min(last_sample, skip_us);
246                    in_buffer[rem_skip..rem_samples].copy_from_slice(&audio[last_audio_s..]);
247
248                    last_sample = tot_samples; // We end this here
249                    
250                    packet_writer.write_packet(
251                        encode_no_skip(&mut opus_encoder, &in_buffer, 0, frame_sizes[0])?,
252                        serial, 
253                        ogg::PacketWriteEndInfo::EndStream,
254                        granule::<S_PS>((skip_us + audio.len())/(NUM_CHANNELS as usize))
255                    )?;
256                    
257                }
258            }
259            
260        }
261
262    if cfg!(test) {set_final_range(opus_encoder.final_range().unwrap())}
263
264    Ok(buffer)
265}