fn literals

3 files changed, 144 insertions, 0 deletions

diff --git a/ast/ast.go b/ast/ast.go
index 31b1be1..424b4ac 100644
--- a/ast/ast.go
+++ b/ast/ast.go
@@ -3,6 +3,7 @@ package ast
 import (
 	"bytes"
 	"mana/tokens"
+	"strings"
 )
 
 type Node interface {
@@ -247,3 +248,29 @@ func (ie *IfExpression) String() string {
 
 	return out.String()
 }
+
+type FunctionLiteral struct {
+	Token 		tokens.Token // the 'fn' token
+	Parameters 	[]*Identifier
+	Body 		*BlockStatement
+}
+
+func (fl *FunctionLiteral) expressionNode() 	{}
+func (fl *FunctionLiteral) TokenLiteral() string { return fl.Token.Literal }
+func (fl *FunctionLiteral) String() string {
+	var out bytes.Buffer
+
+	params := []string{}
+
+	for _, p := range fl.Parameters {
+		params = append(params, p.String())
+	}
+
+	out.WriteString(fl.TokenLiteral())
+	out.WriteString("(")
+	out.WriteString(strings.Join(params, ", "))
+	out.WriteString(")")
+	out.WriteString(fl.Body.String())
+
+	return out.String()
+}
diff --git a/parser/parser.go b/parser/parser.go
index e38e54c..afbdf24 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -68,6 +68,7 @@ func New(l *lexer.Lexer) *Parser {
 	p.registerPrefix(tokens.TRUE, p.parseBoolean)
 	p.registerPrefix(tokens.FALSE, p.parseBoolean)
 	p.registerPrefix(tokens.IF, p.parseIfExpression)
+	p.registerPrefix(tokens.FUNCTION, p.parseFunctionLiteral)
 
 	// Initialize the infix parse functions.
 	p.infixParseFns = make(map[tokens.TokenType]infixParseFn)
@@ -330,6 +331,48 @@ func (p *Parser) parseBlockStatement() *ast.BlockStatement {
 	return block
 }
 
+// parseFunctionLiteral parses a function literal.
+func (p *Parser) parseFunctionLiteral() ast.Expression {
+	lit := &ast.FunctionLiteral{Token: p.curToken}
+
+	if !p.expectPeek(tokens.LPAREN) { return nil }
+
+	lit.Parameters = p.parseFunctionParameters()
+
+	if !p.expectPeek(tokens.LBRACE) { return nil }
+
+	lit.Body = p.parseBlockStatement()
+
+	return lit
+
+}
+
+// parseFunctionParameters parses function parameters.
+func (p *Parser) parseFunctionParameters() []*ast.Identifier {
+	identifiers := []*ast.Identifier{}
+
+	if p.peekTokenIs(tokens.RPAREN) {
+		p.nextToken()
+		return identifiers
+	}
+
+	p.nextToken()
+
+	ident := &ast.Identifier{Token: p.curToken, Value: p.curToken.Literal}
+	identifiers = append(identifiers, ident)
+
+	for p.peekTokenIs(tokens.COMMA) {
+		p.nextToken()
+		p.nextToken()
+		ident := &ast.Identifier{Token: p.curToken, Value: p.curToken.Literal}
+		identifiers = append(identifiers, ident)
+	}
+
+	if !p.expectPeek(tokens.RPAREN) { return nil }
+
+	return identifiers
+}
+
 // curTokenIs returns true if the current token is of the given type.
 func (p *Parser) curTokenIs(t tokens.TokenType) bool {
 	return p.curToken.Type == t
diff --git a/parser/parser_test.go b/parser/parser_test.go
index f26fa1e..415b8cb 100644
--- a/parser/parser_test.go
+++ b/parser/parser_test.go
@@ -652,3 +652,77 @@ func TestIfElseExpression(t *testing.T) {
 		return
 	}
 }
+
+func TestFunctionLiteralParsing(t *testing.T) {
+	input := "fn(x, y) { x + y; }"
+
+	var l *lexer.Lexer = lexer.New(input)
+	var p *Parser = New(l)
+	var program *ast.Program = p.ParseProgram()
+	checkParserErrors(t, p)
+
+	if len(program.Statements) != 1 {
+		t.Fatalf("program.Statements does not contain %d statements. got=%d", 1, len(program.Statements))
+	}
+
+	stmt, ok := program.Statements[0].(*ast.ExpressionStatement)
+
+	if !ok {
+		t.Fatalf("program.Statements[0] is not ast.ExpressionStatement. got=%T", program.Statements[0])
+	}
+
+	function, ok := stmt.Expression.(*ast.FunctionLiteral)
+
+	if !ok {
+		t.Fatalf("stmt.Expression is not ast.FunctionLiteral. got=%T", stmt.Expression)
+	}
+
+	if len(function.Parameters) != 2 {
+		t.Fatalf("function literal parameters wrong. want 2, got=%d\n", len(function.Parameters))
+	}
+
+	testLiteralExpression(t, function.Parameters[0], "x")
+	testLiteralExpression(t, function.Parameters[1], "y")
+
+	if len(function.Body.Statements) != 1 {
+		t.Fatalf("function.Body.Statements has not 1 statements. got=%d", len(function.Body.Statements))
+	}
+
+	bodyStmt, ok := function.Body.Statements[0].(*ast.ExpressionStatement)
+
+	if !ok {
+		t.Fatalf("function body stmt is not ast.ExpressionStatement. got=%T", function.Body.Statements[0])
+	}
+
+	testInfixExpression(t, bodyStmt.Expression, "x", "+", "y")
+}
+
+func TestFunctionParameterParsing(t *testing.T) {
+	var tests = []struct {
+		input          string
+		expectedParams []string
+	}{
+		{input: "fn() {};", expectedParams: []string{}},
+		{input: "fn(x) {};", expectedParams: []string{"x"}},
+		{input: "fn(x, y, z) {};", expectedParams: []string{"x", "y", "z"}},
+	}
+
+	for _, tt := range tests {
+		var l *lexer.Lexer = lexer.New(tt.input)
+		var p *Parser = New(l)
+
+		var program *ast.Program = p.ParseProgram()
+		checkParserErrors(t, p)
+		
+		stmt := program.Statements[0].(*ast.ExpressionStatement)
+		function := stmt.Expression.(*ast.FunctionLiteral)
+
+		if len(function.Parameters) != len(tt.expectedParams) {
+			t.Errorf("length parameters wrong. want %d, got=%d\n", len(tt.expectedParams), len(function.Parameters))
+		}
+
+		for i, ident := range tt.expectedParams {
+			testLiteralExpression(t, function.Parameters[i], ident)
+		}
+	}
+}