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