3 files changed, 15 insertions, 9 deletions
diff --git a/Userland/Libraries/LibVideo/VP9/Parser.cpp b/Userland/Libraries/LibVideo/VP9/Parser.cpp
index cb33f79084..8903f5023e 100644
--- a/Userland/Libraries/LibVideo/VP9/Parser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/Parser.cpp
@@ -1057,7 +1057,7 @@ DecoderErrorOr<void> Parser::intra_frame_mode_info()
         }
         m_y_mode = m_default_intra_mode;
     }
-    m_uv_mode = TRY_READ(m_tree_parser->parse_tree<PredictionMode>(SyntaxElementType::DefaultUVMode));
+    m_uv_mode = TRY_READ(TreeParser::parse_default_uv_mode(*m_bit_stream, *m_probability_tables, m_y_mode));
     return {};
 }
 
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
index 2cb90bc91f..bebf97d344 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.cpp
@@ -134,6 +134,19 @@ ErrorOr<PredictionMode> TreeParser::parse_default_intra_mode(BitStream& bit_stre
     return value;
 }
 
+ErrorOr<PredictionMode> TreeParser::parse_default_uv_mode(BitStream& bit_stream, ProbabilityTables const& probability_table, PredictionMode y_mode)
+{
+    // Tree
+    TreeParser::TreeSelection tree = { intra_mode_tree };
+
+    // Probabilities
+    u8 const* probabilities = probability_table.kf_uv_mode_prob()[to_underlying(y_mode)];
+
+    auto value = TRY(parse_tree_new<PredictionMode>(bit_stream, tree, [&](u8 node) { return probabilities[node]; }));
+    // Default UV mode is not counted.
+    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
  */
@@ -195,8 +208,6 @@ TreeParser::TreeSelection TreeParser::select_tree(SyntaxElementType type)
 u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
 {
     switch (type) {
-    case SyntaxElementType::DefaultUVMode:
-        return calculate_default_uv_mode_probability(node);
     case SyntaxElementType::IntraMode:
         return calculate_intra_mode_probability(node);
     case SyntaxElementType::SubIntraMode:
@@ -268,11 +279,6 @@ u8 TreeParser::select_tree_probability(SyntaxElementType type, u8 node)
 #define ABOVE_SINGLE m_decoder.m_above_single
 #define LEFT_SINGLE m_decoder.m_left_single
 
-u8 TreeParser::calculate_default_uv_mode_probability(u8 node)
-{
-    return m_decoder.m_probability_tables->kf_uv_mode_prob()[to_underlying(m_decoder.m_y_mode)][node];
-}
-
 u8 TreeParser::calculate_intra_mode_probability(u8 node)
 {
     m_ctx = size_group_lookup[m_decoder.m_mi_size];
diff --git a/Userland/Libraries/LibVideo/VP9/TreeParser.h b/Userland/Libraries/LibVideo/VP9/TreeParser.h
index 7f031c0382..411883ac0e 100644
--- a/Userland/Libraries/LibVideo/VP9/TreeParser.h
+++ b/Userland/Libraries/LibVideo/VP9/TreeParser.h
@@ -64,6 +64,7 @@ public:
 
     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);
 
     void set_default_intra_mode_variables(u8 idx, u8 idy)
     {
@@ -97,7 +98,6 @@ public:
     }
 
 private:
-    u8 calculate_default_uv_mode_probability(u8 node);
     u8 calculate_intra_mode_probability(u8 node);
     u8 calculate_sub_intra_mode_probability(u8 node);
     u8 calculate_uv_mode_probability(u8 node);