LibVideo/VP9: Successfully parse partition syntax element

6 files changed, 165 insertions, 66 deletions

diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.cpp b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
index 89f6195cfd..acb7b2ca27 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
@@ -23,6 +23,7 @@ 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;
@@ -743,13 +744,20 @@ bool Decoder::decode_tiles()
     return true;
 }
 
+template<typename T>
+void clear_context(T* context, size_t size)
+{
+    if (!(*context))
+        *context = static_cast<T>(malloc(size));
+    else
+        __builtin_memset(*context, 0, size);
+}
+
 bool Decoder::clear_above_context()
 {
-    // FIXME
-    // When this function is invoked the arrays AboveNonzeroContext, AbovePartitionContext, AboveSegPredContext should be set equal to 0.
-    // AboveNonzeroContext[0..2][0..MiCols*2-1] = 0
-    // AboveSegPredContext[0..MiCols-1] = 0
-    // AbovePartitionContext[0..Sb64Cols*8-1] = 0
+    clear_context(&m_above_nonzero_context, sizeof(u8) * 3 * m_mi_cols * 2);
+    clear_context(&m_above_seg_pred_context, sizeof(u8) * m_mi_cols);
+    clear_context(&m_above_partition_context, sizeof(u8) * m_sb64_cols * 8);
     return true;
 }
 
@@ -777,11 +785,9 @@ bool Decoder::decode_tile()
 
 bool Decoder::clear_left_context()
 {
-    // FIXME
-    // When this function is invoked the arrays LeftNonzeroContext, LeftPartitionContext, LeftSegPredContext should be set equal to 0.
-    // LeftNonzeroContext[0..2][0..MiRows*2-1] = 0
-    // LeftSegPredContext[0..MiRows-1] = 0
-    // LeftPartitionContext[0..Sb64Rows*8-1] = 0
+    clear_context(&m_left_nonzero_context, sizeof(u8) * 3 * m_mi_rows * 2);
+    clear_context(&m_left_seg_pred_context, sizeof(u8) * m_mi_rows);
+    clear_context(&m_left_partition_context, sizeof(u8) * m_sb64_rows * 8);
     return true;
 }
 
@@ -815,4 +821,20 @@ void Decoder::dump_info()
     dbgln("Interpolation filter: {}", (u8)m_interpolation_filter);
 }
 
+Decoder::~Decoder()
+{
+    if (m_above_nonzero_context)
+        free(m_above_nonzero_context);
+    if (m_left_nonzero_context)
+        free(m_left_nonzero_context);
+    if (m_above_seg_pred_context)
+        free(m_above_seg_pred_context);
+    if (m_left_seg_pred_context)
+        free(m_left_seg_pred_context);
+    if (m_above_partition_context)
+        free(m_above_partition_context);
+    if (m_left_partition_context)
+        free(m_left_partition_context);
+}
+
 }
diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.h b/Userland/Libraries/LibVideo/VP9/Decoder.h
index 6eaea166f5..f7740eacf8 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.h
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.h
@@ -19,6 +19,7 @@ namespace Video::VP9 {
 class Decoder {
 public:
     Decoder();
+    ~Decoder();
     bool parse_frame(const ByteBuffer&);
     void dump_info();
 
@@ -128,6 +129,12 @@ private:
     i8 m_loop_filter_ref_deltas[MAX_REF_FRAMES];
     i8 m_loop_filter_mode_deltas[2];
 
+    u8** m_above_nonzero_context { nullptr };
+    u8** m_left_nonzero_context { nullptr };
+    u8* m_above_seg_pred_context { nullptr };
+    u8* m_left_seg_pred_context { nullptr };
+    u8* m_above_partition_context { nullptr };
+    u8* m_left_partition_context { nullptr };
     u32 m_mi_row_start { 0 };
     u32 m_mi_row_end { 0 };
     u32 m_mi_col_start { 0 };
diff --git a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.cpp b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.cpp
index f6ad6bd820..d3a095dd38 100644
--- a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.cpp
+++ b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.cpp
@@ -10,28 +10,28 @@ namespace Video::VP9 {
 
 void SyntaxElementCounter::clear_counts()
 {
-    __builtin_memset(counts_intra_mode, 0, BLOCK_SIZE_GROUPS * INTRA_MODES);
-    __builtin_memset(counts_uv_mode, 0, INTRA_MODES * INTRA_MODES);
-    __builtin_memset(counts_partition, 0, PARTITION_CONTEXTS * PARTITION_TYPES);
-    __builtin_memset(counts_interp_filter, 0, INTERP_FILTER_CONTEXTS * SWITCHABLE_FILTERS);
-    __builtin_memset(counts_inter_mode, 0, INTER_MODE_CONTEXTS * INTER_MODES);
-    __builtin_memset(counts_tx_size, 0, TX_SIZES * TX_SIZE_CONTEXTS * TX_SIZES);
-    __builtin_memset(counts_is_inter, 0, IS_INTER_CONTEXTS * 2);
-    __builtin_memset(counts_comp_mode, 0, COMP_MODE_CONTEXTS * 2);
-    __builtin_memset(counts_single_ref, 0, REF_CONTEXTS * 2 * 2);
-    __builtin_memset(counts_comp_ref, 0, REF_CONTEXTS * 2);
-    __builtin_memset(counts_skip, 0, SKIP_CONTEXTS * 2);
-    __builtin_memset(counts_mv_joint, 0, MV_JOINTS);
-    __builtin_memset(counts_mv_sign, 0, 2 * 2);
-    __builtin_memset(counts_mv_class, 0, 2 * MV_CLASSES);
-    __builtin_memset(counts_mv_class0_bit, 0, 2 * CLASS0_SIZE);
-    __builtin_memset(counts_mv_class0_fr, 0, 2 * CLASS0_SIZE * MV_FR_SIZE);
-    __builtin_memset(counts_mv_class0_hp, 0, 2 * 2);
-    __builtin_memset(counts_mv_bits, 0, 2 * MV_OFFSET_BITS * 2);
-    __builtin_memset(counts_mv_fr, 0, 2 * MV_FR_SIZE);
-    __builtin_memset(counts_mv_hp, 0, 2 * 2);
-    __builtin_memset(counts_token, 0, TX_SIZES * BLOCK_TYPES * REF_TYPES * COEF_BANDS * PREV_COEF_CONTEXTS * UNCONSTRAINED_NODES);
-    __builtin_memset(counts_more_coefs, 0, TX_SIZES * BLOCK_TYPES * REF_TYPES * COEF_BANDS * PREV_COEF_CONTEXTS * 2);
+    __builtin_memset(m_counts_intra_mode, 0, BLOCK_SIZE_GROUPS * INTRA_MODES);
+    __builtin_memset(m_counts_uv_mode, 0, INTRA_MODES * INTRA_MODES);
+    __builtin_memset(m_counts_partition, 0, PARTITION_CONTEXTS * PARTITION_TYPES);
+    __builtin_memset(m_counts_interp_filter, 0, INTERP_FILTER_CONTEXTS * SWITCHABLE_FILTERS);
+    __builtin_memset(m_counts_inter_mode, 0, INTER_MODE_CONTEXTS * INTER_MODES);
+    __builtin_memset(m_counts_tx_size, 0, TX_SIZES * TX_SIZE_CONTEXTS * TX_SIZES);
+    __builtin_memset(m_counts_is_inter, 0, IS_INTER_CONTEXTS * 2);
+    __builtin_memset(m_counts_comp_mode, 0, COMP_MODE_CONTEXTS * 2);
+    __builtin_memset(m_counts_single_ref, 0, REF_CONTEXTS * 2 * 2);
+    __builtin_memset(m_counts_comp_ref, 0, REF_CONTEXTS * 2);
+    __builtin_memset(m_counts_skip, 0, SKIP_CONTEXTS * 2);
+    __builtin_memset(m_counts_mv_joint, 0, MV_JOINTS);
+    __builtin_memset(m_counts_mv_sign, 0, 2 * 2);
+    __builtin_memset(m_counts_mv_class, 0, 2 * MV_CLASSES);
+    __builtin_memset(m_counts_mv_class0_bit, 0, 2 * CLASS0_SIZE);
+    __builtin_memset(m_counts_mv_class0_fr, 0, 2 * CLASS0_SIZE * MV_FR_SIZE);
+    __builtin_memset(m_counts_mv_class0_hp, 0, 2 * 2);
+    __builtin_memset(m_counts_mv_bits, 0, 2 * MV_OFFSET_BITS * 2);
+    __builtin_memset(m_counts_mv_fr, 0, 2 * MV_FR_SIZE);
+    __builtin_memset(m_counts_mv_hp, 0, 2 * 2);
+    __builtin_memset(m_counts_token, 0, TX_SIZES * BLOCK_TYPES * REF_TYPES * COEF_BANDS * PREV_COEF_CONTEXTS * UNCONSTRAINED_NODES);
+    __builtin_memset(m_counts_more_coefs, 0, TX_SIZES * BLOCK_TYPES * REF_TYPES * COEF_BANDS * PREV_COEF_CONTEXTS * 2);
 }
 
 }
diff --git a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
index 78ec1c2865..336e3897f8 100644
--- a/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
+++ b/Userland/Libraries/LibVideo/VP9/SyntaxElementCounter.h
@@ -46,29 +46,28 @@ class SyntaxElementCounter final {
 public:
     void clear_counts();
 
-private:
-    u8 counts_intra_mode[BLOCK_SIZE_GROUPS][INTRA_MODES];
-    u8 counts_uv_mode[INTRA_MODES][INTRA_MODES];
-    u8 counts_partition[PARTITION_CONTEXTS][PARTITION_TYPES];
-    u8 counts_interp_filter[INTERP_FILTER_CONTEXTS][SWITCHABLE_FILTERS];
-    u8 counts_inter_mode[INTER_MODE_CONTEXTS][INTER_MODES];
-    u8 counts_tx_size[TX_SIZES][TX_SIZE_CONTEXTS][TX_SIZES];
-    u8 counts_is_inter[IS_INTER_CONTEXTS][2];
-    u8 counts_comp_mode[COMP_MODE_CONTEXTS][2];
-    u8 counts_single_ref[REF_CONTEXTS][2][2];
-    u8 counts_comp_ref[REF_CONTEXTS][2];
-    u8 counts_skip[SKIP_CONTEXTS][2];
-    u8 counts_mv_joint[MV_JOINTS];
-    u8 counts_mv_sign[2][2];
-    u8 counts_mv_class[2][MV_CLASSES];
-    u8 counts_mv_class0_bit[2][CLASS0_SIZE];
-    u8 counts_mv_class0_fr[2][CLASS0_SIZE][MV_FR_SIZE];
-    u8 counts_mv_class0_hp[2][2];
-    u8 counts_mv_bits[2][MV_OFFSET_BITS][2];
-    u8 counts_mv_fr[2][MV_FR_SIZE];
-    u8 counts_mv_hp[2][2];
-    u8 counts_token[TX_SIZES][BLOCK_TYPES][REF_TYPES][COEF_BANDS][PREV_COEF_CONTEXTS][UNCONSTRAINED_NODES];
-    u8 counts_more_coefs[TX_SIZES][BLOCK_TYPES][REF_TYPES][COEF_BANDS][PREV_COEF_CONTEXTS][2];
+    u8 m_counts_intra_mode[BLOCK_SIZE_GROUPS][INTRA_MODES];
+    u8 m_counts_uv_mode[INTRA_MODES][INTRA_MODES];
+    u8 m_counts_partition[PARTITION_CONTEXTS][PARTITION_TYPES];
+    u8 m_counts_interp_filter[INTERP_FILTER_CONTEXTS][SWITCHABLE_FILTERS];
+    u8 m_counts_inter_mode[INTER_MODE_CONTEXTS][INTER_MODES];
+    u8 m_counts_tx_size[TX_SIZES][TX_SIZE_CONTEXTS][TX_SIZES];
+    u8 m_counts_is_inter[IS_INTER_CONTEXTS][2];
+    u8 m_counts_comp_mode[COMP_MODE_CONTEXTS][2];
+    u8 m_counts_single_ref[REF_CONTEXTS][2][2];
+    u8 m_counts_comp_ref[REF_CONTEXTS][2];
+    u8 m_counts_skip[SKIP_CONTEXTS][2];
+    u8 m_counts_mv_joint[MV_JOINTS];
+    u8 m_counts_mv_sign[2][2];
+    u8 m_counts_mv_class[2][MV_CLASSES];
+    u8 m_counts_mv_class0_bit[2][CLASS0_SIZE];
+    u8 m_counts_mv_class0_fr[2][CLASS0_SIZE][MV_FR_SIZE];
+    u8 m_counts_mv_class0_hp[2][2];
+    u8 m_counts_mv_bits[2][MV_OFFSET_BITS][2];
+    u8 m_counts_mv_fr[2][MV_FR_SIZE];
+    u8 m_counts_mv_hp[2][2];
+    u8 m_counts_token[TX_SIZES][BLOCK_TYPES][REF_TYPES][COEF_BANDS][PREV_COEF_CONTEXTS][UNCONSTRAINED_NODES];
+    u8 m_counts_more_coefs[TX_SIZES][BLOCK_TYPES][REF_TYPES][COEF_BANDS][PREV_COEF_CONTEXTS][2];
 };
 
 }
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index ac06e66b2f..28e9def16c 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -12,14 +12,19 @@ namespace Video::VP9 {
 int TreeParser::parse_tree(SyntaxElementType type)
 {
     auto tree_selection = select_tree(type);
-    if (tree_selection.is_single_value())
-        return tree_selection.get_single_value();
-    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))];
-    } while (n > 0);
-    return -n;
+    int value;
+    if (tree_selection.is_single_value()) {
+        value = tree_selection.get_single_value();
+    } else {
+        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))];
+        } while (n > 0);
+        value = -n;
+    }
+    count_syntax_element(type, value);
+    return value;
 }
 
 /*
@@ -171,10 +176,70 @@ u8 TreeParser::calculate_partition_probability(u8 node)
     }
     above = (above & (1 << block_offset)) > 0;
     left = (left & (1 << block_offset)) > 0;
-    auto ctx = bsl * 4 + left * 2 + above;
+    m_ctx = bsl * 4 + left * 2 + above;
     if (m_frame_is_intra)
-        return m_probability_tables.kf_partition_probs()[ctx][node2];
-    return m_probability_tables.partition_probs()[ctx][node2];
+        return m_probability_tables.kf_partition_probs()[m_ctx][node2];
+    return m_probability_tables.partition_probs()[m_ctx][node2];
+}
+
+void TreeParser::count_syntax_element(SyntaxElementType type, int value)
+{
+    switch (type) {
+    case SyntaxElementType::Partition:
+        m_syntax_element_counter->m_counts_partition[m_ctx][value]++;
+        break;
+    case SyntaxElementType::IntraMode:
+        break;
+    case SyntaxElementType::SubIntraMode:
+        break;
+    case SyntaxElementType::UVMode:
+        break;
+    case SyntaxElementType::Skip:
+        break;
+    case SyntaxElementType::IsInter:
+        break;
+    case SyntaxElementType::CompMode:
+        break;
+    case SyntaxElementType::CompRef:
+        break;
+    case SyntaxElementType::SingleRefP1:
+        break;
+    case SyntaxElementType::SingleRefP2:
+        break;
+    case SyntaxElementType::MVSign:
+        break;
+    case SyntaxElementType::MVClass0Bit:
+        break;
+    case SyntaxElementType::MVBit:
+        break;
+    case SyntaxElementType::TXSize:
+        break;
+    case SyntaxElementType::InterMode:
+        break;
+    case SyntaxElementType::InterpFilter:
+        break;
+    case SyntaxElementType::MVJoint:
+        break;
+    case SyntaxElementType::MVClass:
+        break;
+    case SyntaxElementType::MVClass0FR:
+        break;
+    case SyntaxElementType::MVClass0HP:
+        break;
+    case SyntaxElementType::MVFR:
+        break;
+    case SyntaxElementType::MVHP:
+        break;
+    case SyntaxElementType::Token:
+        break;
+    case SyntaxElementType::MoreCoefs:
+        break;
+    case SyntaxElementType::DefaultIntraMode:
+    case SyntaxElementType::DefaultUVMode:
+    case SyntaxElementType::SegmentID:
+    case SyntaxElementType::SegIDPredicted:
+        break;
+    }
 }
 
 TreeParser::TreeSelection::TreeSelection(const int* values)
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 2fb09841ee..d444b136a3 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -42,6 +42,7 @@ public:
     int parse_tree(SyntaxElementType type);
     TreeSelection select_tree(SyntaxElementType type);
     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; }
@@ -55,12 +56,17 @@ public:
     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; }
 
 private:
     u8 calculate_partition_probability(u8 node);
 
     ProbabilityTables& m_probability_tables;
     BitStream* m_bit_stream { nullptr };
+    SyntaxElementCounter* m_syntax_element_counter { nullptr };
+
+    // 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 };