feat(parser): parse block captures in fn bodys properly

goruby · Feb 20, 2018 · 22c1778 · 22c1778
1 parent cd93ac6
commit 22c1778
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/parser/parser.go b/parser/parser.go
@@ -207,6 +207,7 @@ func (p *parser) init(fset *gotoken.FileSet, filename string, src []byte, mode M
 	p.registerInfix(token.INT, p.parseCallArgument)
 	p.registerInfix(token.STRING, p.parseCallArgument)
 	p.registerInfix(token.SYMBEG, p.parseCallArgument)
+	p.registerInfix(token.CAPTURE, p.parseCallArgument)
 	p.registerInfix(token.SELF, p.parseCallArgument)
 	p.registerInfix(token.LBRACE, p.parseCallBlock)
 	p.registerInfix(token.DO, p.parseCallBlock)

diff --git a/parser/parser_test.go b/parser/parser_test.go
@@ -98,6 +98,31 @@ func TestBlockCapture(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "block capture as arg on call in func body",
+			input: `
+			def foo
+				each &block
+			end`,
+			result: &ast.FunctionLiteral{
+				Name:       &ast.Identifier{Value: "foo"},
+				Parameters: []*ast.FunctionParameter{},
+				Body: &ast.BlockStatement{
+					Statements: []ast.Statement{
+						&ast.ExpressionStatement{
+							Expression: &ast.ContextCallExpression{
+								Function: &ast.Identifier{Value: "each"},
+								Arguments: []ast.Expression{
+									&ast.BlockCapture{
+										Name: &ast.Identifier{Value: "block"},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, tt := range tests {