rav1e/util/
cdf.rs

1// Copyright (c) 2017-2021, The rav1e contributors. All rights reserved
2//
3// This source code is subject to the terms of the BSD 2 Clause License and
4// the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
5// was not distributed with this source code in the LICENSE file, you can
6// obtain it at www.aomedia.org/license/software. If the Alliance for Open
7// Media Patent License 1.0 was not distributed with this source code in the
8// PATENTS file, you can obtain it at www.aomedia.org/license/patent.
9
10pub const fn cdf<const VARS: usize, const CDF_LEN: usize>(
11  vars: [u16; VARS],
12) -> [u16; CDF_LEN] {
13  // Ensure that at least one zero is kept at the end
14  assert!(CDF_LEN > VARS);
15
16  let mut out = [0; CDF_LEN];
17  let mut i = 0;
18  while i < vars.len() {
19    assert!(vars[i] <= 32768);
20    out[i] = 32768 - vars[i];
21    i += 1;
22  }
23
24  out
25}
26
27pub const fn cdf_2d<
28  const VARS: usize,
29  const CDF_LEN: usize,
30  const N_2D: usize,
31>(
32  vars: [[u16; VARS]; N_2D],
33) -> [[u16; CDF_LEN]; N_2D] {
34  let mut out = [[0u16; CDF_LEN]; N_2D];
35  let mut c = 0;
36  while c < vars.len() {
37    out[c] = cdf(vars[c]);
38    c += 1;
39  }
40
41  out
42}
43
44pub const fn cdf_3d<
45  const VARS: usize,
46  const CDF_LEN: usize,
47  const N_2D: usize,
48  const N_3D: usize,
49>(
50  vars: [[[u16; VARS]; N_2D]; N_3D],
51) -> [[[u16; CDF_LEN]; N_2D]; N_3D] {
52  let mut out = [[[0u16; CDF_LEN]; N_2D]; N_3D];
53  let mut c = 0;
54  while c < vars.len() {
55    out[c] = cdf_2d(vars[c]);
56    c += 1;
57  }
58
59  out
60}
61
62pub const fn cdf_4d<
63  const VARS: usize,
64  const CDF_LEN: usize,
65  const N_2D: usize,
66  const N_3D: usize,
67  const N_4D: usize,
68>(
69  vars: [[[[u16; VARS]; N_2D]; N_3D]; N_4D],
70) -> [[[[u16; CDF_LEN]; N_2D]; N_3D]; N_4D] {
71  let mut out = [[[[0u16; CDF_LEN]; N_2D]; N_3D]; N_4D];
72  let mut c = 0;
73  while c < vars.len() {
74    out[c] = cdf_3d(vars[c]);
75    c += 1;
76  }
77
78  out
79}
80
81pub const fn cdf_5d<
82  const VARS: usize,
83  const CDF_LEN: usize,
84  const N_2D: usize,
85  const N_3D: usize,
86  const N_4D: usize,
87  const N_5D: usize,
88>(
89  vars: [[[[[u16; VARS]; N_2D]; N_3D]; N_4D]; N_5D],
90) -> [[[[[u16; CDF_LEN]; N_2D]; N_3D]; N_4D]; N_5D] {
91  let mut out = [[[[[0u16; CDF_LEN]; N_2D]; N_3D]; N_4D]; N_5D];
92  let mut c = 0;
93  while c < vars.len() {
94    out[c] = cdf_4d(vars[c]);
95    c += 1;
96  }
97
98  out
99}
100
101#[cfg(test)]
102mod test {
103  use super::*;
104
105  #[test]
106  fn cdf_len_ok() {
107    let _: [u16; 5] = cdf([]);
108    let _: [u16; 5] = cdf([1]);
109    let _: [u16; 5] = cdf([1, 2, 3, 4]);
110  }
111
112  #[test]
113  #[should_panic]
114  fn cdf_len_panics() {
115    let _: [u16; 5] = cdf([1, 2, 3, 4, 5]);
116  }
117
118  #[test]
119  #[should_panic]
120  fn cdf_val_panics() {
121    let _: [u16; 5] = cdf([40000]);
122  }
123
124  #[test]
125  fn cdf_vals_ok() {
126    let cdf: [u16; 5] = cdf([2000, 10000, 32768, 0]);
127    assert_eq!(cdf, [30768, 22768, 0, 32768, 0]);
128  }
129
130  #[test]
131  fn cdf_5d_ok() {
132    let cdf: [[[[[u16; 4]; 2]; 1]; 1]; 1] =
133      cdf_5d([[[[[1000, 2000], [3000, 4000]]]]]);
134    assert_eq!(cdf, [[[[[31768, 30768, 0, 0], [29768, 28768, 0, 0],]]]])
135  }
136}