LibVideo: Parameterize parsing compound references in the VP9 decoder

3 files changed, 100 insertions, 93 deletions

diff --git a/Userland/Libraries/LibVideo/VP9/Parser.cpp b/Userland/Libraries/LibVideo/VP9/Parser.cpp
index 752d1981ab..3c607efa1a 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Parser.cpp
@@ -1279,22 +1279,28 @@ DecoderErrorOr<void> Parser::read_ref_frames()
         return {};
     }
     ReferenceMode comp_mode;
+    Optional<bool> above_single = m_available_u ? m_above_single : Optional<bool>();
+    Optional<bool> left_single = m_available_l ? m_left_single : Optional<bool>();
+    Optional<bool> above_intra = m_available_u ? m_above_intra : Optional<bool>();
+    Optional<bool> left_intra = m_available_l ? m_left_intra : Optional<bool>();
+    Optional<ReferenceFrameType> above_ref_frame_0 = m_available_u ? m_above_ref_frame[0] : Optional<ReferenceFrameType>();
+    Optional<ReferenceFrameType> left_ref_frame_0 = m_available_l ? m_left_ref_frame[0] : Optional<ReferenceFrameType>();
     if (m_reference_mode == ReferenceModeSelect) {
-        Optional<bool> above_single = m_available_u ? m_above_single : Optional<bool>();
-        Optional<bool> left_single = m_available_l ? m_left_single : Optional<bool>();
-        Optional<bool> above_intra = m_available_u ? m_above_intra : Optional<bool>();
-        Optional<bool> left_intra = m_available_l ? m_left_intra : Optional<bool>();
-        Optional<ReferenceFrameType> above_ref_frame_0 = m_available_u ? m_above_ref_frame[0] : Optional<ReferenceFrameType>();
-        Optional<ReferenceFrameType> left_ref_frame_0 = m_available_l ? m_left_ref_frame[0] : Optional<ReferenceFrameType>();
         comp_mode = TRY_READ(TreeParser::parse_comp_mode(*m_bit_stream, *m_probability_tables, *m_syntax_element_counter, m_comp_fixed_ref, above_single, left_single, above_intra, left_intra, above_ref_frame_0, left_ref_frame_0));
     } else {
         comp_mode = m_reference_mode;
     }
     if (comp_mode == CompoundReference) {
-        auto idx = m_ref_frame_sign_bias[m_comp_fixed_ref];
-        auto comp_ref = TRY_READ(m_tree_parser->parse_tree(SyntaxElementType::CompRef));
-        m_ref_frame[idx] = m_comp_fixed_ref;
-        m_ref_frame[!idx] = m_comp_var_ref[comp_ref];
+        auto biased_reference_index = m_ref_frame_sign_bias[m_comp_fixed_ref];
+        auto inverse_biased_reference_index = biased_reference_index == 0 ? 1 : 0;
+
+        Optional<ReferenceFrameType> above_ref_frame_biased = m_available_u ? m_above_ref_frame[inverse_biased_reference_index] : Optional<ReferenceFrameType>();
+        Optional<ReferenceFrameType> left_ref_frame_biased = m_available_l ? m_left_ref_frame[inverse_biased_reference_index] : Optional<ReferenceFrameType>();
+        // FIXME: Create an enum for compound frame references using names Primary and Secondary.
+        auto comp_ref = TRY_READ(TreeParser::parse_comp_ref(*m_bit_stream, *m_probability_tables, *m_syntax_element_counter, m_comp_fixed_ref, m_comp_var_ref, above_single, left_single, above_intra, left_intra, above_ref_frame_0, left_ref_frame_0, above_ref_frame_biased, left_ref_frame_biased));
+
+        m_ref_frame[biased_reference_index] = m_comp_fixed_ref;
+        m_ref_frame[inverse_biased_reference_index] = m_comp_var_ref[comp_ref];
         return {};
     }
     auto single_ref_p1 = TRY_READ(m_tree_parser->parse_tree<bool>(SyntaxElementType::SingleRefP1));
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index f8e718f7fe..0375630ac9 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -346,13 +346,94 @@ ErrorOr<ReferenceMode> TreeParser::parse_comp_mode(BitStream& bit_stream, Probab
     return value;
 }
 
+ErrorOr<bool> TreeParser::parse_comp_ref(BitStream& bit_stream, ProbabilityTables const& probability_table, SyntaxElementCounter& counter, ReferenceFrameType comp_fixed_ref, ReferenceFramePair comp_var_ref, Optional<bool> above_single, Optional<bool> left_single, Optional<bool> above_intra, Optional<bool> left_intra, Optional<ReferenceFrameType> above_ref_frame_0, Optional<ReferenceFrameType> left_ref_frame_0, Optional<ReferenceFrameType> above_ref_frame_biased, Optional<ReferenceFrameType> left_ref_frame_biased)
+{
+    // FIXME: Above and left contexts should be in structs.
+
+    // Probabilities
+    u8 context;
+    if (above_intra.has_value() && left_intra.has_value()) {
+        if (above_intra.value() && left_intra.value()) {
+            context = 2;
+        } else if (left_intra.value()) {
+            if (above_single.value()) {
+                context = 1 + 2 * (above_ref_frame_0.value() != comp_var_ref[1]);
+            } else {
+                context = 1 + 2 * (above_ref_frame_biased.value() != comp_var_ref[1]);
+            }
+        } else if (above_intra.value()) {
+            if (left_single.value()) {
+                context = 1 + 2 * (left_ref_frame_0.value() != comp_var_ref[1]);
+            } else {
+                context = 1 + 2 * (left_ref_frame_biased != comp_var_ref[1]);
+            }
+        } else {
+            auto var_ref_above = above_single.value() ? above_ref_frame_0 : above_ref_frame_biased;
+            auto var_ref_left = left_single.value() ? left_ref_frame_0 : left_ref_frame_biased;
+            if (var_ref_above == var_ref_left && comp_var_ref[1] == var_ref_above) {
+                context = 0;
+            } else if (left_single.value() && above_single.value()) {
+                if ((var_ref_above == comp_fixed_ref && var_ref_left == comp_var_ref[0])
+                    || (var_ref_left == comp_fixed_ref && var_ref_above == comp_var_ref[0])) {
+                    context = 4;
+                } else if (var_ref_above == var_ref_left) {
+                    context = 3;
+                } else {
+                    context = 1;
+                }
+            } else if (left_single.value() || above_single.value()) {
+                auto vrfc = left_single.value() ? var_ref_above : var_ref_left;
+                auto rfs = above_single.value() ? var_ref_above : var_ref_left;
+                if (vrfc == comp_var_ref[1] && rfs != comp_var_ref[1]) {
+                    context = 1;
+                } else if (rfs == comp_var_ref[1] && vrfc != comp_var_ref[1]) {
+                    context = 2;
+                } else {
+                    context = 4;
+                }
+            } else if (var_ref_above == var_ref_left) {
+                context = 4;
+            } else {
+                context = 2;
+            }
+        }
+    } else if (above_intra.has_value()) {
+        if (above_intra.value()) {
+            context = 2;
+        } else {
+            if (above_single.value()) {
+                context = 3 * static_cast<u8>(above_ref_frame_0.value() != comp_var_ref[1]);
+            } else {
+                context = 4 * static_cast<u8>(above_ref_frame_biased.value() != comp_var_ref[1]);
+            }
+        }
+    } else if (left_intra.has_value()) {
+        if (left_intra.value()) {
+            context = 2;
+        } else {
+            if (left_single.value()) {
+                context = 3 * static_cast<u8>(left_ref_frame_0.value() != comp_var_ref[1]);
+            } else {
+                context = 4 * static_cast<u8>(left_ref_frame_biased != comp_var_ref[1]);
+            }
+        }
+    } else {
+        context = 2;
+    }
+
+    u8 probability = probability_table.comp_ref_prob()[context];
+
+    auto value = TRY(parse_tree_new<bool>(bit_stream, { binary_tree }, [&](u8) { return probability; }));
+    increment_counter(counter.m_counts_comp_ref[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
  */
 TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
 {
     switch (type) {
-    case SyntaxElementType::CompRef:
     case SyntaxElementType::SingleRefP1:
     case SyntaxElementType::SingleRefP2:
     case SyntaxElementType::MVSign:
@@ -386,8 +467,6 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
 u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
 {
     switch (type) {
-    case SyntaxElementType::CompRef:
-        return calculate_comp_ref_probability();
     case SyntaxElementType::SingleRefP1:
         return calculate_single_ref_p1_probability();
     case SyntaxElementType::SingleRefP2:
@@ -435,82 +514,6 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
 #define ABOVE_SINGLE m_decoder.m_above_single
 #define LEFT_SINGLE m_decoder.m_left_single
 
-u8 TreeParser::calculate_comp_ref_probability()
-{
-    auto fix_ref_idx = m_decoder.m_ref_frame_sign_bias[m_decoder.m_comp_fixed_ref];
-    auto var_ref_idx = !fix_ref_idx;
-    if (AVAIL_U && AVAIL_L) {
-        if (ABOVE_INTRA && LEFT_INTRA) {
-            m_ctx = 2;
-        } else if (LEFT_INTRA) {
-            if (ABOVE_SINGLE) {
-                m_ctx = 1 + 2 * (ABOVE_FRAME_0 != m_decoder.m_comp_var_ref[1]);
-            } else {
-                m_ctx = 1 + 2 * (m_decoder.m_above_ref_frame[var_ref_idx] != m_decoder.m_comp_var_ref[1]);
-            }
-        } else if (ABOVE_INTRA) {
-            if (LEFT_SINGLE) {
-                m_ctx = 1 + 2 * (LEFT_FRAME_0 != m_decoder.m_comp_var_ref[1]);
-            } else {
-                m_ctx = 1 + 2 * (m_decoder.m_left_ref_frame[var_ref_idx] != m_decoder.m_comp_var_ref[1]);
-            }
-        } else {
-            auto var_ref_above = m_decoder.m_above_ref_frame[ABOVE_SINGLE ? 0 : var_ref_idx];
-            auto var_ref_left = m_decoder.m_left_ref_frame[LEFT_SINGLE ? 0 : var_ref_idx];
-            if (var_ref_above == var_ref_left && m_decoder.m_comp_var_ref[1] == var_ref_above) {
-                m_ctx = 0;
-            } else if (LEFT_SINGLE && ABOVE_SINGLE) {
-                if ((var_ref_above == m_decoder.m_comp_fixed_ref && var_ref_left == m_decoder.m_comp_var_ref[0])
-                    || (var_ref_left == m_decoder.m_comp_fixed_ref && var_ref_above == m_decoder.m_comp_var_ref[0])) {
-                    m_ctx = 4;
-                } else if (var_ref_above == var_ref_left) {
-                    m_ctx = 3;
-                } else {
-                    m_ctx = 1;
-                }
-            } else if (LEFT_SINGLE || ABOVE_SINGLE) {
-                auto vrfc = LEFT_SINGLE ? var_ref_above : var_ref_left;
-                auto rfs = ABOVE_SINGLE ? var_ref_above : var_ref_left;
-                if (vrfc == m_decoder.m_comp_var_ref[1] && rfs != m_decoder.m_comp_var_ref[1]) {
-                    m_ctx = 1;
-                } else if (rfs == m_decoder.m_comp_var_ref[1] && vrfc != m_decoder.m_comp_var_ref[1]) {
-                    m_ctx = 2;
-                } else {
-                    m_ctx = 4;
-                }
-            } else if (var_ref_above == var_ref_left) {
-                m_ctx = 4;
-            } else {
-                m_ctx = 2;
-            }
-        }
-    } else if (AVAIL_U) {
-        if (ABOVE_INTRA) {
-            m_ctx = 2;
-        } else {
-            if (ABOVE_SINGLE) {
-                m_ctx = 3 * (ABOVE_FRAME_0 != m_decoder.m_comp_var_ref[1]);
-            } else {
-                m_ctx = 4 * (m_decoder.m_above_ref_frame[var_ref_idx] != m_decoder.m_comp_var_ref[1]);
-            }
-        }
-    } else if (AVAIL_L) {
-        if (LEFT_INTRA) {
-            m_ctx = 2;
-        } else {
-            if (LEFT_SINGLE) {
-                m_ctx = 3 * (LEFT_FRAME_0 != m_decoder.m_comp_var_ref[1]);
-            } else {
-                m_ctx = 4 * (m_decoder.m_left_ref_frame[var_ref_idx] != m_decoder.m_comp_var_ref[1]);
-            }
-        }
-    } else {
-        m_ctx = 2;
-    }
-
-    return m_decoder.m_probability_tables->comp_ref_prob()[m_ctx];
-}
-
 u8 TreeParser::calculate_single_ref_p1_probability()
 {
     if (AVAIL_U && AVAIL_L) {
@@ -725,9 +728,6 @@ void TreeParser::count_syntax_element(SyntaxElementType type, int value)
         increment_counter(count);
     };
     switch (type) {
-    case SyntaxElementType::CompRef:
-        increment(m_decoder.m_syntax_element_counter->m_counts_comp_ref[m_ctx][value]);
-        return;
     case SyntaxElementType::SingleRefP1:
         increment(m_decoder.m_syntax_element_counter->m_counts_single_ref[m_ctx][0][value]);
         return;
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 02add35c23..06b001e0ad 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -8,6 +8,7 @@
 #pragma once
 
 #include "BitStream.h"
+#include "Context.h"
 #include "Enums.h"
 #include "ProbabilityTables.h"
 #include "SyntaxElementCounter.h"
@@ -76,6 +77,7 @@ public:
     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);
     static ErrorOr<bool> parse_is_inter(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, Optional<bool> above_intra, Optional<bool> left_intra);
     static ErrorOr<ReferenceMode> parse_comp_mode(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, ReferenceFrameType comp_fixed_ref, Optional<bool> above_single, Optional<bool> left_single, Optional<bool> above_intra, Optional<bool> left_intra, Optional<ReferenceFrameType> above_ref_frame_0, Optional<ReferenceFrameType> left_ref_frame_0);
+    static ErrorOr<bool> parse_comp_ref(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, ReferenceFrameType comp_fixed_ref, ReferenceFramePair comp_var_ref, Optional<bool> above_single, Optional<bool> left_single, Optional<bool> above_intra, Optional<bool> left_intra, Optional<ReferenceFrameType> above_ref_frame_0, Optional<ReferenceFrameType> left_ref_frame_0, Optional<ReferenceFrameType> above_ref_frame_biased, Optional<ReferenceFrameType> left_ref_frame_biased);
 
     void set_default_intra_mode_variables(u8 idx, u8 idy)
     {
@@ -109,7 +111,6 @@ public:
     }
 
 private:
-    u8 calculate_comp_ref_probability();
     u8 calculate_single_ref_p1_probability();
     u8 calculate_single_ref_p2_probability();
     u8 calculate_token_probability(u8 node);