cdf.rs

// Copyright (c) 2017-2021, The rav1e contributors. All rights reserved
//
// This source code is subject to the terms of the BSD 2 Clause License and
// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
// was not distributed with this source code in the LICENSE file, you can
// obtain it at www.aomedia.org/license/software. If the Alliance for Open
// Media Patent License 1.0 was not distributed with this source code in the
// PATENTS file, you can obtain it at www.aomedia.org/license/patent.

pub const fn cdf<const VARS: usize, const CDF_LEN: usize>(
  vars: [u16; VARS],
) -> [u16; CDF_LEN] {
  // Ensure that at least one zero is kept at the end
  assert!(CDF_LEN > VARS);

  let mut out = [0; CDF_LEN];
  let mut i = 0;
  while i < vars.len() {
    assert!(vars[i] <= 32768);
    out[i] = 32768 - vars[i];
    i += 1;
  }

  out
}

pub const fn cdf_2d<
  const VARS: usize,
  const CDF_LEN: usize,
  const N_2D: usize,
>(
  vars: [[u16; VARS]; N_2D],
) -> [[u16; CDF_LEN]; N_2D] {
  let mut out = [[0u16; CDF_LEN]; N_2D];
  let mut c = 0;
  while c < vars.len() {
    out[c] = cdf(vars[c]);
    c += 1;
  }

  out
}

pub const fn cdf_3d<
  const VARS: usize,
  const CDF_LEN: usize,
  const N_2D: usize,
  const N_3D: usize,
>(
  vars: [[[u16; VARS]; N_2D]; N_3D],
) -> [[[u16; CDF_LEN]; N_2D]; N_3D] {
  let mut out = [[[0u16; CDF_LEN]; N_2D]; N_3D];
  let mut c = 0;
  while c < vars.len() {
    out[c] = cdf_2d(vars[c]);
    c += 1;
  }

  out
}

pub const fn cdf_4d<
  const VARS: usize,
  const CDF_LEN: usize,
  const N_2D: usize,
  const N_3D: usize,
  const N_4D: usize,
>(
  vars: [[[[u16; VARS]; N_2D]; N_3D]; N_4D],
) -> [[[[u16; CDF_LEN]; N_2D]; N_3D]; N_4D] {
  let mut out = [[[[0u16; CDF_LEN]; N_2D]; N_3D]; N_4D];
  let mut c = 0;
  while c < vars.len() {
    out[c] = cdf_3d(vars[c]);
    c += 1;
  }

  out
}

pub const fn cdf_5d<
  const VARS: usize,
  const CDF_LEN: usize,
  const N_2D: usize,
  const N_3D: usize,
  const N_4D: usize,
  const N_5D: usize,
>(
  vars: [[[[[u16; VARS]; N_2D]; N_3D]; N_4D]; N_5D],
) -> [[[[[u16; CDF_LEN]; N_2D]; N_3D]; N_4D]; N_5D] {
  let mut out = [[[[[0u16; CDF_LEN]; N_2D]; N_3D]; N_4D]; N_5D];
  let mut c = 0;
  while c < vars.len() {
    out[c] = cdf_4d(vars[c]);
    c += 1;
  }

  out
}

#[cfg(test)]
mod test {
  use super::*;

  #[test]
  fn cdf_len_ok() {
    let _: [u16; 5] = cdf([]);
    let _: [u16; 5] = cdf([1]);
    let _: [u16; 5] = cdf([1, 2, 3, 4]);
  }

  #[test]
  #[should_panic]
  fn cdf_len_panics() {
    let _: [u16; 5] = cdf([1, 2, 3, 4, 5]);
  }

  #[test]
  #[should_panic]
  fn cdf_val_panics() {
    let _: [u16; 5] = cdf([40000]);
  }

  #[test]
  fn cdf_vals_ok() {
    let cdf: [u16; 5] = cdf([2000, 10000, 32768, 0]);
    assert_eq!(cdf, [30768, 22768, 0, 32768, 0]);
  }

  #[test]
  fn cdf_5d_ok() {
    let cdf: [[[[[u16; 4]; 2]; 1]; 1]; 1] =
      cdf_5d([[[[[1000, 2000], [3000, 4000]]]]]);
    assert_eq!(cdf, [[[[[31768, 30768, 0, 0], [29768, 28768, 0, 0],]]]])
  }
}