From fb59a0826f6b3b8f140675221b799b2a2943eb44 Mon Sep 17 00:00:00 2001
From: Soha Jin <soha@jin.sh>
Date: Mon, 17 Jun 2024 20:10:24 +0800
Subject: [PATCH] Add `getContestProblem` to Contest for getting an associated
 ContestProblem

(cherry picked from commit 56104f63024f11599719b1e4553d732cf8e6cbce)
---
 webapp/src/Entity/Contest.php     | 10 ++++++++++
 webapp/src/Twig/TwigExtension.php | 11 ++++-------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/webapp/src/Entity/Contest.php b/webapp/src/Entity/Contest.php
index 28ced7e8476..da37be007b7 100644
--- a/webapp/src/Entity/Contest.php
+++ b/webapp/src/Entity/Contest.php
@@ -862,6 +862,16 @@ public function getProblems(): Collection
         return $this->problems;
     }
 
+    public function getContestProblem(Problem $problem): ?ContestProblem
+    {
+        foreach ($this->getProblems() as $contestProblem) {
+            if ($contestProblem->getProblem() === $problem) {
+                return $contestProblem;
+            }
+        }
+        return null;
+    }
+
     public function addClarification(Clarification $clarification): Contest
     {
         $this->clarifications[] = $clarification;
diff --git a/webapp/src/Twig/TwigExtension.php b/webapp/src/Twig/TwigExtension.php
index e28461d863a..619830eef84 100644
--- a/webapp/src/Twig/TwigExtension.php
+++ b/webapp/src/Twig/TwigExtension.php
@@ -1108,14 +1108,11 @@ public function problemBadge(ContestProblem $problem): string
         );
     }
 
-    public function problemBadgeForProblemAndContest(Problem $problem, ?Contest $contest): string
+    public function problemBadgeForProblemAndContest(Problem $problem, ?Contest $contest = null): string
     {
-        foreach ($problem->getContestProblems() as $contestProblem) {
-            if ($contestProblem->getContest() === $contest) {
-                return $this->problemBadge($contestProblem);
-            }
-        }
-        return '';
+        $contest ??= $this->dj->getCurrentContest();
+        $contestProblem = $contest?->getContestProblem($problem);
+        return $contestProblem === null ? '' : $this->problemBadge($contestProblem);
     }
 
     public function printMetadata(?string $metadata): string