nihav-ms/src/codecs/imaadpcmenc.rs

   1 use nihav_core::codecs::*;
   2 use nihav_core::io::byteio::*;
   3 use nihav_codec_support::codecs::imaadpcm::*;
   4
   5 struct NibbleWriter<'a> {
   6     bw:     ByteWriter<'a>,
   7     val:    u8,
   8     first:  bool,
   9 }
  10
  11 impl<'a> NibbleWriter<'a> {
  12     fn new(bw: ByteWriter<'a>) -> Self {
  13         Self { bw, val: 0, first: true }
  14     }
  15     fn write(&mut self, nib: u8) -> EncoderResult<()> {
  16         if self.first {
  17             self.val = nib;
  18             self.first = false;
  19         } else {
  20             self.val |= nib << 4;
  21             self.bw.write_byte(self.val)?;
  22             self.first = true;
  23         }
  24         Ok(())
  25     }
  26 }
  27
  28 #[derive(Clone,Copy)]
  29 struct TrellisNode {
  30     state:  IMAState,
  31     nib:    u8,
  32     error:  i64,
  33 }
  34
  35 impl Default for TrellisNode {
  36     fn default() -> Self {
  37         TrellisNode {
  38             state:  IMAState::new(),
  39             nib:    0,
  40             error:  0,
  41         }
  42     }
  43 }
  44
  45 #[derive(Default)]
  46 struct IMAADPCMEncoder {
  47     stream:     Option<NAStreamRef>,
  48     samples:    Vec<i16>,
  49     block_len:  usize,
  50     block_size: usize,
  51     channels:   usize,
  52     flush:      bool,
  53     srate:      u32,
  54     trellis:    bool,
  55     nodes:      Vec<[TrellisNode; 16]>,
  56     first:      Vec<IMAState>,
  57     nibs:       Vec<Vec<u8>>,
  58 }
  59
  60 const DEFAULT_BLOCK_LEN: usize = 256;
  61
  62 impl IMAADPCMEncoder {
  63     fn new() -> Self { Self::default() }
  64     fn encode_packet(&mut self) -> EncoderResult<NAPacket> {
  65         if self.samples.is_empty() {
  66             return Err(EncoderError::TryAgain);
  67         }
  68         let cur_len = (self.samples.len() / self.channels).min(self.block_len);
  69         if cur_len < self.block_len && !self.flush {
  70             return Err(EncoderError::TryAgain);
  71         }
  72         if cur_len < self.block_len {
  73             self.samples.resize(self.block_len * self.channels, 0);
  74         }
  75
  76         let mut dbuf = vec![0u8; self.block_size];
  77         let mut mw = MemoryWriter::new_write(dbuf.as_mut_slice());
  78         let mut bw = ByteWriter::new(&mut mw);
  79
  80         for ch in 0..self.channels {
  81             self.first[ch].predictor = i32::from(self.samples[ch]);
  82             self.first[ch].step = Self::calc_step(self.samples[ch], self.samples[ch + self.channels]) as usize;
  83         }
  84
  85         if !self.trellis {
  86             for ch in 0..self.channels {
  87                 bw.write_u16le(self.first[ch].predictor as u16)?;
  88                 bw.write_byte(self.first[ch].step as u8)?;
  89                 bw.write_byte(0)?;
  90             }
  91             let mut nw = NibbleWriter::new(bw);
  92             for samples in self.samples.chunks(self.channels).take(self.block_len).skip(1) {
  93                 for (state, &samp) in self.first.iter_mut().zip(samples.iter()) {
  94                     let nib = state.compress_sample(samp);
  95                     state.expand_sample(nib);
  96                     nw.write(nib)?;
  97                 }
  98             }
  99             drop(nw);
 100         } else {
 101             self.nodes.reserve(self.block_len);
 102             self.nibs.resize(self.channels, Vec::new());
 103             for nibs in self.nibs.iter_mut() {
 104                 nibs.resize(self.block_len, 0);
 105             }
 106             let step = self.channels;
 107             let mut state = [TrellisNode::default(); 16];
 108             for ch in 0..self.channels {
 109                 self.nodes.clear();
 110                 for i in 0..16 {
 111                     let step = (((self.first[ch].step + i) as isize) - 8).max(0).min(IMA_MAX_STEP as isize) as usize;
 112                     state[i].state.predictor = self.first[ch].predictor;
 113                     state[i].state.step = step;
 114                     state[i].error = 0;
 115                     state[i].nib = step as u8;
 116                 }
 117                 self.nodes.push(state);
 118                 let mut sidx = ch + step;
 119                 for _i in 1..self.block_len {
 120                     let sample = self.samples[sidx];
 121
 122                     for (nnib, cstate) in state.iter_mut().enumerate() {
 123                         for nib in 0..16 {
 124                             let pnode = &self.nodes[self.nodes.len() - 1][nib];
 125                             let mut ima = pnode.state;
 126                             let nsamp = ima.expand_sample(nnib as u8);
 127                             let diff = i64::from(i32::from(sample) - i32::from(nsamp));
 128                             let error = pnode.error + diff * diff;
 129                             if (nib == 0) || error < cstate.error {
 130                                 cstate.state = ima;
 131                                 cstate.nib = nib as u8;
 132                                 cstate.error = error;
 133                             }
 134                         }
 135                     }
 136                     self.nodes.push(state);
 137
 138                     sidx += step;
 139                 }
 140
 141                 let mut idx = 0;
 142                 let mut best_err = self.nodes[self.nodes.len() - 1][0].error;
 143                 for (i, node) in self.nodes[self.nodes.len() - 1].iter().enumerate() {
 144                     if node.error < best_err {
 145                         best_err = node.error;
 146                         idx = i;
 147                     }
 148                 }
 149                 let mut dst = self.nibs[ch].iter_mut().rev();
 150                 while let Some(nodes) = self.nodes.pop() {
 151                     *dst.next().unwrap() = idx as u8;
 152                     idx = nodes[idx].nib as usize;
 153                 }
 154                 self.nibs[ch][0] = idx as u8;
 155             }
 156             for ch in 0..self.channels {
 157                 bw.write_u16le(self.first[ch].predictor as u16)?;
 158                 bw.write_byte(self.nibs[ch][0])?;
 159                 bw.write_byte(0)?;
 160             }
 161             let mut nw = NibbleWriter::new(bw);
 162             for i in 1..self.block_len {
 163                 for ch in 0..self.channels {
 164                     nw.write(self.nibs[ch][i])?;
 165                 }
 166             }
 167             drop(nw);
 168         }
 169
 170         self.samples.drain(..self.block_len * self.channels);
 171         let ts = NATimeInfo::new(None, None, Some(1), 1, self.srate);
 172         Ok(NAPacket::new(self.stream.clone().unwrap(), ts, true, dbuf))
 173     }
 174     fn calc_block_size(nsamps: usize, channels: usize) -> usize {
 175         ((nsamps - 1) * channels + 1) / 2 + 4 * channels
 176     }
 177     fn calc_step(samp1: i16, samp2: i16) -> u8 {
 178         let diff = (i32::from(samp1) - i32::from(samp2)).abs();
 179         for (i, &step) in IMA_STEP_TABLE.iter().enumerate() {
 180             if step >= diff {
 181                 return i as u8;
 182             }
 183         }
 184         IMA_MAX_STEP
 185     }
 186 }
 187
 188 impl NAEncoder for IMAADPCMEncoder {
 189     fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
 190         match encinfo.format {
 191             NACodecTypeInfo::None => {
 192                 let mut ofmt = EncodeParameters::default();
 193                 ofmt.format = NACodecTypeInfo::Audio(NAAudioInfo::new(0, 1, SND_S16_FORMAT, DEFAULT_BLOCK_LEN));
 194                 Ok(ofmt)
 195             },
 196             NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError),
 197             NACodecTypeInfo::Audio(ainfo) => {
 198                 let mut outinfo = ainfo;
 199                 outinfo.channels = outinfo.channels.max(1);
 200                 if outinfo.format != SND_S16P_FORMAT && outinfo.format != SND_S16_FORMAT {
 201                     outinfo.format = SND_S16_FORMAT;
 202                 }
 203                 if outinfo.block_len == 0 {
 204                     outinfo.block_len = DEFAULT_BLOCK_LEN;
 205                 }
 206                 if outinfo.block_len < 2 {
 207                     outinfo.block_len = 2;
 208                 }
 209                 if (outinfo.block_len & 7) != 0 {
 210                     outinfo.block_len = (outinfo.block_len & 7) + 7;
 211                 }
 212                 let mut ofmt = *encinfo;
 213                 ofmt.format = NACodecTypeInfo::Audio(outinfo);
 214                 Ok(ofmt)
 215             }
 216         }
 217     }
 218     fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
 219         match encinfo.format {
 220             NACodecTypeInfo::None => Err(EncoderError::FormatError),
 221             NACodecTypeInfo::Video(_) => Err(EncoderError::FormatError),
 222             NACodecTypeInfo::Audio(ainfo) => {
 223                 if ainfo.format != SND_S16P_FORMAT && ainfo.format != SND_S16_FORMAT {
 224                     return Err(EncoderError::FormatError);
 225                 }
 226                 if ainfo.block_len < 2 || ((ainfo.block_len * (ainfo.channels as usize)) & 7) != 0 {
 227                     return Err(EncoderError::FormatError);
 228                 }
 229                 self.channels = ainfo.channels as usize;
 230                 self.block_len = ainfo.block_len;
 231                 self.block_size = Self::calc_block_size(self.block_len, ainfo.channels as usize);
 232                 self.first = Vec::with_capacity(self.channels);
 233                 for _ in 0..self.channels {
 234                     self.first.push(IMAState::new());
 235                 }
 236
 237                 let soniton = NASoniton::new(4, 0);
 238                 let out_ainfo = NAAudioInfo::new(ainfo.sample_rate, ainfo.channels, soniton, Self::calc_block_size(self.block_len, self.channels));
 239                 let info = NACodecInfo::new("ima-adpcm-ms", NACodecTypeInfo::Audio(out_ainfo), None);
 240                 let mut stream = NAStream::new(StreamType::Audio, stream_id, info, self.block_len as u32, ainfo.sample_rate, 0);
 241                 stream.set_num(stream_id as usize);
 242                 let stream = stream.into_ref();
 243
 244                 self.stream = Some(stream.clone());
 245                 self.samples = Vec::with_capacity(self.block_len * self.channels);
 246                 self.srate = ainfo.sample_rate;
 247                 self.flush = false;
 248
 249                 Ok(stream)
 250             },
 251         }
 252     }
 253     fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
 254         let buf = frm.get_buffer();
 255         if let Some(ref abuf) = buf.get_abuf_i16() {
 256             let src = abuf.get_data();
 257             let len = abuf.get_length();
 258             let ch  = abuf.get_chmap().num_channels();
 259             if abuf.get_step() > 1 || ch == 1 {
 260                 self.samples.extend_from_slice(&src[..len * ch]);
 261             } else {
 262                 let astride = abuf.get_stride();
 263                 self.samples.reserve(len * ch);
 264                 for i in 0..len {
 265                     for ch in 0..self.channels {
 266                         self.samples.push(src[ch * astride + i]);
 267                     }
 268                 }
 269             }
 270             Ok(())
 271         } else {
 272             Err(EncoderError::InvalidParameters)
 273         }
 274     }
 275     fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
 276         if let Ok(pkt) = self.encode_packet() {
 277             Ok(Some(pkt))
 278         } else {
 279             Ok(None)
 280         }
 281     }
 282     fn flush(&mut self) -> EncoderResult<()> {
 283         self.flush = true;
 284         Ok(())
 285     }
 286 }
 287
 288 const ENCODER_OPTS: &[NAOptionDefinition] = &[
 289     NAOptionDefinition {
 290         name: "trellis", description: "Use trellis search",
 291         opt_type: NAOptionDefinitionType::Bool },
 292 ];
 293
 294 impl NAOptionHandler for IMAADPCMEncoder {
 295     fn get_supported_options(&self) -> &[NAOptionDefinition] { ENCODER_OPTS }
 296     fn set_options(&mut self, options: &[NAOption]) {
 297         for option in options.iter() {
 298             for opt_def in ENCODER_OPTS.iter() {
 299                 if opt_def.check(option).is_ok() {
 300                     match option.name {
 301                         "trellis" => {
 302                             if let NAValue::Bool(val) = option.value {
 303                                 self.trellis = val;
 304                             }
 305                         },
 306                         _ => {},
 307                     };
 308                 }
 309             }
 310         }
 311     }
 312     fn query_option_value(&self, name: &str) -> Option<NAValue> {
 313         match name {
 314             "trellis" => Some(NAValue::Bool(self.trellis)),
 315             _ => None,
 316         }
 317     }
 318 }
 319
 320 pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
 321     Box::new(IMAADPCMEncoder::new())
 322 }
 323
 324 #[cfg(test)]
 325 mod test {
 326     use nihav_core::codecs::*;
 327     use nihav_core::demuxers::*;
 328     use nihav_core::muxers::*;
 329     use nihav_codec_support::test::enc_video::*;
 330     use crate::*;
 331     use nihav_commonfmt::*;
 332
 333     fn test_ima_adpcm_ms_encoder(name: &'static str, trellis: bool, hash: &[u32; 4]) {
 334         let mut dmx_reg = RegisteredDemuxers::new();
 335         generic_register_all_demuxers(&mut dmx_reg);
 336         let mut dec_reg = RegisteredDecoders::new();
 337         generic_register_all_decoders(&mut dec_reg);
 338         ms_register_all_decoders(&mut dec_reg);
 339         let mut mux_reg = RegisteredMuxers::new();
 340         generic_register_all_muxers(&mut mux_reg);
 341         let mut enc_reg = RegisteredEncoders::new();
 342         ms_register_all_encoders(&mut enc_reg);
 343
 344         let dec_config = DecoderTestParams {
 345                 demuxer:        "avi",
 346                 in_name:        "assets/Duck/ot171_vp40.avi",
 347                 stream_type:    StreamType::Audio,
 348                 limit:          None,
 349                 dmx_reg, dec_reg,
 350             };
 351         let enc_config = EncoderTestParams {
 352                 muxer:          "wav",
 353                 enc_name:       "ima-adpcm-ms",
 354                 out_name:       name,
 355                 mux_reg, enc_reg,
 356             };
 357         let dst_ainfo = NAAudioInfo {
 358                 sample_rate:    0,
 359                 channels:       0,
 360                 format:         SND_S16_FORMAT,
 361                 block_len:      128,
 362             };
 363         let enc_params = EncodeParameters {
 364                 format:  NACodecTypeInfo::Audio(dst_ainfo),
 365                 quality: 0,
 366                 bitrate: 0,
 367                 tb_num:  0,
 368                 tb_den:  0,
 369                 flags:   0,
 370             };
 371         let enc_options = &[
 372                 NAOption{name: "trellis", value: NAValue::Bool(trellis)},
 373             ];
 374         test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
 375
 376         test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
 377                           hash);
 378     }
 379     #[test]
 380     fn test_ima_adpcm_ms_encoder_notrellis() {
 381         test_ima_adpcm_ms_encoder("msimaadpcm-notr.wav", false,
 382             &[0x59909f10, 0xf0420dd2, 0xcfee7ef5, 0x7623caa3]);
 383     }
 384     #[test]
 385     fn test_ima_adpcm_ms_encoder_trellis() {
 386         test_ima_adpcm_ms_encoder("msimaadpcm-tr.wav", true,
 387             &[0x99e373e2, 0x80439b4b, 0xcb4f0b78, 0xeb1e9a51]);
 388     }
 389 }