Desugar pattern matching nodes

Author: allcre

ast/desugar/PrismDesugar.cc

@@ -541,21 +541,6 @@ void desugarPatternMatchingVars(InsSeq::STATS_store &vars, DesugarContext dctx,
     }
 }
 
-// Desugar `in` and `=>` oneline pattern matching
-ExpressionPtr desugarOnelinePattern(DesugarContext dctx, core::LocOffsets loc, parser::Node *match) {
-    auto matchExpr = MK::RaiseUnimplemented(loc);
-    auto bodyExpr = MK::RaiseUnimplemented(loc);
-    auto elseExpr = MK::EmptyTree();
-
-    InsSeq::STATS_store vars;
-    desugarPatternMatchingVars(vars, dctx, match);
-    if (!vars.empty()) {
-        bodyExpr = MK::InsSeq(match->loc, move(vars), move(bodyExpr));
-    }
-
-    return MK::If(loc, move(matchExpr), move(bodyExpr), move(elseExpr));
-}
-
 bool locReported = false;
 
 ClassDef::RHS_store scopeNodeToBody(DesugarContext dctx, unique_ptr<parser::Node> node) {
@@ -2153,14 +2138,8 @@ ExpressionPtr node2TreeImplBody(DesugarContext dctx, parser::Node *what) {
             [&](parser::MatchCurLine *matchCurLine) { desugaredByPrismTranslator(matchCurLine); },
             [&](parser::Redo *redo) { desugaredByPrismTranslator(redo); },
             [&](parser::EncodingLiteral *encodingLiteral) { desugaredByPrismTranslator(encodingLiteral); },
-            [&](parser::MatchPattern *pattern) {
-                auto res = desugarOnelinePattern(dctx, pattern->loc, pattern->rhs.get());
-                result = move(res);
-            },
-            [&](parser::MatchPatternP *pattern) {
-                auto res = desugarOnelinePattern(dctx, pattern->loc, pattern->rhs.get());
-                result = move(res);
-            },
+            [&](parser::MatchPattern *pattern) { desugaredByPrismTranslator(pattern); },
+            [&](parser::MatchPatternP *pattern) { desugaredByPrismTranslator(pattern); },
             [&](parser::EmptyElse *else_) { result = MK::EmptyTree(); },
             [&](parser::ResolvedConst *resolvedConst) {
                 result = make_expression<ConstantLit>(resolvedConst->loc, move(resolvedConst->symbol));

parser/prism/Translator.cc

@@ -87,6 +87,23 @@ static void collectPatternMatchingVars(ast::InsSeq::STATS_store &vars, parser::N
     }
 }
 
+// Desugar `in` and `=>` oneline pattern matching (mirrors desugarOnelinePattern in PrismDesugar.cc)
+static ast::ExpressionPtr desugarOnelinePattern(core::LocOffsets loc, parser::Node *match) {
+    auto matchExpr = MK::RaiseUnimplemented(loc);
+    auto bodyExpr = MK::RaiseUnimplemented(loc);
+    auto elseExpr = MK::EmptyTree();
+
+    ast::InsSeq::STATS_store vars;
+    collectPatternMatchingVars(vars, match);
+
+    if (!vars.empty()) {
+        auto matchLoc = match != nullptr ? match->loc : loc;
+        bodyExpr = MK::InsSeq(matchLoc, move(vars), move(bodyExpr));
+    }
+
+    return MK::If(loc, move(matchExpr), move(bodyExpr), move(elseExpr));
+}
+
 // Allocates a new `NodeWithExpr` with a pre-computed `ExpressionPtr` AST.
 template <typename SorbetNode, typename... TArgs>
 unique_ptr<parser::Node> Translator::make_node_with_expr(ast::ExpressionPtr desugaredExpr, TArgs &&...args) const {
@@ -2912,15 +2929,25 @@ unique_ptr<parser::Node> Translator::translate(pm_node_t *node, bool preserveCon
             auto value = patternTranslate(matchRequiredNode->value);
             auto pattern = patternTranslate(matchRequiredNode->pattern);
 
-            return make_unique<parser::MatchPattern>(location, move(value), move(pattern));
+            if (!directlyDesugar || !hasExpr(value, pattern)) {
+                return make_unique<parser::MatchPattern>(location, move(value), move(pattern));
+            }
+
+            auto expr = desugarOnelinePattern(location, pattern.get());
+            return make_node_with_expr<parser::MatchPattern>(move(expr), location, move(value), move(pattern));
         }
         case PM_MATCH_PREDICATE_NODE: {
             auto matchPredicateNode = down_cast<pm_match_predicate_node>(node);
 
             auto value = patternTranslate(matchPredicateNode->value);
             auto pattern = patternTranslate(matchPredicateNode->pattern);
 
-            return make_unique<parser::MatchPatternP>(location, move(value), move(pattern));
+            if (!directlyDesugar || !hasExpr(value, pattern)) {
+                return make_unique<parser::MatchPatternP>(location, move(value), move(pattern));
+            }
+
+            auto expr = desugarOnelinePattern(location, pattern.get());
+            return make_node_with_expr<parser::MatchPatternP>(move(expr), location, move(value), move(pattern));
         }
         case PM_MATCH_WRITE_NODE: { // A regex match that assigns to a local variable, like `a =~ /wat/`
             auto matchWriteNode = down_cast<pm_match_write_node>(node);