LibVideo/VP9: Refactor how TreeParser accesses decoder data

The TreeParser requires information about a lot of the decoder's
current state in order to parse syntax tree elements correctly, so
there has to be some communication between the Decoder and the
TreeParser. Previously, the Decoder would copy its state to the
TreeParser when it changed, however, this was a poor choice. Now,
the TreeParser simply has a reference to its owning Decoder, and
accesses its state directly.

4 files changed, 47 insertions, 88 deletions

diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.cpp b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
index f627aa27f0..cf817dce87 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
@@ -14,17 +14,14 @@ namespace Video::VP9 {
 
 Decoder::Decoder()
     : m_probability_tables(make<ProbabilityTables>())
-    , m_tree_parser(make<TreeParser>(*m_probability_tables))
+    , m_tree_parser(make<TreeParser>(*this))
 {
-    m_tree_parser->set_segmentation_tree_probs(m_segmentation_tree_probs);
 }
 
 bool Decoder::parse_frame(const ByteBuffer& frame_data)
 {
     m_bit_stream = make<BitStream>(frame_data.data(), frame_data.size());
     m_syntax_element_counter = make<SyntaxElementCounter>();
-    m_tree_parser->set_bit_stream(m_bit_stream);
-    m_tree_parser->set_syntax_element_counter(m_syntax_element_counter);
 
     if (!uncompressed_header())
         return false;
@@ -130,8 +127,6 @@ bool Decoder::uncompressed_header()
         }
     }
 
-    m_tree_parser->set_frame_is_intra(m_frame_is_intra);
-
     if (!m_error_resilient_mode) {
         m_refresh_frame_context = m_bit_stream->read_bit();
         m_frame_parallel_decoding_mode = m_bit_stream->read_bit();
@@ -522,9 +517,8 @@ u8 Decoder::inv_recenter_nonneg(u8 v, u8 m)
 
 bool Decoder::read_coef_probs()
 {
-    auto max_tx_size = tx_mode_to_biggest_tx_size[m_tx_mode];
-    m_tree_parser->set_max_tx_size(max_tx_size);
-    for (auto tx_size = TX_4x4; tx_size <= max_tx_size; tx_size = static_cast<TXSize>(static_cast<int>(tx_size) + 1)) {
+    m_max_tx_size = tx_mode_to_biggest_tx_size[m_tx_mode];
+    for (auto tx_size = TX_4x4; tx_size <= m_max_tx_size; tx_size = static_cast<TXSize>(static_cast<int>(tx_size) + 1)) {
         auto update_probs = m_bit_stream->read_literal(1);
         if (update_probs == 1) {
             for (auto i = 0; i < 2; i++) {
@@ -774,9 +768,9 @@ bool Decoder::decode_tile()
     for (auto row = m_mi_row_start; row < m_mi_row_end; row += 8) {
         if (!clear_left_context())
             return false;
-        m_tree_parser->set_row(row);
+        m_row = row;
         for (auto col = m_mi_col_start; col < m_mi_col_end; col += 8) {
-            m_tree_parser->set_col(col);
+            m_col = col;
             if (!decode_partition(row, col, Block_64x64))
                 return false;
         }
@@ -796,15 +790,11 @@ bool Decoder::decode_partition(u32 row, u32 col, u8 block_subsize)
 {
     if (row >= m_mi_rows || col >= m_mi_cols)
         return false;
-    auto num_8x8 = num_8x8_blocks_wide_lookup[block_subsize];
-    auto half_block_8x8 = num_8x8 >> 1;
-    auto has_rows = (row + half_block_8x8) < m_mi_rows;
-    auto has_cols = (col + half_block_8x8) < m_mi_cols;
-
-    m_tree_parser->set_has_rows(has_rows);
-    m_tree_parser->set_has_cols(has_cols);
-    m_tree_parser->set_block_subsize(block_subsize);
-    m_tree_parser->set_num_8x8(num_8x8);
+    m_block_subsize = block_subsize;
+    m_num_8x8 = num_8x8_blocks_wide_lookup[block_subsize];
+    auto half_block_8x8 = m_num_8x8 >> 1;
+    m_has_rows = (row + half_block_8x8) < m_mi_rows;
+    m_has_cols = (col + half_block_8x8) < m_mi_cols;
 
     auto partition = m_tree_parser->parse_tree(SyntaxElementType::Partition);
     dbgln("Parsed partition value {}", partition);
@@ -827,14 +817,10 @@ bool Decoder::decode_partition(u32 row, u32 col, u8 block_subsize)
 bool Decoder::decode_block(u32 row, u32 col, u8 subsize)
 {
     m_mi_row = row;
-    m_tree_parser->set_mi_row(m_mi_row);
     m_mi_col = col;
-    m_tree_parser->set_mi_col(m_mi_col);
     m_mi_size = subsize;
     m_available_u = row > 0;
-    m_tree_parser->set_available_u(m_available_u);
     m_available_l = col > m_mi_col_start;
-    m_tree_parser->set_available_l(m_available_l);
     if (!mode_info())
         return false;
     // FIXME: Finish implementing
diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.h b/Userland/Libraries/LibVideo/VP9/Decoder.h
index 159d993d0d..d6405e861b 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.h
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.h
@@ -17,6 +17,8 @@
 namespace Video::VP9 {
 
 class Decoder {
+    friend class TreeParser;
+
 public:
     Decoder();
     ~Decoder();
@@ -158,6 +160,13 @@ private:
     bool m_available_l { false };
     int m_segment_id { 0 };
     int m_skip { false };
+    u8 m_num_8x8 { 0 };
+    bool m_has_rows { false };
+    bool m_has_cols { false };
+    TXSize m_max_tx_size { TX_4x4 };
+    u8 m_block_subsize { 0 };
+    u32 m_row { 0 };
+    u32 m_col { 0 };
 
     bool m_use_hp { false };
 
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index bbb3069212..bb8f90e8c9 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -5,6 +5,7 @@
  */
 
 #include "TreeParser.h"
+#include "Decoder.h"
 #include "LookupTables.h"
 
 namespace Video::VP9 {
@@ -19,7 +20,7 @@ int TreeParser::parse_tree(SyntaxElementType type)
         auto tree = tree_selection.get_tree_value();
         int n = 0;
         do {
-            n = tree[n + m_bit_stream->read_bool(select_tree_probability(type, n >> 1))];
+            n = tree[n + m_decoder.m_bit_stream->read_bool(select_tree_probability(type, n >> 1))];
         } while (n > 0);
         value = -n;
     }
@@ -34,11 +35,11 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
 {
     switch (type) {
     case SyntaxElementType::Partition:
-        if (m_has_rows && m_has_cols)
+        if (m_decoder.m_has_rows && m_decoder.m_has_cols)
             return { partition_tree };
-        if (m_has_cols)
+        if (m_decoder.m_has_cols)
             return { cols_partition_tree };
-        if (m_has_rows)
+        if (m_decoder.m_has_rows)
             return { rows_partition_tree };
         return { PartitionSplit };
     case SyntaxElementType::DefaultIntraMode:
@@ -62,9 +63,9 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
     case SyntaxElementType::MoreCoefs:
         return { binary_tree };
     case SyntaxElementType::TXSize:
-        if (m_max_tx_size == TX_32x32)
+        if (m_decoder.m_max_tx_size == TX_32x32)
             return { tx_size_32_tree };
-        if (m_max_tx_size == TX_16x16)
+        if (m_decoder.m_max_tx_size == TX_16x16)
             return { tx_size_16_tree };
         return { tx_size_8_tree };
     case SyntaxElementType::InterMode:
@@ -80,7 +81,7 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
         return { mv_fr_tree };
     case SyntaxElementType::MVClass0HP:
     case SyntaxElementType::MVHP:
-        if (m_use_hp)
+        if (m_decoder.m_use_hp)
             return { binary_tree };
         return { 1 };
     case SyntaxElementType::Token:
@@ -108,7 +109,7 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
     case SyntaxElementType::UVMode:
         break;
     case SyntaxElementType::SegmentID:
-        return m_segmentation_tree_probs[node];
+        return m_decoder.m_segmentation_tree_probs[node];
     case SyntaxElementType::Skip:
         return calculate_skip_probability();
     case SyntaxElementType::SegIDPredicted:
@@ -158,9 +159,9 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
 u8 TreeParser::calculate_partition_probability(u8 node)
 {
     int node2;
-    if (m_has_rows && m_has_cols) {
+    if (m_decoder.m_has_rows && m_decoder.m_has_cols) {
         node2 = node;
-    } else if (m_has_cols) {
+    } else if (m_decoder.m_has_cols) {
         node2 = 1;
     } else {
         node2 = 2;
@@ -168,37 +169,37 @@ u8 TreeParser::calculate_partition_probability(u8 node)
 
     u32 above = 0;
     u32 left = 0;
-    auto bsl = mi_width_log2_lookup[m_block_subsize];
+    auto bsl = mi_width_log2_lookup[m_decoder.m_block_subsize];
     auto block_offset = mi_width_log2_lookup[Block_64x64] - bsl;
-    for (auto i = 0; i < m_num_8x8; i++) {
-        above |= m_above_partition_context[m_col + i];
-        left |= m_left_partition_context[m_row + i];
+    for (auto i = 0; i < m_decoder.m_num_8x8; i++) {
+        above |= m_decoder.m_above_partition_context[m_decoder.m_col + i];
+        left |= m_decoder.m_left_partition_context[m_decoder.m_row + i];
     }
     above = (above & (1 << block_offset)) > 0;
     left = (left & (1 << block_offset)) > 0;
     m_ctx = bsl * 4 + left * 2 + above;
-    if (m_frame_is_intra)
-        return m_probability_tables.kf_partition_probs()[m_ctx][node2];
-    return m_probability_tables.partition_probs()[m_ctx][node2];
+    if (m_decoder.m_frame_is_intra)
+        return m_decoder.m_probability_tables->kf_partition_probs()[m_ctx][node2];
+    return m_decoder.m_probability_tables->partition_probs()[m_ctx][node2];
 }
 
 u8 TreeParser::calculate_skip_probability()
 {
     m_ctx = 0;
-    if (m_available_u) {
+    if (m_decoder.m_available_u) {
         // FIXME: m_ctx += m_skips[m_mi_row - 1][m_mi_col];
     }
-    if (m_available_l) {
+    if (m_decoder.m_available_l) {
         // FIXME: m_ctx += m_skips[m_mi_row][m_mi_col - 1];
     }
-    return m_probability_tables.skip_prob()[m_ctx];
+    return m_decoder.m_probability_tables->skip_prob()[m_ctx];
 }
 
 void TreeParser::count_syntax_element(SyntaxElementType type, int value)
 {
     switch (type) {
     case SyntaxElementType::Partition:
-        m_syntax_element_counter->m_counts_partition[m_ctx][value]++;
+        m_decoder.m_syntax_element_counter->m_counts_partition[m_ctx][value]++;
         break;
     case SyntaxElementType::IntraMode:
         break;
@@ -207,7 +208,7 @@ void TreeParser::count_syntax_element(SyntaxElementType type, int value)
     case SyntaxElementType::UVMode:
         break;
     case SyntaxElementType::Skip:
-        m_syntax_element_counter->m_counts_skip[m_ctx][value]++;
+        m_decoder.m_syntax_element_counter->m_counts_skip[m_ctx][value]++;
         break;
     case SyntaxElementType::IsInter:
         break;
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 272ee729da..d22e378a52 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -13,10 +13,12 @@
 
 namespace Video::VP9 {
 
+class Decoder;
+
 class TreeParser {
 public:
-    explicit TreeParser(ProbabilityTables& probability_tables)
-        : m_probability_tables(probability_tables)
+    explicit TreeParser(Decoder& decoder)
+        : m_decoder(decoder)
     {
     }
 
@@ -44,52 +46,13 @@ public:
     u8 select_tree_probability(SyntaxElementType type, u8 node);
     void count_syntax_element(SyntaxElementType type, int value);
 
-    void set_bit_stream(BitStream* bit_stream) { m_bit_stream = bit_stream; }
-    void set_has_rows(bool has_rows) { m_has_rows = has_rows; }
-    void set_has_cols(bool has_cols) { m_has_cols = has_cols; }
-    void set_max_tx_size(TXSize max_tx_size) { m_max_tx_size = max_tx_size; }
-    void set_use_hp(bool use_hp) { m_use_hp = use_hp; }
-    void set_block_subsize(u8 block_subsize) { m_block_subsize = block_subsize; }
-    void set_num_8x8(u8 num_8x8) { m_num_8x8 = num_8x8; }
-    void set_above_partition_context(u8* above_partition_context) { m_above_partition_context = above_partition_context; }
-    void set_left_partition_context(u8* left_partition_context) { m_left_partition_context = left_partition_context; }
-    void set_col(u32 col) { m_col = col; }
-    void set_row(u32 row) { m_row = row; }
-    void set_frame_is_intra(bool frame_is_intra) { m_frame_is_intra = frame_is_intra; }
-    void set_syntax_element_counter(SyntaxElementCounter* syntax_element_counter) { m_syntax_element_counter = syntax_element_counter; }
-    void set_segmentation_tree_probs(u8* segmentation_tree_probs) { m_segmentation_tree_probs = segmentation_tree_probs; }
-    void set_available_u(bool available_u) { m_available_u = available_u; }
-    void set_available_l(bool available_l) { m_available_l = available_l; }
-    void set_mi_row(u32 mi_row) { m_mi_row = mi_row; }
-    void set_mi_col(u32 mi_col) { m_mi_col = mi_col; }
-
 private:
     u8 calculate_partition_probability(u8 node);
     u8 calculate_skip_probability();
 
-    ProbabilityTables& m_probability_tables;
-    BitStream* m_bit_stream { nullptr };
-    SyntaxElementCounter* m_syntax_element_counter { nullptr };
-
+    Decoder& m_decoder;
     // m_ctx is a member variable because it is required for syntax element counting (section 9.3.4)
     u8 m_ctx { 0 };
-
-    bool m_has_rows { false };
-    bool m_has_cols { false };
-    TXSize m_max_tx_size { TX_4x4 };
-    bool m_use_hp { false };
-    u8 m_block_subsize { 0 };
-    u8 m_num_8x8 { 0 };
-    u8* m_above_partition_context { nullptr };
-    u8* m_left_partition_context { nullptr };
-    u32 m_col { 0 };
-    u32 m_row { 0 };
-    bool m_frame_is_intra { false };
-    u8* m_segmentation_tree_probs { nullptr };
-    bool m_available_u { false };
-    bool m_available_l { false };
-    u32 m_mi_col { 0 };
-    u32 m_mi_row { 0 };
 };
 
 }