LibVideo/VP9: Implement parsing Token and MoreCoefs trees

These elements were being used in the new tokens implementation, so
support for them in the TreeParser has been added.

Additionally, this uncovered a bug where the nonzero contexts were
being cleared with the wrong size.

3 files changed, 86 insertions, 10 deletions

diff --git a/Userland/Libraries/LibVideo/VP9/Parser.cpp b/Userland/Libraries/LibVideo/VP9/Parser.cpp
index 66e35eb779..2dd82741b3 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Parser.cpp
@@ -752,7 +752,7 @@ void Parser::clear_context(Vector<Vector<u8>>& context, size_t outer_size, size_
 
 bool Parser::clear_above_context()
 {
-    clear_context(m_above_nonzero_context, 2 * m_mi_cols, 3);
+    clear_context(m_above_nonzero_context, 3, 2 * m_mi_cols);
     clear_context(m_above_seg_pred_context, m_mi_cols);
     clear_context(m_above_partition_context, m_sb64_cols * 8);
     return true;
@@ -780,7 +780,7 @@ bool Parser::decode_tile()
 
 bool Parser::clear_left_context()
 {
-    clear_context(m_left_nonzero_context, 2 * m_mi_rows, 3);
+    clear_context(m_left_nonzero_context, 3, 2 * m_mi_rows);
     clear_context(m_left_seg_pred_context, m_mi_rows);
     clear_context(m_left_partition_context, m_sb64_rows * 8);
     return true;
@@ -1193,7 +1193,7 @@ BlockSubsize Parser::get_plane_block_size(u32 subsize, u8 plane)
 
 bool Parser::tokens(size_t plane, u32 start_x, u32 start_y, TXSize tx_size, u32 block_index)
 {
-    m_tree_parser->set_more_coefs_variables(start_x, start_y);
+    m_tree_parser->set_start_x_and_y(start_x, start_y);
     size_t segment_eob = 16 << (tx_size << 1);
     auto scan = get_scan(plane, tx_size, block_index);
     auto check_eob = true;
@@ -1201,7 +1201,7 @@ bool Parser::tokens(size_t plane, u32 start_x, u32 start_y, TXSize tx_size, u32
     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_band(band);
+        m_tree_parser->set_tokens_variables(band, c, plane, tx_size, pos);
         if (check_eob) {
             auto more_coefs = m_tree_parser->parse_tree<bool>(SyntaxElementType::MoreCoefs);
             if (!more_coefs)
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index 41d3ddf320..47b7a14e19 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -159,9 +159,9 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
     case SyntaxElementType::MVHP:
         break;
     case SyntaxElementType::Token:
-        break;
+        return calculate_token_probability(node);
     case SyntaxElementType::MoreCoefs:
-        break;
+        return calculate_more_coefs_probability();
     }
     TODO();
 }
@@ -561,6 +561,67 @@ u8 TreeParser::calculate_interp_filter_probability(u8 node)
     return m_decoder.m_probability_tables->interp_filter_probs()[m_ctx][node];
 }
 
+u8 TreeParser::calculate_token_probability(u8 node)
+{
+    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;
+}
+
+u8 TreeParser::calculate_more_coefs_probability()
+{
+    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;
+        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];
+            if (y4 + i < max_y)
+                left |= m_decoder.m_left_nonzero_context[m_plane][y4 + i];
+        }
+        m_ctx = 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 a = (i - 1) * n + j;
+        auto a2 = i * n + j - 1;
+        if (i > 0 && j > 0) {
+            if (m_decoder.m_tx_type == DCT_ADST) {
+                neighbor_0 = a;
+                neighbor_1 = a;
+            } else if (m_decoder.m_tx_type == ADST_DCT) {
+                neighbor_0 = a2;
+                neighbor_1 = a2;
+            } else {
+                neighbor_0 = a;
+                neighbor_1 = a2;
+            }
+        } else if (i > 0) {
+            neighbor_0 = a;
+            neighbor_1 = a;
+        } else {
+            neighbor_0 = a2;
+            neighbor_1 = a2;
+        }
+        m_ctx = (1 + m_decoder.m_token_cache[neighbor_0] + m_decoder.m_token_cache[neighbor_1]) >> 1;
+    }
+    return m_decoder.m_probability_tables->coef_probs()[m_tx_size][m_plane > 0][m_decoder.m_is_inter][m_band][m_ctx][0];
+}
+
 void TreeParser::count_syntax_element(SyntaxElementType type, int value)
 {
     switch (type) {
@@ -621,9 +682,11 @@ void TreeParser::count_syntax_element(SyntaxElementType type, int value)
     case SyntaxElementType::MVHP:
         break;
     case SyntaxElementType::Token:
-        break;
+        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:
-        break;
+        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;
     case SyntaxElementType::DefaultIntraMode:
     case SyntaxElementType::DefaultUVMode:
     case SyntaxElementType::SegmentID:
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 8c0084518a..46640dbf35 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -57,9 +57,16 @@ public:
         m_idy = idy;
     }
 
-    void set_band(u8 band) { m_band = band; }
+    void 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;
+    }
 
-    void set_more_coefs_variables(u32 start_x, u32 start_y)
+    void set_start_x_and_y(u32 start_x, u32 start_y)
     {
         m_start_x = start_x;
         m_start_y = start_y;
@@ -83,6 +90,8 @@ private:
     u8 calculate_tx_size_probability(u8 node);
     u8 calculate_inter_mode_probability(u8 node);
     u8 calculate_interp_filter_probability(u8 node);
+    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)
@@ -95,6 +104,10 @@ private:
     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 };
 };
 
 }