diff --git a/lib/prism/translation/parser/compiler.rb b/lib/prism/translation/parser/compiler.rb
index 1369fe6f8..7643c4af0 100644
--- a/lib/prism/translation/parser/compiler.rb
+++ b/lib/prism/translation/parser/compiler.rb
@@ -1607,11 +1607,7 @@ def visit_when_node(node)
           builder.when(
             token(node.keyword_loc),
             visit_all(node.conditions),
-            if node.then_keyword_loc
-              token(node.then_keyword_loc)
-            else
-              srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset || (node.conditions.last.location.end_offset + 1), [";"])
-            end,
+            srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset || node.statements&.then_keyword_loc, [";", "then"]),
             visit(node.statements)
           )
         end