From c1e5feca13e1c9f504fbfef4b58ab6f229dcef6e Mon Sep 17 00:00:00 2001
From: Martin Benjak <m_benjak@apple.com>
Date: Thu, 6 Aug 2026 14:05:42 -0700
Subject: [PATCH] Remove encoder-side limitation for restricted switch frames

The function prune_ref_by_selective_ref_frame() returned 1 for all blocks that reference restricted references. This caused WARPMV to be the only prediction mode to be considered for restricted references by the encoder, severely hindering the compression performance of restricted switch frames.

Removed this limitation and set display_order_hint to REF_RESTRICTED_DOH to ensure that all now accessible prediction modes are correctly encoded. Added is_restricted checks to av2_get_refresh_frame_flags() because the display_order_hint of restricted references is always >= the current display_order_hint.
---
 av2/encoder/encode_strategy.c |  3 ++-
 av2/encoder/encoder.c         |  1 +
 av2/encoder/rdopt.h           | 14 +++++++-------
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/av2/encoder/encode_strategy.c b/av2/encoder/encode_strategy.c
index 6e6eb63d39..73510db30d 100644
--- a/av2/encoder/encode_strategy.c
+++ b/av2/encoder/encode_strategy.c
@@ -794,6 +794,7 @@ int av2_get_refresh_frame_flags(
           // been output yet and whose DOH is at least the current
           // frame's DOH. (DOH requirement)
           if (cm->ref_frame_map[i] != NULL &&
+              !cm->ref_frame_map[i]->is_restricted &&
               cm->ref_frame_map[i]->implicit_output_picture &&
               !cm->ref_frame_map[i]->frame_output_done &&
               (int)cm->ref_frame_map[i]->display_order_hint >= cur_disp_order) {
@@ -847,7 +848,7 @@ int av2_get_refresh_frame_flags(
   // at least the current frame's DOH. (DOH requirement)
   for (int i = 0; i < cpi->common.seq_params.ref_frames; i++) {
     const RefCntBuffer *const buf = cpi->common.ref_frame_map[i];
-    if (buf != NULL && buf->implicit_output_picture &&
+    if (buf != NULL && !buf->is_restricted && buf->implicit_output_picture &&
         !buf->frame_output_done &&
         (int)buf->display_order_hint >= cur_disp_order) {
       olk_flags_to_keep |= (1 << i);
diff --git a/av2/encoder/encoder.c b/av2/encoder/encoder.c
index aa9c53fb8e..689b9e36b2 100644
--- a/av2/encoder/encoder.c
+++ b/av2/encoder/encoder.c
@@ -5118,6 +5118,7 @@ int av2_encode(AV2_COMP *const cpi, uint8_t *const dest,
                                                cm->ref_frame_map[i]->mlayer_id,
                                                cm->mlayer_id)) {
             cm->ref_frame_map[i]->is_restricted = true;
+            cm->ref_frame_map[i]->display_order_hint = REF_RESTRICTED_DOH;
           }
         }
       }
diff --git a/av2/encoder/rdopt.h b/av2/encoder/rdopt.h
index 4607390495..c1d4124a77 100644
--- a/av2/encoder/rdopt.h
+++ b/av2/encoder/rdopt.h
@@ -215,13 +215,6 @@ static INLINE int prune_ref_by_selective_ref_frame(
   assert(ref_frame[0] != NONE_FRAME);
   if (ref_frame[0] == INTRA_FRAME) return 0;
 
-  RefCntBuffer *ref0_buf = get_ref_frame_buf(cm, ref_frame[0]);
-  if (ref0_buf)
-    if (ref0_buf->is_restricted) return 1;
-  RefCntBuffer *ref1_buf = get_ref_frame_buf(cm, ref_frame[1]);
-  if (ref1_buf)
-    if (ref1_buf->is_restricted) return 1;
-
   const int comp_pred = is_inter_ref_frame(ref_frame[1]);
 
   if (comp_pred && ref_frame[0] >= RANKED_REF0_TO_PRUNE) return 1;
@@ -259,6 +252,13 @@ static INLINE int prune_ref_by_selective_ref_frame(
     }
   }
 
+  RefCntBuffer *ref0_buf = get_ref_frame_buf(cm, ref_frame[0]);
+  if (ref0_buf)
+    if (ref0_buf->is_restricted) return 0;
+  RefCntBuffer *ref1_buf = get_ref_frame_buf(cm, ref_frame[1]);
+  if (ref1_buf)
+    if (ref1_buf->is_restricted) return 0;
+
   int dir_refrank0[2] = { -1, -1 };
   int dir_refrank1[2] = { -1, -1 };
   int d0 = get_dir_rank(cm, ref_frame[0], dir_refrank0);