Merge pull request #301 from Shopify/fix-#204

Translate conditionals in `in` statements
Shopify · Oct 25, 2024 · a146baa · a146baa
2 parents 0a57681 + af24f04
commit a146baa
Show file tree

Hide file tree

Showing 7 changed files with 294 additions and 16 deletions.
diff --git a/parser/prism/Translator.cc b/parser/prism/Translator.cc
@@ -1343,12 +1343,36 @@ unique_ptr<parser::Node> Translator::patternTranslate(pm_node_t *node) {
         case PM_IN_NODE: { // An `in` pattern such as in a `case` statement, or as a standalone expression.
             auto inNode = down_cast<pm_in_node>(node);
 
-            auto sorbetPattern = patternTranslate(inNode->pattern);
-
+            auto prismPattern = inNode->pattern;
+            unique_ptr<parser::Node> sorbetPattern;
+            unique_ptr<parser::Node> sorbetGuard;
             auto inlineIfSingle = true;
             auto statements = translateStatements(inNode->statements, inlineIfSingle);
 
-            return make_unique<parser::InPattern>(location, move(sorbetPattern), nullptr, move(statements));
+            if (prismPattern != nullptr &&
+                (PM_NODE_TYPE_P(prismPattern, PM_IF_NODE) || PM_NODE_TYPE_P(prismPattern, PM_UNLESS_NODE))) {
+                pm_statements_node *conditionalStatements = nullptr;
+
+                if (PM_NODE_TYPE_P(prismPattern, PM_IF_NODE)) {
+                    auto ifNode = down_cast<pm_if_node>(prismPattern);
+                    conditionalStatements = ifNode->statements;
+                    sorbetGuard = make_unique<parser::IfGuard>(location, translate(ifNode->predicate));
+                } else { // PM_UNLESS_NODE
+                    auto unlessNode = down_cast<pm_unless_node>(prismPattern);
+                    conditionalStatements = unlessNode->statements;
+                    sorbetGuard = make_unique<parser::UnlessGuard>(location, translate(unlessNode->predicate));
+                }
+
+                ENFORCE(
+                    conditionalStatements->body.size == 1,
+                    "In pattern-matching's `in` clause, a conditional (if/unless) guard must have a single statement.");
+
+                sorbetPattern = patternTranslate(conditionalStatements->body.nodes[0]);
+            } else {
+                sorbetPattern = patternTranslate(prismPattern);
+            }
+
+            return make_unique<parser::InPattern>(location, move(sorbetPattern), move(sorbetGuard), move(statements));
         }
         case PM_LOCAL_VARIABLE_TARGET_NODE: { // A variable binding in a pattern, like the `head` in `[head, *tail]`
             auto localVarTargetNode = down_cast<pm_local_variable_target_node>(node);

diff --git a/test/prism_regression/case_match.parse-tree.exp b/test/prism_regression/case_match.parse-tree.exp
@@ -354,5 +354,199 @@ Begin {
       ]
       elseBody = NULL
     }
+    CaseMatch {
+      expr = Send {
+        receiver = NULL
+        method = <U bar>
+        args = [
+        ]
+      }
+      inBodies = [
+        InPattern {
+          pattern = MatchVar {
+            name = <U x>
+          }
+          guard = IfGuard {
+            condition = Send {
+              receiver = LVar {
+                name = <U x>
+              }
+              method = <U ==>
+              args = [
+                Integer {
+                  val = "1"
+                }
+              ]
+            }
+          }
+          body = String {
+            val = <U in with if>
+          }
+        }
+        InPattern {
+          pattern = ArrayPattern {
+            elts = [
+              MatchVar {
+                name = <U a>
+              }
+              MatchVar {
+                name = <U b>
+              }
+            ]
+          }
+          guard = IfGuard {
+            condition = Send {
+              receiver = LVar {
+                name = <U b>
+              }
+              method = <U ==>
+              args = [
+                Integer {
+                  val = "2"
+                }
+              ]
+            }
+          }
+          body = String {
+            val = <U in with 2 args and if>
+          }
+        }
+        InPattern {
+          pattern = ArrayPattern {
+            elts = [
+              MatchVar {
+                name = <U c>
+              }
+              MatchVar {
+                name = <U d>
+              }
+            ]
+          }
+          guard = NULL
+          body = Begin {
+            stmts = [
+              If {
+                condition = Send {
+                  receiver = LVar {
+                    name = <U c>
+                  }
+                  method = <U ==>
+                  args = [
+                    Integer {
+                      val = "3"
+                    }
+                  ]
+                }
+                then_ = LVar {
+                  name = <U c>
+                }
+                else_ = NULL
+              }
+              String {
+                val = <U in with 2 args, semicolon, and if>
+              }
+            ]
+          }
+        }
+      ]
+      elseBody = NULL
+    }
+    CaseMatch {
+      expr = Send {
+        receiver = NULL
+        method = <U baz>
+        args = [
+        ]
+      }
+      inBodies = [
+        InPattern {
+          pattern = MatchVar {
+            name = <U x>
+          }
+          guard = UnlessGuard {
+            condition = Send {
+              receiver = LVar {
+                name = <U x>
+              }
+              method = <U ==>
+              args = [
+                Integer {
+                  val = "1"
+                }
+              ]
+            }
+          }
+          body = String {
+            val = <U in with unless>
+          }
+        }
+        InPattern {
+          pattern = ArrayPattern {
+            elts = [
+              MatchVar {
+                name = <U a>
+              }
+              MatchVar {
+                name = <U b>
+              }
+            ]
+          }
+          guard = UnlessGuard {
+            condition = Send {
+              receiver = LVar {
+                name = <U b>
+              }
+              method = <U ==>
+              args = [
+                Integer {
+                  val = "2"
+                }
+              ]
+            }
+          }
+          body = String {
+            val = <U in with 2 args and unless>
+          }
+        }
+        InPattern {
+          pattern = ArrayPattern {
+            elts = [
+              MatchVar {
+                name = <U c>
+              }
+              MatchVar {
+                name = <U d>
+              }
+            ]
+          }
+          guard = NULL
+          body = Begin {
+            stmts = [
+              If {
+                condition = Send {
+                  receiver = LVar {
+                    name = <U c>
+                  }
+                  method = <U ==>
+                  args = [
+                    Integer {
+                      val = "3"
+                    }
+                  ]
+                }
+                then_ = NULL
+                else_ = LVar {
+                  name = <U c>
+                }
+              }
+              String {
+                val = <U in with 2 args, semicolon, and unless>
+              }
+            ]
+          }
+        }
+      ]
+      elseBody = NULL
+    }
   ]
 }
diff --git a/test/prism_regression/case_match.rb b/test/prism_regression/case_match.rb
@@ -43,3 +43,23 @@
   "one!"
   puts "surprise, multi-line!"
 end
+
+# pattern matching with if guards
+case bar
+in x if x == 1
+  "in with if"
+in a, b if b == 2
+  "in with 2 args and if"
+in c, d; c if c == 3
+  "in with 2 args, semicolon, and if"
+end
+
+# pattern matching with unless guards
+case baz
+in x unless x == 1
+  "in with unless"
+in a, b unless b == 2
+  "in with 2 args and unless"
+in c, d; c unless c == 3
+  "in with 2 args, semicolon, and unless"
+end
diff --git a/test/prism_regression/if.parse-tree.exp b/test/prism_regression/if.parse-tree.exp
@@ -1,6 +1,26 @@
-If {
-  condition = True {
-  }
-  then_ = NULL
-  else_ = NULL
+Begin {
+  stmts = [
+    If {
+      condition = Send {
+        receiver = NULL
+        method = <U condition>
+        args = [
+        ]
+      }
+      then_ = NULL
+      else_ = NULL
+    }
+    If {
+      condition = Send {
+        receiver = NULL
+        method = <U condition>
+        args = [
+        ]
+      }
+      then_ = String {
+        val = <U value>
+      }
+      else_ = NULL
+    }
+  ]
 }
diff --git a/test/prism_regression/if.rb b/test/prism_regression/if.rb
@@ -1,3 +1,5 @@
 # typed: false
 
-if true; end
+if condition; end
+
+"value" if condition
diff --git a/test/prism_regression/unless.parse-tree.exp b/test/prism_regression/unless.parse-tree.exp
@@ -1,8 +1,24 @@
-If {
-  condition = True {
-  }
-  then_ = NULL
-  else_ = String {
-    val = <U body>
-  }
+Begin {
+  stmts = [
+    If {
+      condition = True {
+      }
+      then_ = NULL
+      else_ = String {
+        val = <U body>
+      }
+    }
+    If {
+      condition = Send {
+        receiver = NULL
+        method = <U condition>
+        args = [
+        ]
+      }
+      then_ = NULL
+      else_ = String {
+        val = <U value>
+      }
+    }
+  ]
 }
diff --git a/test/prism_regression/unless.rb b/test/prism_regression/unless.rb
@@ -3,3 +3,5 @@
 unless true
   "body"
 end
+
+"value" unless condition
-Original file line number
+Diff line change
@@ Expand Up / @@ -3,3 +3,5 @@ @@
     unless true
       "body"
     end
+    "value" unless condition