From 2e055519b53beed326ce977d1452c939319a86ef Mon Sep 17 00:00:00 2001
From: Emily Samp <emily.samp@shopify.com>
Date: Fri, 23 Aug 2024 16:06:25 -0500
Subject: [PATCH] Implement translation for `case` and `when`

---
 parser/prism/Translator.cc                | 41 +++++++++-
 test/prism_regression/case.parse-tree.exp | 98 +++++++++++++++++++++++
 test/prism_regression/case.rb             | 19 +++++
 3 files changed, 156 insertions(+), 2 deletions(-)
 create mode 100644 test/prism_regression/case.parse-tree.exp
 create mode 100644 test/prism_regression/case.rb

diff --git a/parser/prism/Translator.cc b/parser/prism/Translator.cc
index 2a3796fe11f..b8b0441b8c0 100644
--- a/parser/prism/Translator.cc
+++ b/parser/prism/Translator.cc
@@ -142,6 +142,26 @@ std::unique_ptr<parser::Node> Translator::translate(pm_node_t *node) {
                 return sendNode;
             }
         }
+        case PM_CASE_NODE: {
+            auto caseNode = reinterpret_cast<pm_case_node *>(node);
+            pm_location_t *loc = &caseNode->base.location;
+
+            auto predicate = translate(caseNode->predicate);
+
+            auto prismConditions = absl::MakeSpan(caseNode->conditions.nodes, caseNode->conditions.size);
+
+            parser::NodeVec sorbetConditions{};
+            sorbetConditions.reserve(prismConditions.size());
+
+            for (auto &prismCondition : prismConditions) {
+                sorbetConditions.emplace_back(translate(prismCondition));
+            }
+
+            auto consequent = translate(reinterpret_cast<pm_node_t *>(caseNode->consequent));
+
+            return make_unique<Case>(parser.translateLocation(loc), std::move(predicate), std::move(sorbetConditions),
+                                     std::move(consequent));
+        }
         case PM_CLASS_NODE: { // Class declarations, not including singleton class declarations (`class <<`)
             auto classNode = reinterpret_cast<pm_class_node *>(node);
             pm_location_t *loc = &classNode->base.location;
@@ -581,6 +601,25 @@ std::unique_ptr<parser::Node> Translator::translate(pm_node_t *node) {
 
             return make_unique<parser::Until>(parser.translateLocation(loc), std::move(predicate), std::move(body));
         }
+        case PM_WHEN_NODE: {
+            auto whenNode = reinterpret_cast<pm_when_node *>(node);
+            auto *loc = &whenNode->base.location;
+
+            auto prismConditions = absl::MakeSpan(whenNode->conditions.nodes, whenNode->conditions.size);
+
+            parser::NodeVec sorbetConditions{};
+            sorbetConditions.reserve(prismConditions.size());
+
+            for (auto &prismCondition : prismConditions) {
+                sorbetConditions.emplace_back(translate(prismCondition));
+            }
+
+            auto inlineIfSingle = true;
+            auto statements = translateStatements(whenNode->statements, inlineIfSingle);
+
+            return make_unique<parser::When>(parser.translateLocation(loc), std::move(sorbetConditions),
+                                             std::move(statements));
+        }
         case PM_YIELD_NODE: {
             auto yieldNode = reinterpret_cast<pm_yield_node *>(node);
             pm_location_t *loc = &yieldNode->base.location;
@@ -604,7 +643,6 @@ std::unique_ptr<parser::Node> Translator::translate(pm_node_t *node) {
         case PM_CALL_TARGET_NODE:
         case PM_CAPTURE_PATTERN_NODE:
         case PM_CASE_MATCH_NODE:
-        case PM_CASE_NODE:
         case PM_CLASS_VARIABLE_AND_WRITE_NODE:
         case PM_CLASS_VARIABLE_OPERATOR_WRITE_NODE:
         case PM_CLASS_VARIABLE_OR_WRITE_NODE:
@@ -684,7 +722,6 @@ std::unique_ptr<parser::Node> Translator::translate(pm_node_t *node) {
         case PM_SOURCE_FILE_NODE:
         case PM_SOURCE_LINE_NODE:
         case PM_UNDEF_NODE:
-        case PM_WHEN_NODE:
         case PM_WHILE_NODE:
         case PM_X_STRING_NODE:
         case PM_SCOPE_NODE:
diff --git a/test/prism_regression/case.parse-tree.exp b/test/prism_regression/case.parse-tree.exp
new file mode 100644
index 00000000000..90aa4c6782c
--- /dev/null
+++ b/test/prism_regression/case.parse-tree.exp
@@ -0,0 +1,98 @@
+Begin {
+  stmts = [
+    DefMethod {
+      name = <U foo>
+      args = NULL
+      body = NULL
+    }
+    Case {
+      condition = Send {
+        receiver = NULL
+        method = <U foo>
+        args = [
+        ]
+      }
+      whens = [
+        When {
+          patterns = [
+            Integer {
+              val = "1"
+            }
+          ]
+          body = Send {
+            receiver = NULL
+            method = <U puts>
+            args = [
+              String {
+                val = <U one!>
+              }
+            ]
+          }
+        }
+        When {
+          patterns = [
+            Integer {
+              val = "2"
+            }
+            Integer {
+              val = "3"
+            }
+          ]
+          body = Send {
+            receiver = NULL
+            method = <U puts>
+            args = [
+              String {
+                val = <U two or three!>
+              }
+            ]
+          }
+        }
+      ]
+      else_ = Send {
+        receiver = NULL
+        method = <U puts>
+        args = [
+          String {
+            val = <U Who knows!>
+          }
+        ]
+      }
+    }
+    Case {
+      condition = Send {
+        receiver = NULL
+        method = <U foo>
+        args = [
+        ]
+      }
+      whens = [
+        When {
+          patterns = [
+            Const {
+              scope = NULL
+              name = <C <U Integer>>
+            }
+          ]
+          body = Begin {
+            stmts = [
+              Integer {
+                val = "4"
+              }
+              Send {
+                receiver = NULL
+                method = <U puts>
+                args = [
+                  String {
+                    val = <U surprise, multi-line!>
+                  }
+                ]
+              }
+            ]
+          }
+        }
+      ]
+      else_ = NULL
+    }
+  ]
+}
diff --git a/test/prism_regression/case.rb b/test/prism_regression/case.rb
new file mode 100644
index 00000000000..3092359d35b
--- /dev/null
+++ b/test/prism_regression/case.rb
@@ -0,0 +1,19 @@
+# typed: true
+
+def foo; end
+
+case foo
+when 1
+  puts "one!"
+when 2, 3
+  puts "two or three!"
+else
+  puts "Who knows!"
+end
+
+# no else
+case foo
+when Integer
+  4
+  puts "surprise, multi-line!"
+end