LibVideo/VP9: Implement sections 6.1.2 and 8.4.1-8.4.4

These section implement the behavior to refresh the probability
tables after parsing a frame.

8 files changed, 188 insertions, 9 deletions

diff --git a/Userland/Libraries/LibVideo/CMakeLists.txt b/Userland/Libraries/LibVideo/CMakeLists.txt
index d7a3f4a21a..ecc1b0483b 100644
--- a/Userland/Libraries/LibVideo/CMakeLists.txt
+++ b/Userland/Libraries/LibVideo/CMakeLists.txt
@@ -10,7 +10,7 @@ set(SOURCES
     VP9/Symbols.h
     VP9/SyntaxElementCounter.cpp
     VP9/TreeParser.cpp
-    VP9/Utilities.h
+    VP9/Utilities.cpp
 )
 
 serenity_lib(LibVideo video)
diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.cpp b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
index 68601efcef..725d9ac29b 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.cpp
@@ -30,6 +30,126 @@ void Decoder::dump_frame_info()
     m_parser->dump_info();
 }
 
+u8 Decoder::merge_prob(u8 pre_prob, u8 count_0, u8 count_1, u8 count_sat, u8 max_update_factor)
+{
+    auto total_decode_count = count_0 + count_1;
+    auto prob = (total_decode_count == 0) ? 128 : clip_3(1, 255, (count_0 * 256 + (total_decode_count >> 1)) / total_decode_count);
+    auto count = min(total_decode_count, count_sat);
+    auto factor = (max_update_factor * count) / count_sat;
+    return round_2(pre_prob * (256 - factor) + (prob * factor), 8);
+}
+
+u8 Decoder::merge_probs(const int* tree, int index, u8* probs, u8* counts, u8 count_sat, u8 max_update_factor)
+{
+    auto s = tree[index];
+    auto left_count = (s <= 0) ? counts[-s] : merge_probs(tree, s, probs, counts, count_sat, max_update_factor);
+    auto r = tree[index + 1];
+    auto right_count = (r <= 0) ? counts[-r] : merge_probs(tree, r, probs, counts, count_sat, max_update_factor);
+    probs[index >> 1] = merge_prob(probs[index >> 1], left_count, right_count, count_sat, max_update_factor);
+    return left_count + right_count;
+}
+
+bool Decoder::adapt_coef_probs()
+{
+    u8 update_factor;
+    if (m_parser->m_frame_is_intra || m_parser->m_last_frame_type != KeyFrame)
+        update_factor = 112;
+    else
+        update_factor = 128;
+
+    for (size_t t = 0; t < 4; t++) {
+        for (size_t i = 0; i < 2; i++) {
+            for (size_t j = 0; j < 2; j++) {
+                for (size_t k = 0; k < 6; k++) {
+                    size_t max_l = (k == 0) ? 3 : 6;
+                    for (size_t l = 0; l < max_l; l++) {
+                        auto& coef_probs = m_parser->m_probability_tables->coef_probs()[t][i][j][k][l];
+                        merge_probs(small_token_tree, 2, coef_probs,
+                            m_parser->m_syntax_element_counter->m_counts_token[t][i][j][k][l],
+                            24, update_factor);
+                        merge_probs(binary_tree, 0, coef_probs,
+                            m_parser->m_syntax_element_counter->m_counts_more_coefs[t][i][j][k][l],
+                            24, update_factor);
+                    }
+                }
+            }
+        }
+    }
+
+    return true;
+}
+
+#define ADAPT_PROB_TABLE(name, size)                                     \
+    do {                                                                 \
+        for (size_t i = 0; i < (size); i++) {                            \
+            auto table = probs.name##_prob();                            \
+            table[i] = adapt_prob(table[i], counter.m_counts_##name[i]); \
+        }                                                                \
+    } while (0)
+
+#define ADAPT_TREE(tree_name, prob_name, count_name, size)                                                 \
+    do {                                                                                                   \
+        for (size_t i = 0; i < (size); i++) {                                                              \
+            adapt_probs(tree_name##_tree, probs.prob_name##_probs()[i], counter.m_counts_##count_name[i]); \
+        }                                                                                                  \
+    } while (0)
+
+bool Decoder::adapt_non_coef_probs()
+{
+    auto& probs = *m_parser->m_probability_tables;
+    auto& counter = *m_parser->m_syntax_element_counter;
+    ADAPT_PROB_TABLE(is_inter, IS_INTER_CONTEXTS);
+    ADAPT_PROB_TABLE(comp_mode, COMP_MODE_CONTEXTS);
+    ADAPT_PROB_TABLE(comp_ref, REF_CONTEXTS);
+    for (size_t i = 0; i < REF_CONTEXTS; i++) {
+        for (size_t j = 0; j < 2; j++)
+            probs.single_ref_prob()[i][j] = adapt_prob(probs.single_ref_prob()[i][j], counter.m_counts_single_ref[i][j]);
+    }
+    ADAPT_TREE(inter_mode, inter_mode, inter_mode, INTER_MODE_CONTEXTS);
+    ADAPT_TREE(intra_mode, y_mode, intra_mode, INTER_MODE_CONTEXTS);
+    ADAPT_TREE(intra_mode, uv_mode, uv_mode, INTER_MODE_CONTEXTS);
+    ADAPT_TREE(partition, partition, partition, INTER_MODE_CONTEXTS);
+    ADAPT_PROB_TABLE(skip, SKIP_CONTEXTS);
+    if (m_parser->m_interpolation_filter == Switchable) {
+        ADAPT_TREE(interp_filter, interp_filter, interp_filter, INTERP_FILTER_CONTEXTS);
+    }
+    if (m_parser->m_tx_mode == TXModeSelect) {
+        for (size_t i = 0; i < TX_SIZE_CONTEXTS; i++) {
+            auto& tx_probs = probs.tx_probs();
+            auto& tx_counts = counter.m_counts_tx_size;
+            adapt_probs(tx_size_8_tree, tx_probs[TX_8x8][i], tx_counts[TX_8x8][i]);
+            adapt_probs(tx_size_16_tree, tx_probs[TX_16x16][i], tx_counts[TX_16x16][i]);
+            adapt_probs(tx_size_32_tree, tx_probs[TX_32x32][i], tx_counts[TX_32x32][i]);
+        }
+    }
+    adapt_probs(mv_joint_tree, probs.mv_joint_probs(), counter.m_counts_mv_joint);
+    for (size_t i = 0; i < 2; i++) {
+        probs.mv_sign_prob()[i] = adapt_prob(probs.mv_sign_prob()[i], counter.m_counts_mv_sign[i]);
+        adapt_probs(mv_class_tree, probs.mv_class_probs()[i], counter.m_counts_mv_class[i]);
+        probs.mv_class0_bit_prob()[i] = adapt_prob(probs.mv_class0_bit_prob()[i], counter.m_counts_mv_class0_bit[i]);
+        for (size_t j = 0; j < MV_OFFSET_BITS; j++)
+            probs.mv_bits_prob()[i][j] = adapt_prob(probs.mv_bits_prob()[i][j], counter.m_counts_mv_bits[i][j]);
+        for (size_t j = 0; j < CLASS0_SIZE; j++)
+            adapt_probs(mv_fr_tree, probs.mv_class0_fr_probs()[i][j], counter.m_counts_mv_class0_fr[i][j]);
+        adapt_probs(mv_fr_tree, probs.mv_fr_probs()[i], counter.m_counts_mv_fr[i]);
+        if (m_parser->m_allow_high_precision_mv) {
+            probs.mv_class0_hp_prob()[i] = adapt_prob(probs.mv_class0_hp_prob()[i], counter.m_counts_mv_class0_hp[i]);
+            probs.mv_hp_prob()[i] = adapt_prob(probs.mv_hp_prob()[i], counter.m_counts_mv_hp[i]);
+        }
+    }
+    return true;
+}
+
+void Decoder::adapt_probs(int const* tree, u8* probs, u8* counts)
+{
+    merge_probs(tree, 0, probs, counts, COUNT_SAT, MAX_UPDATE_FACTOR);
+}
+
+u8 Decoder::adapt_prob(u8 prob, u8 counts[2])
+{
+    return merge_prob(prob, counts[0], counts[1], COUNT_SAT, MAX_UPDATE_FACTOR);
+}
+
 bool Decoder::predict_intra(size_t, u32, u32, bool, bool, bool, TXSize, u32)
 {
     // TODO: Implement
diff --git a/Userland/Libraries/LibVideo/VP9/Decoder.h b/Userland/Libraries/LibVideo/VP9/Decoder.h
index 1883af1964..0732483896 100644
--- a/Userland/Libraries/LibVideo/VP9/Decoder.h
+++ b/Userland/Libraries/LibVideo/VP9/Decoder.h
@@ -20,6 +20,14 @@ public:
     void dump_frame_info();
 
 private:
+    /* (8.4) Probability Adaptation Process */
+    u8 merge_prob(u8 pre_prob, u8 count_0, u8 count_1, u8 count_sat, u8 max_update_factor);
+    u8 merge_probs(int const* tree, int index, u8* probs, u8* counts, u8 count_sat, u8 max_update_factor);
+    bool adapt_coef_probs();
+    bool adapt_non_coef_probs();
+    void adapt_probs(int const* tree, u8* probs, u8* counts);
+    u8 adapt_prob(u8 prob, u8 counts[2]);
+
     /* (8.5) Prediction Processes */
     bool predict_intra(size_t plane, u32 x, u32 y, bool have_left, bool have_above, bool not_on_right, TXSize tx_size, u32 block_index);
     bool predict_inter(size_t plane, u32 x, u32 y, u32 w, u32 h, u32 block_index);
diff --git a/Userland/Libraries/LibVideo/VP9/LookupTables.h b/Userland/Libraries/LibVideo/VP9/LookupTables.h
index df0edeb185..228c0e4033 100644
--- a/Userland/Libraries/LibVideo/VP9/LookupTables.h
+++ b/Userland/Libraries/LibVideo/VP9/LookupTables.h
@@ -216,4 +216,10 @@ static constexpr BlockSubsize ss_size_lookup[BLOCK_SIZES][2][2] = {
     { { Block_64x64, Block_64x32 }, { Block_32x64, Block_32x32 } },
 };
 
+static constexpr int small_token_tree[6] = {
+    0, 0, // Unused
+    -ZeroToken, 4,
+    -OneToken, -TwoToken
+};
+
 }
diff --git a/Userland/Libraries/LibVideo/VP9/Parser.cpp b/Userland/Libraries/LibVideo/VP9/Parser.cpp
index 4a291a6630..6c5b1c07d5 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Parser.cpp
@@ -45,11 +45,34 @@ bool Parser::parse_frame(ByteBuffer const& frame_data)
     SAFE_CALL(m_bit_stream->exit_bool());
 
     SAFE_CALL(decode_tiles());
+    SAFE_CALL(refresh_probs());
 
     dbgln("Finished reading frame!");
     return true;
 }
 
+bool Parser::trailing_bits()
+{
+    while (m_bit_stream->get_position() & 7u)
+        RESERVED_ZERO;
+    return true;
+}
+
+bool Parser::refresh_probs()
+{
+    if (!m_error_resilient_mode && !m_frame_parallel_decoding_mode) {
+        m_probability_tables->load_probs(m_frame_context_idx);
+        SAFE_CALL(m_decoder.adapt_coef_probs());
+        if (!m_frame_is_intra) {
+            m_probability_tables->load_probs2(m_frame_context_idx);
+            SAFE_CALL(m_decoder.adapt_non_coef_probs());
+        }
+    }
+    if (m_refresh_frame_context)
+        m_probability_tables->save_probs(m_frame_context_idx);
+    return true;
+}
+
 /* (6.2) */
 bool Parser::uncompressed_header()
 {
@@ -397,13 +420,6 @@ bool Parser::setup_past_independence()
     return true;
 }
 
-bool Parser::trailing_bits()
-{
-    while (m_bit_stream->get_position() & 7u)
-        RESERVED_ZERO;
-    return true;
-}
-
 bool Parser::compressed_header()
 {
     SAFE_CALL(read_tx_mode());
diff --git a/Userland/Libraries/LibVideo/VP9/Parser.h b/Userland/Libraries/LibVideo/VP9/Parser.h
index f9b256dc45..cc0df8139b 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.h
+++ b/Userland/Libraries/LibVideo/VP9/Parser.h
@@ -46,6 +46,10 @@ private:
     void clear_context(Vector<u8>& context, size_t size);
     void clear_context(Vector<Vector<u8>>& context, size_t outer_size, size_t inner_size);
 
+    /* (6.1) Frame Syntax */
+    bool trailing_bits();
+    bool refresh_probs();
+
     /* (6.2) Uncompressed Header Syntax */
     bool uncompressed_header();
     bool frame_sync_code();
@@ -64,7 +68,6 @@ private:
     u16 calc_min_log2_tile_cols();
     u16 calc_max_log2_tile_cols();
     bool setup_past_independence();
-    bool trailing_bits();
 
     /* (6.3) Compressed Header Syntax */
     bool compressed_header();
diff --git a/Userland/Libraries/LibVideo/VP9/Utilities.cpp b/Userland/Libraries/LibVideo/VP9/Utilities.cpp
new file mode 100644
index 0000000000..70336efc6d
--- /dev/null
+++ b/Userland/Libraries/LibVideo/VP9/Utilities.cpp
@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) 2021, Hunter Salyer <thefalsehonesty@gmail.com>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#include "Utilities.h"
+
+namespace Video::VP9 {
+
+u8 clip_3(u8 x, u8 y, u8 z)
+{
+    return clamp(z, x, y);
+}
+
+u8 round_2(u8 x, u8 n)
+{
+    return (x + (1 << (n - 1))) >> n;
+}
+
+}
diff --git a/Userland/Libraries/LibVideo/VP9/Utilities.h b/Userland/Libraries/LibVideo/VP9/Utilities.h
index 9222cddb95..26ac862c8f 100644
--- a/Userland/Libraries/LibVideo/VP9/Utilities.h
+++ b/Userland/Libraries/LibVideo/VP9/Utilities.h
@@ -6,6 +6,8 @@
 
 #pragma once
 
+#include <AK/Types.h>
+
 namespace Video::VP9 {
 
 #define SAFE_CALL(call)             \
@@ -16,4 +18,7 @@ namespace Video::VP9 {
         }                           \
     } while (0)
 
+u8 clip_3(u8 x, u8 y, u8 z);
+u8 round_2(u8 x, u8 n);
+
 }