From 77d8728db4963d7dd0dd1af5692b39ce20ecc622 Mon Sep 17 00:00:00 2001
From: Alexandre Terrasa <alexandre.terrasa@shopify.com>
Date: Thu, 13 Aug 2020 14:35:25 -0400
Subject: [PATCH] Fix parsing context for lambda and modules

Signed-off-by: Alexandre Terrasa <alexandre.terrasa@shopify.com>
---
 test/whitequark/test_context_lambda.rb        |  4 ++++
 ...ontext_lambda.rb.parse-tree-whitequark.exp | 23 +++++++++++++++++++
 third_party/parser/cc/grammars/typedruby.ypp  | 11 +++++++--
 .../parser/include/ruby_parser/context.hh     |  1 +
 4 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 test/whitequark/test_context_lambda.rb
 create mode 100644 test/whitequark/test_context_lambda.rb.parse-tree-whitequark.exp

diff --git a/test/whitequark/test_context_lambda.rb b/test/whitequark/test_context_lambda.rb
new file mode 100644
index 00000000000..b223593cba3
--- /dev/null
+++ b/test/whitequark/test_context_lambda.rb
@@ -0,0 +1,4 @@
+# typed: strict
+
+module A; ->(a = (return; 1)) {}; end
+class B; ->(a = (return; 1)) {}; end
diff --git a/test/whitequark/test_context_lambda.rb.parse-tree-whitequark.exp b/test/whitequark/test_context_lambda.rb.parse-tree-whitequark.exp
new file mode 100644
index 00000000000..7140399e34d
--- /dev/null
+++ b/test/whitequark/test_context_lambda.rb.parse-tree-whitequark.exp
@@ -0,0 +1,23 @@
+s(:begin,
+  s(:module,
+    s(:const, nil, :A),
+    s(:block,
+      s(:send,
+        s(:const, nil, :Kernel), :lambda),
+      s(:args,
+        s(:optarg, :a,
+          s(:begin,
+            s(:begin,
+              s(:return),
+              s(:int, "1"))))), nil)),
+  s(:class,
+    s(:const, nil, :B), nil,
+    s(:block,
+      s(:send,
+        s(:const, nil, :Kernel), :lambda),
+      s(:args,
+        s(:optarg, :a,
+          s(:begin,
+            s(:begin,
+              s(:return),
+              s(:int, "1"))))), nil)))
diff --git a/third_party/parser/cc/grammars/typedruby.ypp b/third_party/parser/cc/grammars/typedruby.ypp
index d4f2c963551..dedc83416f5 100644
--- a/third_party/parser/cc/grammars/typedruby.ypp
+++ b/third_party/parser/cc/grammars/typedruby.ypp
@@ -1481,10 +1481,14 @@ int yylex(parser::semantic_type *lval, ruby_parser::typedruby27 &driver) {
                         delimited_block->end);
                       DIAGCHECK();
                     }
-                | tLAMBDA lambda
+                | tLAMBDA
+                    {
+                      driver.lex.context.push(Context::State::LAMBDA);
+                    }
+                    lambda
                     {
                       auto lambda_call = driver.build.callLambda(self, $1);
-                      auto &lambda = $2;
+                      auto &lambda = $3;
                       $$ = driver.build.block(self, lambda_call,
                         lambda->begin,
                         lambda->args,
@@ -1584,6 +1588,7 @@ int yylex(parser::semantic_type *lval, ruby_parser::typedruby27 &driver) {
                     {
                       driver.lex.extend_static();
                       driver.lex.cmdarg.push(false);
+                      driver.lex.context.push(Context::State::MODULE);
                     }
                     bodystmt kEND
                     {
@@ -1599,6 +1604,7 @@ int yylex(parser::semantic_type *lval, ruby_parser::typedruby27 &driver) {
 
                       driver.lex.cmdarg.pop();
                       driver.lex.unextend();
+                      driver.lex.context.pop();
                     }
                 | kDEF fname
                     {
@@ -1955,6 +1961,7 @@ opt_block_args_tail:
                     }
                   f_larglist
                     {
+                      driver.lex.context.pop();
                       driver.lex.cmdarg.push(false);
                     }
                   lambda_body
diff --git a/third_party/parser/include/ruby_parser/context.hh b/third_party/parser/include/ruby_parser/context.hh
index 202334d2afa..5bb1de0d7aa 100644
--- a/third_party/parser/include/ruby_parser/context.hh
+++ b/third_party/parser/include/ruby_parser/context.hh
@@ -11,6 +11,7 @@ class Context {
 public:
     enum class State {
         CLASS,
+        MODULE,
         SCLASS,
         DEF,
         DEFS,