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