LibVideo: Parameterize TXSize parsing for the VP9 decoder

Userland/Libraries/LibVideo/VP9/Parser.cpp

@@ -1090,10 +1090,15 @@ bool Parser::seg_feature_active(u8 feature)
 DecoderErrorOr<void> Parser::read_tx_size(bool allow_select)
 {
     m_max_tx_size = max_txsize_lookup[m_mi_size];
-    if (allow_select && m_tx_mode == TXModeSelect && m_mi_size >= Block_8x8)
-        m_tx_size = TRY_READ(m_tree_parser->parse_tree<TXSize>(SyntaxElementType::TXSize));
-    else
+    if (allow_select && m_tx_mode == TXModeSelect && m_mi_size >= Block_8x8) {
+        Optional<bool> above_skip = m_available_u ? m_skips[get_image_index(m_mi_row - 1, m_mi_col)] : Optional<bool>();
+        Optional<bool> left_skip = m_available_l ? m_skips[get_image_index(m_mi_row, m_mi_col - 1)] : Optional<bool>();
+        Optional<TXSize> above_tx_size = m_available_u ? m_tx_sizes[get_image_index(m_mi_row - 1, m_mi_col)] : Optional<TXSize>();
+        Optional<TXSize> left_tx_size = m_available_l ? m_tx_sizes[get_image_index(m_mi_row, m_mi_col - 1)] : Optional<TXSize>();
+        m_tx_size = TRY_READ(TreeParser::parse_tx_size(*m_bit_stream, *m_probability_tables, *m_syntax_element_counter, m_max_tx_size, above_skip, left_skip, above_tx_size, left_tx_size));
+    } else {
         m_tx_size = min(m_max_tx_size, tx_mode_to_biggest_tx_size[m_tx_mode]);
+    }
     return {};
 }
 

Userland/Libraries/LibVideo/VP9/TreeParser.cpp

@@ -257,6 +257,38 @@ ErrorOr<bool> TreeParser::parse_skip(BitStream& bit_stream, ProbabilityTables co
     return value;
 }
 
+ErrorOr<TXSize> TreeParser::parse_tx_size(BitStream& bit_stream, ProbabilityTables const& probability_table, SyntaxElementCounter& counter, TXSize max_tx_size, Optional<bool> above_skip, Optional<bool> left_skip, Optional<TXSize> above_tx_size, Optional<TXSize> left_tx_size)
+{
+    // FIXME: Above and left contexts should be in structs.
+
+    // Tree
+    TreeParser::TreeSelection tree { tx_size_8_tree };
+    if (max_tx_size == TX_16x16)
+        tree = { tx_size_16_tree };
+    if (max_tx_size == TX_32x32)
+        tree = { tx_size_32_tree };
+
+    // Probabilities
+    auto above = max_tx_size;
+    auto left = max_tx_size;
+    if (above_skip.has_value() && !above_skip.value()) {
+        above = above_tx_size.value();
+    }
+    if (left_skip.has_value() && !left_skip.value()) {
+        left = left_tx_size.value();
+    }
+    if (!left_skip.has_value())
+        left = above;
+    if (!above_skip.has_value())
+        above = left;
+    auto context = (above + left) > max_tx_size;
+    u8 const* probabilities = probability_table.tx_probs()[max_tx_size][context];
+
+    auto value = TRY(parse_tree_new<TXSize>(bit_stream, tree, [&](u8 node) { return probabilities[node]; }));
+    increment_counter(counter.m_counts_tx_size[max_tx_size][context][value]);
+    return value;
+}
+
 /*
  * Select a tree value based on the type of syntax element being parsed, as well as some parser state, as specified in section 9.3.1
  */
@@ -273,12 +305,6 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
     case SyntaxElementType::MVBit:
     case SyntaxElementType::MoreCoefs:
         return { binary_tree };
-    case SyntaxElementType::TXSize:
-        if (m_decoder.m_max_tx_size == TX_32x32)
-            return { tx_size_32_tree };
-        if (m_decoder.m_max_tx_size == TX_16x16)
-            return { tx_size_16_tree };
-        return { tx_size_8_tree };
     case SyntaxElementType::MVJoint:
         return { mv_joint_tree };
     case SyntaxElementType::MVClass:
@@ -322,8 +348,6 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
     case SyntaxElementType::MVBit:
         VERIFY(m_mv_bit < MV_OFFSET_BITS);
         return m_decoder.m_probability_tables->mv_bits_prob()[m_mv_component][m_mv_bit];
-    case SyntaxElementType::TXSize:
-        return calculate_tx_size_probability(node);
     case SyntaxElementType::MVJoint:
         return m_decoder.m_probability_tables->mv_joint_probs()[node];
     case SyntaxElementType::MVClass:
@@ -619,28 +643,6 @@ u8 TreeParser::calculate_single_ref_p2_probability()
     return m_decoder.m_probability_tables->single_ref_prob()[m_ctx][1];
 }
 
-u8 TreeParser::calculate_tx_size_probability(u8 node)
-{
-    auto above = m_decoder.m_max_tx_size;
-    auto left = m_decoder.m_max_tx_size;
-    if (AVAIL_U) {
-        auto u_pos = (m_decoder.m_mi_row - 1) * m_decoder.m_mi_cols + m_decoder.m_mi_col;
-        if (!m_decoder.m_skips[u_pos])
-            above = m_decoder.m_tx_sizes[u_pos];
-    }
-    if (AVAIL_L) {
-        auto l_pos = m_decoder.m_mi_row * m_decoder.m_mi_cols + m_decoder.m_mi_col - 1;
-        if (!m_decoder.m_skips[l_pos])
-            left = m_decoder.m_tx_sizes[l_pos];
-    }
-    if (!AVAIL_L)
-        left = above;
-    if (!AVAIL_U)
-        above = left;
-    m_ctx = (above + left) > m_decoder.m_max_tx_size;
-    return m_decoder.m_probability_tables->tx_probs()[m_decoder.m_max_tx_size][m_ctx][node];
-}
-
 void TreeParser::set_tokens_variables(u8 band, u32 c, u32 plane, TXSize tx_size, u32 pos)
 {
     m_band = band;
@@ -744,9 +746,6 @@ void TreeParser::count_syntax_element(SyntaxElementType type, int value)
         increment(m_decoder.m_syntax_element_counter->m_counts_mv_bits[m_mv_component][m_mv_bit][value]);
         m_mv_bit = 0xFF;
         return;
-    case SyntaxElementType::TXSize:
-        increment(m_decoder.m_syntax_element_counter->m_counts_tx_size[m_decoder.m_max_tx_size][m_ctx][value]);
-        return;
     case SyntaxElementType::MVJoint:
         increment(m_decoder.m_syntax_element_counter->m_counts_mv_joint[value]);
         return;

Userland/Libraries/LibVideo/VP9/TreeParser.h

@@ -73,6 +73,7 @@ public:
     static ErrorOr<PredictionMode> parse_inter_mode(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, u8 mode_context_for_ref_frame_0);
     static ErrorOr<InterpolationFilter> parse_interpolation_filter(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, Optional<ReferenceFrameType> above_ref_frame, Optional<ReferenceFrameType> left_ref_frame, Optional<InterpolationFilter> above_interpolation_filter, Optional<InterpolationFilter> left_interpolation_filter);
     static ErrorOr<bool> parse_skip(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, Optional<bool> const& above_skip, Optional<bool> const& left_skip);
+    static ErrorOr<TXSize> parse_tx_size(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, TXSize max_tx_size, Optional<bool> above_skip, Optional<bool> left_skip, Optional<TXSize> above_tx_size, Optional<TXSize> left_tx_size);
 
     void set_default_intra_mode_variables(u8 idx, u8 idy)
     {
@@ -111,7 +112,6 @@ private:
     u8 calculate_comp_ref_probability();
     u8 calculate_single_ref_p1_probability();
     u8 calculate_single_ref_p2_probability();
-    u8 calculate_tx_size_probability(u8 node);
     u8 calculate_token_probability(u8 node);
     u8 calculate_more_coefs_probability();