Fix selection of transform function for 4x4 blocks

DST function was returned for inter luma transform blocks of size 4x4
even though they must use DCT. Fixed by checking the prediction mode of
the block in addition to whether it is chroma or luma.

src/strategies/avx2/quant-avx2.c

@@ -375,7 +375,7 @@ int kvz_quantize_residual_avx2(encoder_state_t *const state,
     kvz_transformskip(state->encoder_control, residual, coeff, width);
   }
   else {
-    kvz_transform2d(state->encoder_control, residual, coeff, width, (color == COLOR_Y ? 0 : 65535));
+    kvz_transform2d(state->encoder_control, residual, coeff, width, color, cur_cu->type);
   }
 
   // Quantize coeffs. (coeff -> coeff_out)
@@ -408,7 +408,7 @@ int kvz_quantize_residual_avx2(encoder_state_t *const state,
       kvz_itransformskip(state->encoder_control, residual, coeff, width);
     }
     else {
-      kvz_itransform2d(state->encoder_control, residual, coeff, width, (color == COLOR_Y ? 0 : 65535));
+      kvz_itransform2d(state->encoder_control, residual, coeff, width, color, cur_cu->type);
     }
 
     // Get quantized reconstruction. (residual + pred_in -> rec_out)

src/strategies/generic/quant-generic.c

@@ -208,7 +208,7 @@ int kvz_quantize_residual_generic(encoder_state_t *const state,
     kvz_transformskip(state->encoder_control, residual, coeff, width);
   }
   else {
-    kvz_transform2d(state->encoder_control, residual, coeff, width, (color == COLOR_Y ? 0 : 65535));
+    kvz_transform2d(state->encoder_control, residual, coeff, width, color, cur_cu->type);
   }
 
   // Quantize coeffs. (coeff -> coeff_out)
@@ -246,7 +246,7 @@ int kvz_quantize_residual_generic(encoder_state_t *const state,
       kvz_itransformskip(state->encoder_control, residual, coeff, width);
     }
     else {
-      kvz_itransform2d(state->encoder_control, residual, coeff, width, (color == COLOR_Y ? 0 : 65535));
+      kvz_itransform2d(state->encoder_control, residual, coeff, width, color, cur_cu->type);
     }
 
     // Get quantized reconstruction. (residual + pred_in -> rec_out)

src/strategies/strategies-dct.c

@@ -55,21 +55,22 @@ int kvz_strategy_register_dct(void* opaque, uint8_t bitdepth) {
 
 
 /**
-* \brief  Get a function that calculates SAD for NxN block.
+ * \brief  Get a function that performs the transform for a block.
  *
-* \param n  Width of the region for which SAD is calculated.
+ * \param width    Width of the region
+ * \param color    Color plane
+ * \param type     Prediction type
  *
-* \returns  Pointer to cost_16bit_nxn_func.
+ * \returns Pointer to the function.
  */
-dct_func * kvz_get_dct_func(int8_t width, int32_t mode)
+dct_func * kvz_get_dct_func(int8_t width, color_t color, cu_type_t type)
 {
   switch (width) {
   case 4:
-    switch (mode){
-    case 65535:
-      return kvz_dct_4x4;
-    default:
+    if (color == COLOR_Y && type == CU_INTRA) {
       return kvz_fast_forward_dst_4x4;
+    } else {
+      return kvz_dct_4x4;
     }
   case 8:
     return kvz_dct_8x8;
@@ -83,21 +84,22 @@ dct_func * kvz_get_dct_func(int8_t width, int32_t mode)
 }
 
 /**
-* \brief  Get a function that calculates SAD for NxN block.
+ * \brief  Get a function that performs the inverse transform for a block.
  *
-* \param n  Width of the region for which SAD is calculated.
+ * \param width    Width of the region
+ * \param color    Color plane
+ * \param type     Prediction type
  *
-* \returns  Pointer to cost_16bit_nxn_func.
+ * \returns Pointer to the function.
  */
-dct_func * kvz_get_idct_func(int8_t width, int32_t mode)
+dct_func * kvz_get_idct_func(int8_t width, color_t color, cu_type_t type)
 {
   switch (width) {
   case 4:
-    switch (mode){
-    case 65535:
-      return kvz_idct_4x4;
-    default:
+    if (color == COLOR_Y && type == CU_INTRA) {
       return kvz_fast_inverse_dst_4x4;
+    } else {
+      return kvz_idct_4x4;
     }
   case 8:
     return kvz_idct_8x8;

src/strategies/strategies-dct.h

@@ -27,7 +27,7 @@
  */
 
 #include "global.h" // IWYU pragma: keep
-
+#include "cu.h"
 
 typedef unsigned (dct_func)(int8_t bitdepth, const int16_t *input, int16_t *output);
 
@@ -49,8 +49,9 @@ extern dct_func * kvz_idct_32x32;
 
 
 int kvz_strategy_register_dct(void* opaque, uint8_t bitdepth);
-dct_func * kvz_get_dct_func(int8_t width, int32_t mode);
-dct_func * kvz_get_idct_func(int8_t width, int32_t mode);
+dct_func * kvz_get_dct_func(int8_t width, color_t color, cu_type_t type);
+dct_func * kvz_get_idct_func(int8_t width, color_t color, cu_type_t type);
+
 
 
 #define STRATEGIES_DCT_EXPORTS \

src/transform.c

@@ -186,15 +186,25 @@ void kvz_itransformskip(const encoder_control_t * const encoder, int16_t *block,
  * \param coeff transform coefficients
  * \param block_size width of transform
  */
-void kvz_transform2d(const encoder_control_t * const encoder, int16_t *block, int16_t *coeff, int8_t block_size, int32_t mode)
+void kvz_transform2d(const encoder_control_t * const encoder,
+                     int16_t *block,
+                     int16_t *coeff,
+                     int8_t block_size,
+                     color_t color,
+                     cu_type_t type)
 {
-  dct_func *dct_func = kvz_get_dct_func(block_size, mode);  
+  dct_func *dct_func = kvz_get_dct_func(block_size, color, type);
   dct_func(encoder->bitdepth, block, coeff);
 }
 
-void kvz_itransform2d(const encoder_control_t * const encoder, int16_t *block, int16_t *coeff, int8_t block_size, int32_t mode)
+void kvz_itransform2d(const encoder_control_t * const encoder,
+                      int16_t *block,
+                      int16_t *coeff,
+                      int8_t block_size,
+                      color_t color,
+                      cu_type_t type)
 {
-  dct_func *idct_func = kvz_get_idct_func(block_size, mode);
+  dct_func *idct_func = kvz_get_idct_func(block_size, color, type);
   idct_func(encoder->bitdepth, coeff, block);
 }
 

src/transform.h

@@ -38,8 +38,18 @@ extern const int16_t kvz_g_inv_quant_scales[6];
 void kvz_transformskip(const encoder_control_t *encoder, int16_t *block,int16_t *coeff, int8_t block_size);
 void kvz_itransformskip(const encoder_control_t *encoder, int16_t *block,int16_t *coeff, int8_t block_size);
 
-void kvz_transform2d(const encoder_control_t *encoder, int16_t *block,int16_t *coeff, int8_t block_size, int32_t mode);
-void kvz_itransform2d(const encoder_control_t *encoder, int16_t *block,int16_t *coeff, int8_t block_size, int32_t mode);
+void kvz_transform2d(const encoder_control_t * const encoder,
+                     int16_t *block,
+                     int16_t *coeff,
+                     int8_t block_size,
+                     color_t color,
+                     cu_type_t type);
+void kvz_itransform2d(const encoder_control_t * const encoder,
+                      int16_t *block,
+                      int16_t *coeff,
+                      int8_t block_size,
+                      color_t color,
+                      cu_type_t type);
 
 int32_t kvz_get_scaled_qp(int8_t type, int8_t qp, int8_t qp_offset);