LibVideo: Create TokensContext struct for token parsing parameters

With the addition of this struct, both the bool to determine if coefs
should be parsed and the token parse itself can take specific
parameters.

This is the last step in parameterizing all the tree parsing, so the
old functions in TreeParser are now unused. This patch is very
satisfying :^)

There's still more work to be done to clean up how the parameters are
passed from Parser, but that's work for another day.

6 files changed, 62 insertions, 219 deletions

diff --git a/Userland/Libraries/LibVideo/VP9/Context.h b/Userland/Libraries/LibVideo/VP9/Context.h
index 758ce39e04..0408f372f1 100644
--- a/Userland/Libraries/LibVideo/VP9/Context.h
+++ b/Userland/Libraries/LibVideo/VP9/Context.h
@@ -30,4 +30,12 @@ struct Pair {
 typedef Pair<ReferenceFrameType> ReferenceFramePair;
 typedef Pair<MotionVector> MotionVectorPair;
 
+struct TokensContext {
+    TXSize m_tx_size;
+    bool m_is_uv_plane;
+    bool m_is_inter;
+    u8 m_band;
+    u8 m_context_index;
+};
+
 }
diff --git a/Userland/Libraries/LibVideo/VP9/Parser.cpp b/Userland/Libraries/LibVideo/VP9/Parser.cpp
index 2ea9056dfe..c0d894edf8 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Parser.cpp
@@ -19,7 +19,6 @@ namespace Video::VP9 {
 
 Parser::Parser(Decoder& decoder)
     : m_probability_tables(make<ProbabilityTables>())
-    , m_tree_parser(make<TreeParser>(*this))
     , m_decoder(decoder)
 {
 }
@@ -1463,21 +1462,20 @@ BlockSubsize Parser::get_plane_block_size(u32 subsize, u8 plane)
 
 DecoderErrorOr<bool> Parser::tokens(size_t plane, u32 start_x, u32 start_y, TXSize tx_size, u32 block_index)
 {
-    m_tree_parser->set_start_x_and_y(start_x, start_y);
-    size_t segment_eob = 16 << (tx_size << 1);
+    u32 segment_eob = 16 << (tx_size << 1);
     auto scan = get_scan(plane, tx_size, block_index);
     auto check_eob = true;
-    size_t c = 0;
+    u32 c = 0;
     for (; c < segment_eob; c++) {
         auto pos = scan[c];
         auto band = (tx_size == TX_4x4) ? coefband_4x4[c] : coefband_8x8plus[c];
-        m_tree_parser->set_tokens_variables(band, c, plane, tx_size, pos);
+        auto tokens_context = TreeParser::get_tokens_context(m_subsampling_x, m_subsampling_y, m_mi_rows, m_mi_cols, m_above_nonzero_context, m_left_nonzero_context, m_token_cache, tx_size, m_tx_type, plane, start_x, start_y, pos, m_is_inter, band, c);
         if (check_eob) {
-            auto more_coefs = TRY_READ(m_tree_parser->parse_tree<bool>(SyntaxElementType::MoreCoefs));
+            auto more_coefs = TRY_READ(TreeParser::parse_more_coefficients(*m_bit_stream, *m_probability_tables, *m_syntax_element_counter, tokens_context));
             if (!more_coefs)
                 break;
         }
-        auto token = TRY_READ(m_tree_parser->parse_tree<Token>(SyntaxElementType::Token));
+        auto token = TRY_READ(TreeParser::parse_token(*m_bit_stream, *m_probability_tables, *m_syntax_element_counter, tokens_context));
         m_token_cache[pos] = energy_class[token];
         if (token == ZeroToken) {
             m_tokens[pos] = 0;
@@ -1491,7 +1489,7 @@ DecoderErrorOr<bool> Parser::tokens(size_t plane, u32 start_x, u32 start_y, TXSi
     }
     auto non_zero = c > 0;
     m_eob_total += non_zero;
-    for (size_t i = c; i < segment_eob; i++)
+    for (u32 i = c; i < segment_eob; i++)
         m_tokens[scan[i]] = 0;
     return non_zero;
 }
diff --git a/Userland/Libraries/LibVideo/VP9/Parser.h b/Userland/Libraries/LibVideo/VP9/Parser.h
index 8ad62e4d51..c2ab06b65e 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.h
+++ b/Userland/Libraries/LibVideo/VP9/Parser.h
@@ -301,7 +301,6 @@ private:
     OwnPtr<BitStream> m_bit_stream;
     OwnPtr<ProbabilityTables> m_probability_tables;
     OwnPtr<SyntaxElementCounter> m_syntax_element_counter;
-    NonnullOwnPtr<TreeParser> m_tree_parser;
     Decoder& m_decoder;
 };
 
diff --git a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
index 740f7189bb..d7a723a820 100644
--- a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
+++ b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
@@ -11,37 +11,6 @@
 
 namespace Video::VP9 {
 
-enum class SyntaxElementType {
-    Partition,
-    DefaultIntraMode,
-    DefaultUVMode,
-    IntraMode,
-    SubIntraMode,
-    UVMode,
-    SegmentID,
-    Skip,
-    SegIDPredicted,
-    IsInter,
-    CompMode,
-    CompRef,
-    SingleRefP1,
-    SingleRefP2,
-    MVSign,
-    MVClass0Bit,
-    MVBit,
-    TXSize,
-    InterMode,
-    InterpFilter,
-    MVJoint,
-    MVClass,
-    MVClass0FR,
-    MVClass0HP,
-    MVFR,
-    MVHP,
-    Token,
-    MoreCoefs,
-};
-
 class SyntaxElementCounter final {
 public:
     /* (8.3) Clear Counts Process */
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index 03353236af..2ad8515c3f 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -14,40 +14,16 @@
 
 namespace Video::VP9 {
 
-template<typename T>
-ErrorOr<T> TreeParser::parse_tree(SyntaxElementType type)
-{
-    auto tree_selection = select_tree(type);
-    int value;
-    if (tree_selection.is_single_value()) {
-        value = tree_selection.single_value();
-    } else {
-        auto tree = tree_selection.tree();
-        int n = 0;
-        do {
-            n = tree[n + TRY(m_decoder.m_bit_stream->read_bool(select_tree_probability(type, n >> 1)))];
-        } while (n > 0);
-        value = -n;
-    }
-    count_syntax_element(type, value);
-    return static_cast<T>(value);
-}
-
-template ErrorOr<int> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<bool> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<u8> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<u32> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<PredictionMode> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<TXSize> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<InterpolationFilter> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<ReferenceMode> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<Token> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<MvClass> TreeParser::parse_tree(SyntaxElementType);
-template ErrorOr<MvJoint> TreeParser::parse_tree(SyntaxElementType);
+// Parsing of binary trees is handled here, as defined in sections 9.3.
+// Each syntax element is defined in its own section for each overarching section listed here:
+// - 9.3.1: Selection of the binary tree to be used.
+// - 9.3.2: Probability selection based on context and often the node of the tree.
+// - 9.3.4: Counting each syntax element when it is read.
 
 template<typename OutputType>
 inline ErrorOr<OutputType> parse_tree_new(BitStream& bit_stream, TreeParser::TreeSelection tree_selection, Function<u8(u8)> const& probability_getter)
 {
+    // 9.3.3: The tree decoding function.
     if (tree_selection.is_single_value())
         return static_cast<OutputType>(tree_selection.single_value());
 
@@ -652,86 +628,38 @@ ErrorOr<bool> TreeParser::parse_motion_vector_hp(BitStream& bit_stream, Probabil
     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::MoreCoefs:
-        return { binary_tree };
-    case SyntaxElementType::Token:
-        return { token_tree };
-    default:
-        break;
-    }
-    VERIFY_NOT_REACHED();
-}
-
-/*
- * Select a probability with which to read a boolean when decoding a tree, as specified in section 9.3.2
- */
-u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
+TokensContext TreeParser::get_tokens_context(bool subsampling_x, bool subsampling_y, u32 rows, u32 columns, Array<Vector<bool>, 3> const& above_nonzero_context, Array<Vector<bool>, 3> const& left_nonzero_context, u8 token_cache[1024], TXSize tx_size, u8 tx_type, u8 plane, u32 start_x, u32 start_y, u32 position, bool is_inter, u8 band, u32 c)
 {
-    switch (type) {
-    case SyntaxElementType::Token:
-        return calculate_token_probability(node);
-    case SyntaxElementType::MoreCoefs:
-        return calculate_more_coefs_probability();
-    default:
-        break;
-    }
-    VERIFY_NOT_REACHED();
-}
-
-#define ABOVE_FRAME_0 m_decoder.m_above_ref_frame[0]
-#define ABOVE_FRAME_1 m_decoder.m_above_ref_frame[1]
-#define LEFT_FRAME_0 m_decoder.m_left_ref_frame[0]
-#define LEFT_FRAME_1 m_decoder.m_left_ref_frame[1]
-#define AVAIL_U m_decoder.m_available_u
-#define AVAIL_L m_decoder.m_available_l
-#define ABOVE_INTRA m_decoder.m_above_intra
-#define LEFT_INTRA m_decoder.m_left_intra
-#define ABOVE_SINGLE m_decoder.m_above_single
-#define LEFT_SINGLE m_decoder.m_left_single
-
-void TreeParser::set_tokens_variables(u8 band, u32 c, u32 plane, TXSize tx_size, u32 pos)
-{
-    m_band = band;
-    m_c = c;
-    m_plane = plane;
-    m_tx_size = tx_size;
-    m_pos = pos;
-
-    if (m_c == 0) {
-        auto sx = m_plane > 0 ? m_decoder.m_subsampling_x : 0;
-        auto sy = m_plane > 0 ? m_decoder.m_subsampling_y : 0;
-        auto max_x = (2 * m_decoder.m_mi_cols) >> sx;
-        auto max_y = (2 * m_decoder.m_mi_rows) >> sy;
-        u8 numpts = 1 << m_tx_size;
-        auto x4 = m_start_x >> 2;
-        auto y4 = m_start_y >> 2;
+    u8 context;
+    if (c == 0) {
+        auto sx = plane > 0 ? subsampling_x : false;
+        auto sy = plane > 0 ? subsampling_y : false;
+        auto max_x = (2 * columns) >> sx;
+        auto max_y = (2 * rows) >> sy;
+        u8 numpts = 1 << tx_size;
+        auto x4 = start_x >> 2;
+        auto y4 = start_y >> 2;
         u32 above = 0;
         u32 left = 0;
         for (size_t i = 0; i < numpts; i++) {
             if (x4 + i < max_x)
-                above |= m_decoder.m_above_nonzero_context[m_plane][x4 + i];
+                above |= above_nonzero_context[plane][x4 + i];
             if (y4 + i < max_y)
-                left |= m_decoder.m_left_nonzero_context[m_plane][y4 + i];
+                left |= left_nonzero_context[plane][y4 + i];
         }
-        m_ctx = above + left;
+        context = above + left;
     } else {
         u32 neighbor_0, neighbor_1;
-        auto n = 4 << m_tx_size;
-        auto i = m_pos / n;
-        auto j = m_pos % n;
+        auto n = 4 << tx_size;
+        auto i = position / n;
+        auto j = position % n;
         auto a = i > 0 ? (i - 1) * n + j : 0;
         auto a2 = i * n + j - 1;
         if (i > 0 && j > 0) {
-            if (m_decoder.m_tx_type == DCT_ADST) {
+            if (tx_type == DCT_ADST) {
                 neighbor_0 = a;
                 neighbor_1 = a;
-            } else if (m_decoder.m_tx_type == ADST_DCT) {
+            } else if (tx_type == ADST_DCT) {
                 neighbor_0 = a2;
                 neighbor_1 = a2;
             } else {
@@ -745,43 +673,36 @@ void TreeParser::set_tokens_variables(u8 band, u32 c, u32 plane, TXSize tx_size,
             neighbor_0 = a2;
             neighbor_1 = a2;
         }
-        m_ctx = (1 + m_decoder.m_token_cache[neighbor_0] + m_decoder.m_token_cache[neighbor_1]) >> 1;
+        context = (1 + token_cache[neighbor_0] + token_cache[neighbor_1]) >> 1;
     }
-}
 
-u8 TreeParser::calculate_more_coefs_probability()
-{
-    return m_decoder.m_probability_tables->coef_probs()[m_tx_size][m_plane > 0][m_decoder.m_is_inter][m_band][m_ctx][0];
+    return TokensContext { tx_size, plane > 0, is_inter, band, context };
 }
 
-u8 TreeParser::calculate_token_probability(u8 node)
+ErrorOr<bool> TreeParser::parse_more_coefficients(BitStream& bit_stream, ProbabilityTables const& probability_table, SyntaxElementCounter& counter, TokensContext const& context)
 {
-    auto prob = m_decoder.m_probability_tables->coef_probs()[m_tx_size][m_plane > 0][m_decoder.m_is_inter][m_band][m_ctx][min(2, 1 + node)];
-    if (node < 2)
-        return prob;
-    auto x = (prob - 1) / 2;
-    auto& pareto_table = m_decoder.m_probability_tables->pareto_table();
-    if (prob & 1)
-        return pareto_table[x][node - 2];
-    return (pareto_table[x][node - 2] + pareto_table[x + 1][node - 2]) >> 1;
+    auto probability = probability_table.coef_probs()[context.m_tx_size][context.m_is_uv_plane][context.m_is_inter][context.m_band][context.m_context_index][0];
+    auto value = TRY(parse_tree_new<u8>(bit_stream, { binary_tree }, [&](u8) { return probability; }));
+    increment_counter(counter.m_counts_more_coefs[context.m_tx_size][context.m_is_uv_plane][context.m_is_inter][context.m_band][context.m_context_index][value]);
+    return value;
 }
 
-void TreeParser::count_syntax_element(SyntaxElementType type, int value)
+ErrorOr<Token> TreeParser::parse_token(BitStream& bit_stream, ProbabilityTables const& probability_table, SyntaxElementCounter& counter, TokensContext const& context)
 {
-    auto increment = [](u8& count) {
-        increment_counter(count);
+    Function<u8(u8)> probability_getter = [&](u8 node) -> u8 {
+        auto prob = probability_table.coef_probs()[context.m_tx_size][context.m_is_uv_plane][context.m_is_inter][context.m_band][context.m_context_index][min(2, 1 + node)];
+        if (node < 2)
+            return prob;
+        auto x = (prob - 1) / 2;
+        auto const& pareto_table = probability_table.pareto_table();
+        if ((prob & 1) != 0)
+            return pareto_table[x][node - 2];
+        return (pareto_table[x][node - 2] + pareto_table[x + 1][node - 2]) >> 1;
     };
-    switch (type) {
-    case SyntaxElementType::Token:
-        increment(m_decoder.m_syntax_element_counter->m_counts_token[m_tx_size][m_plane > 0][m_decoder.m_is_inter][m_band][m_ctx][min(2, value)]);
-        return;
-    case SyntaxElementType::MoreCoefs:
-        increment(m_decoder.m_syntax_element_counter->m_counts_more_coefs[m_tx_size][m_plane > 0][m_decoder.m_is_inter][m_band][m_ctx][value]);
-        return;
-    default:
-        break;
-    }
-    VERIFY_NOT_REACHED();
+
+    auto value = TRY(parse_tree_new<Token>(bit_stream, { token_tree }, probability_getter));
+    increment_counter(counter.m_counts_token[context.m_tx_size][context.m_is_uv_plane][context.m_is_inter][context.m_band][context.m_context_index][min(2, value)]);
+    return value;
 }
 
 }
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 0227c3a7a6..8d833602c4 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -19,11 +19,6 @@ class Parser;
 
 class TreeParser {
 public:
-    explicit TreeParser(Parser& decoder)
-        : m_decoder(decoder)
-    {
-    }
-
     // FIXME: Move or remove this class once it is unused in the header.
     class TreeSelection {
     public:
@@ -53,16 +48,6 @@ public:
         TreeSelectionValue m_value;
     };
 
-    /* (9.3.3) */
-    template<typename T = int>
-    ErrorOr<T> parse_tree(SyntaxElementType type);
-    /* (9.3.1) */
-    TreeSelection select_tree(SyntaxElementType type);
-    /* (9.3.2) */
-    u8 select_tree_probability(SyntaxElementType type, u8 node);
-    /* (9.3.4) */
-    void count_syntax_element(SyntaxElementType type, int value);
-
     static ErrorOr<Partition> parse_partition(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, bool has_rows, bool has_columns, BlockSubsize block_subsize, u8 num_8x8, Vector<u8> const& above_partition_context, Vector<u8> const& left_partition_context, u32 row, u32 column, bool frame_is_intra);
     static ErrorOr<PredictionMode> parse_default_intra_mode(BitStream&, ProbabilityTables const&, BlockSubsize mi_size, Optional<Array<PredictionMode, 4> const&> above_context, Optional<Array<PredictionMode, 4> const&> left_context, PredictionMode block_sub_modes[4], u8 index_x, u8 index_y);
     static ErrorOr<PredictionMode> parse_default_uv_mode(BitStream&, ProbabilityTables const&, PredictionMode y_mode);
@@ -91,46 +76,9 @@ public:
     static ErrorOr<u8> parse_motion_vector_fr(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, u8 component);
     static ErrorOr<bool> parse_motion_vector_hp(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, u8 component, bool use_hp);
 
-    void set_default_intra_mode_variables(u8 idx, u8 idy)
-    {
-        m_idx = idx;
-        m_idy = idy;
-    }
-
-    void set_tokens_variables(u8 band, u32 c, u32 plane, TXSize tx_size, u32 pos);
-
-    void set_start_x_and_y(u32 start_x, u32 start_y)
-    {
-        m_start_x = start_x;
-        m_start_y = start_y;
-    }
-
-private:
-    u8 calculate_token_probability(u8 node);
-    u8 calculate_more_coefs_probability();
-
-    Parser& m_decoder;
-    // m_ctx is a member variable because it is required for syntax element counting (section 9.3.4)
-    u8 m_ctx { 0 };
-
-    // These are variables necessary for parsing tree data, but aren't useful otherwise, so they're
-    // not stored in the Decoder itself.
-    u8 m_idx { 0 };
-    u8 m_idy { 0 };
-    u8 m_band { 0 };
-    u32 m_start_x { 0 };
-    u32 m_start_y { 0 };
-    u32 m_c { 0 };
-    u32 m_plane { 0 };
-    TXSize m_tx_size;
-    u32 m_pos { 0 };
-    u8 m_mv_component { 0 };
-    // 0xFF indicates the value has not been set.
-    // parse_mv_bit should be called to set this.
-    u8 m_mv_bit { 0xFF };
-    // 0xFF indicates the value has not been set.
-    // parse_mv_class0_fr should be called to set this.
-    u8 m_mv_class0_bit { 0xFF };
+    static TokensContext get_tokens_context(bool subsampling_x, bool subsampling_y, u32 rows, u32 columns, Array<Vector<bool>, 3> const& above_nonzero_context, Array<Vector<bool>, 3> const& left_nonzero_context, u8 token_cache[1024], TXSize tx_size, u8 tx_type, u8 plane, u32 start_x, u32 start_y, u32 position, bool is_inter, u8 band, u32 c);
+    static ErrorOr<bool> parse_more_coefficients(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, TokensContext const& context);
+    static ErrorOr<Token> parse_token(BitStream&, ProbabilityTables const&, SyntaxElementCounter&, TokensContext const& context);
 };
 
 struct PartitionTreeContext {