1 use rustdct::{DCTplanner, TransformType2And3};
2 use transpose::transpose;
3
4 use std::sync::Arc;
5
6 pub const SIZE_MULTIPLIER: u32 = 2;
7 pub const SIZE_MULTIPLIER_U: usize = SIZE_MULTIPLIER as usize;
8
9 pub struct DctCtxt {
10 row_dct: Arc<dyn TransformType2And3<f32>>,
11 col_dct: Arc<dyn TransformType2And3<f32>>,
12 width: usize,
13 height: usize,
14 }
15
16 impl DctCtxt {
new(width: u32, height: u32) -> Self17 pub fn new(width: u32, height: u32) -> Self {
18 let mut planner = DCTplanner::new();
19 let width = width as usize * SIZE_MULTIPLIER_U;
20 let height = height as usize * SIZE_MULTIPLIER_U;
21
22 DctCtxt {
23 row_dct: planner.plan_dct2(width),
24 col_dct: planner.plan_dct2(height),
25 width,
26 height,
27 }
28 }
29
width(&self) -> u3230 pub fn width(&self) -> u32 {
31 self.width as u32
32 }
33
height(&self) -> u3234 pub fn height(&self) -> u32 {
35 self.height as u32
36 }
37
38 /// Perform a 2D DCT on a 1D-packed vector with a given `width x height`.
39 ///
40 /// Assumes `packed_2d` is double-length for scratch space. Returns the vector truncated to
41 /// `width * height`.
42 ///
43 /// ### Panics
44 /// If `self.width * self.height * 2 != packed_2d.len()`
dct_2d(&self, mut packed_2d: Vec<f32>) -> Vec<f32>45 pub fn dct_2d(&self, mut packed_2d: Vec<f32>) -> Vec<f32> {
46 let Self { ref row_dct, ref col_dct, width, height } = *self;
47
48 let trunc_len = width * height;
49 assert_eq!(trunc_len * 2, packed_2d.len());
50
51 {
52 let (packed_2d, scratch) = packed_2d.split_at_mut(trunc_len);
53
54 for (row_in, row_out) in packed_2d.chunks_mut(width)
55 .zip(scratch.chunks_mut(width)) {
56 row_dct.process_dct2(row_in, row_out);
57 }
58
59 transpose(scratch, packed_2d, width, height);
60
61 for (row_in, row_out) in packed_2d.chunks_mut(height)
62 .zip(scratch.chunks_mut(height)) {
63 col_dct.process_dct2(row_in, row_out);
64 }
65
66 transpose(scratch, packed_2d, width, height);
67 }
68
69 packed_2d.truncate(trunc_len);
70 packed_2d
71 }
72
crop_2d(&self, packed: Vec<f32>) -> Vec<f32>73 pub fn crop_2d(&self, packed: Vec<f32>) -> Vec<f32> {
74 crop_2d_dct(packed, self.width)
75 }
76 }
77
78 /// Crop the values off a 1D-packed 2D DCT.
79 ///
80 /// Returns `packed` truncated to the premultiplied size, as determined by `rowstride`
81 ///
82 /// Generic for easier testing
crop_2d_dct<T: Copy>(mut packed: Vec<T>, rowstride: usize) -> Vec<T>83 fn crop_2d_dct<T: Copy>(mut packed: Vec<T>, rowstride: usize) -> Vec<T> {
84 // assert that the rowstride was previously multiplied by SIZE_MULTIPLIER
85 assert_eq!(rowstride % SIZE_MULTIPLIER_U, 0);
86 assert!(rowstride / SIZE_MULTIPLIER_U > 0, "rowstride cannot be cropped: {}", rowstride);
87
88 let new_rowstride = rowstride / SIZE_MULTIPLIER_U;
89
90 for new_row in 0 .. packed.len() / (rowstride * SIZE_MULTIPLIER_U) {
91 let (dest, src) = packed.split_at_mut(new_row * new_rowstride + rowstride);
92 let dest_start = dest.len() - new_rowstride;
93 let src_start = new_rowstride * new_row;
94 let src_end = src_start + new_rowstride;
95 dest[dest_start..].copy_from_slice(&src[src_start..src_end]);
96 }
97
98 let new_len = packed.len() / (SIZE_MULTIPLIER_U * SIZE_MULTIPLIER_U);
99 packed.truncate(new_len);
100
101 packed
102 }
103
104 #[test]
test_crop_2d_dct()105 fn test_crop_2d_dct() {
106 let packed: Vec<i32> = (0 .. 64).collect();
107 assert_eq!(
108 crop_2d_dct(packed.clone(), 8),
109 [
110 0, 1, 2, 3, // 4, 5, 6, 7
111 8, 9, 10, 11, // 12, 13, 14, 15
112 16, 17, 18, 19, // 20, 21, 22, 23,
113 24, 25, 26, 27, // 28, 29, 30, 31,
114 // 32 .. 64
115 ]
116 );
117 }
118
119 #[test]
test_transpose()120 fn test_transpose() {
121
122 }
123