✨ Parse IF-ELSE statements

ast/ast.go

@@ -160,11 +160,19 @@ func (x *InfixExpression) String() string {
 // fixed, fixed-bytes or ufixed. NOT a Contract, Function, mapping (these are
 // the four other types)
 type ElementaryType struct {
-	Pos   token.Pos   // position of the type keyword e.g. `a` in "address"
-	Kind  token.Token // type of the literal e.g. token.ADDRESS, token.UINT_256, token.BOOL
-	Value string      // type literal value e.g. "address", "uint256", "bool" as a string
+	Pos  token.Pos   // position of the type keyword e.g. `a` in "address"
+	Kind token.Token // type of the literal e.g. token.ADDRESS, token.UINT_256, token.BOOL
+	// nil when used as declaration e.g. `address a;`
+	// BUT also used in expressions e.g. `return uint256(a + b)`. Then it
+	// contains the expression a + b
+	Value Expression
 }
 
+// WARNING ElementaryType implements both Type and Expression interfaces.
+// It can be used as a type e.g. in variable declaration `uint256 x;` OR
+// as an expression in return statement `return uint256(a + b);`
+func (x *ElementaryType) expressionNode() {}
+
 // FunctionType represents a Solidity's function type. NOT TO BE CONFUSED WITH
 // FUNCTION DECLARATION. FunctionType is a weird thing that no one uses (lol) e.g.
 // ```solidity
@@ -216,7 +224,7 @@ type ParamList struct {
 // Start() and End() implementations for Expression type Nodes
 
 func (x *ElementaryType) Start() token.Pos { return x.Pos }
-func (x *ElementaryType) End() token.Pos   { return token.Pos(int(x.Pos) + len(x.Value)) }
+func (x *ElementaryType) End() token.Pos   { return token.Pos(int(x.Pos) + int(x.End())) }
 
 // expressionNode() implementations to ensure that only expressions and types
 // can be assigned to an Expression. This is useful if by mistake we try to use
@@ -225,7 +233,13 @@ func (x *ElementaryType) End() token.Pos   { return token.Pos(int(x.Pos) + len(x
 func (*ElementaryType) typeNode() {}
 
 // String() implementations for Types
-func (x *ElementaryType) String() string { return x.Value }
+func (x *ElementaryType) String() string {
+	var out bytes.Buffer
+	out.WriteString(x.Kind.Literal)
+	out.WriteString(" ")
+	out.WriteString(x.Value.String())
+	return out.String()
+}
 
 /*~*~*~*~*~*~*~*~*~*~*~*~* Statements *~*~*~*~*~*~*~*~*~*~*~*~*~*/
 
@@ -279,6 +293,12 @@ type (
 		Expression Expression // expression to be evaluated
 	}
 
+	IfStatement struct {
+		Pos         token.Pos  // position of the "if" keyword
+		Condition   Expression // condition to be evaluated
+		Consequence Statement  // consequence happens if the condition is true; or nil
+		Alternative Statement  // alternative happens if the condition is false; or nil
+	}
 )
 
 // Start() and End() implementations for Statement type Nodes
@@ -298,13 +318,28 @@ func (s *ReturnStatement) End() token.Pos {
 }
 func (s *ExpressionStatement) Start() token.Pos { return s.Pos }
 func (s *ExpressionStatement) End() token.Pos   { return s.Expression.End() }
+func (s *IfStatement) Start() token.Pos         { return s.Pos }
+func (s *IfStatement) End() token.Pos {
+	endPos := s.Pos + 2 // The length of "if".
+
+	if s.Consequence != nil {
+		endPos = s.Consequence.End()
+	}
+
+	if s.Alternative != nil {
+		endPos = s.Alternative.End()
+	}
+
+	return endPos
+}
 
 // statementNode() ensures that only statement nodes can be assigned to a Statement.
 func (*BlockStatement) statementNode()               {}
 func (*UncheckedBlockStatement) statementNode()      {}
 func (*VariableDeclarationStatement) statementNode() {}
 func (*ReturnStatement) statementNode()              {}
 func (*ExpressionStatement) statementNode()          {}
+func (*IfStatement) statementNode()                  {}
 
 // String() implementations for Statements
 
@@ -362,6 +397,21 @@ func (s *ExpressionStatement) String() string {
 	return ""
 }
 
+func (s *IfStatement) String() string {
+	var out bytes.Buffer
+	out.WriteString("if")
+	out.WriteString(s.Condition.String())
+	out.WriteString(" ")
+	out.WriteString(s.Consequence.String())
+
+	if s.Alternative != nil {
+		out.WriteString("else ")
+		out.WriteString(s.Alternative.String())
+	}
+
+	return out.String()
+}
+
 /*~*~*~*~*~*~*~*~*~*~*~*~ Declarations ~*~*~*~*~*~*~*~*~*~*~*~*~*/
 
 // @TODO: Add Contract declaration

parser/exprparsing.go

@@ -274,3 +274,36 @@ func (p *Parser) parseGroupedExpression() ast.Expression {
 
 	return exp
 }
+
+func (p *Parser) parseElementaryType() ast.Expression {
+	if p.trace {
+		defer un(trace("parseElementaryType"))
+	}
+
+	et := &ast.ElementaryType{
+		Pos: p.currTkn.Pos,
+		Kind: token.Token{
+			Type:    p.currTkn.Type,
+			Literal: p.currTkn.Literal,
+			Pos:     p.currTkn.Pos,
+		},
+	}
+
+	if !p.expectPeek(token.LPAREN) {
+		return nil
+	}
+
+	p.nextToken()
+	et.Value = p.parseExpression(LOWEST)
+
+	if !p.expectPeek(token.RPAREN) {
+		return nil
+	}
+
+	// Semi-colon is optional.
+	if p.peekTkn.Type == token.SEMICOLON {
+		p.nextToken()
+	}
+
+	return et
+}

parser/parser.go

@@ -40,7 +40,8 @@ func (p *Parser) Init(file *token.File) {
 	p.registerPrefix(token.HEX_NUMBER, p.parseNumberLiteral)
 	p.registerPrefix(token.TRUE_LITERAL, p.parseBooleanLiteral)
 	p.registerPrefix(token.FALSE_LITERAL, p.parseBooleanLiteral)
-	p.registerPrefix(token.LPAREN, p.parseGroupedExpression)
+
+	registerPrefixElementaryTypes(p)
 
 	// Prefix Expressions
 	p.registerPrefix(token.NOT, p.parsePrefixExpression)
@@ -49,6 +50,7 @@ func (p *Parser) Init(file *token.File) {
 	p.registerPrefix(token.DEC, p.parsePrefixExpression)
 	p.registerPrefix(token.DELETE, p.parsePrefixExpression)
 	p.registerPrefix(token.SUB, p.parsePrefixExpression)
+	p.registerPrefix(token.LPAREN, p.parseGroupedExpression)
 
 	// Infix Expressions
 	p.infixParseFns = make(map[token.TokenType]infixParseFn)
@@ -87,6 +89,12 @@ func (p *Parser) Init(file *token.File) {
 	p.registerInfix(token.ASSIGN_MOD, p.parseInfixExpression)
 }
 
+func registerPrefixElementaryTypes(p *Parser) {
+	for _, tkType := range token.GetElementaryTypes() {
+		p.registerPrefix(tkType, p.parseElementaryType)
+	}
+}
+
 func (p *Parser) ToggleTracing() {
 	p.trace = !p.trace
 }
@@ -193,7 +201,7 @@ func (p *Parser) parseStateVariableDeclaration() *ast.StateVariableDeclaration {
 	decl.Type = &ast.ElementaryType{
 		Pos:   p.currTkn.Pos,
 		Kind:  p.currTkn,
-		Value: p.currTkn.Literal,
+		Value: nil,
 	}
 
 	p.nextToken()
@@ -255,11 +263,14 @@ func (p *Parser) parseStatement() ast.Statement {
 		return p.parseExpressionStatement()
 	case token.IsElementaryType(tkType):
 		// @TODO: Implement other types that variables can have.
+		// @TODO: return address(0) and similar should be handled here
 		return p.parseVariableDeclarationStatement()
 	case tkType == token.LBRACE:
 		return p.parseBlockStatement()
 	case tkType == token.RETURN:
 		return p.parseReturnStatement()
+	case tkType == token.IF:
+		return p.parseIfStatement()
 	}
 }
 
@@ -344,7 +355,7 @@ func (p *Parser) parseVariableDeclarationStatement() *ast.VariableDeclarationSta
 	vdStmt.Type = &ast.ElementaryType{
 		Pos:   p.currTkn.Pos,
 		Kind:  p.currTkn,
-		Value: p.currTkn.Literal,
+		Value: nil,
 	}
 
 	p.nextToken()
@@ -383,8 +394,12 @@ func (p *Parser) parseReturnStatement() *ast.ReturnStatement {
 	retStmt := &ast.ReturnStatement{}
 	retStmt.Pos = p.currTkn.Pos
 
-	// @TODO: Parse expression
-	for !p.currTknIs(token.SEMICOLON) {
+	// Advance to the next token; parse the expression.
+	p.nextToken()
+	retStmt.Result = p.parseExpression(LOWEST)
+
+	// Semicolon is optional on purpose to make it easier to type stuff into
+	if p.peekTkn.Type == token.SEMICOLON {
 		p.nextToken()
 	}
 
@@ -409,6 +424,44 @@ func (p *Parser) parseExpressionStatement() *ast.ExpressionStatement {
 	return exprStmt
 }
 
+func (p *Parser) parseIfStatement() *ast.IfStatement {
+	if p.trace {
+		defer un(trace("parseIfStatement"))
+	}
+
+	// 1. If keyword
+	ifStmt := &ast.IfStatement{
+		Pos: p.currTkn.Pos,
+	}
+
+	// 2. Parenthesis of condition should be next; consume the lparen token.
+	if !p.expectPeek(token.LPAREN) {
+		return nil
+	}
+
+	// 3. Advance to condition; parse the condition.
+	p.nextToken()
+	ifStmt.Condition = p.parseExpression(LOWEST)
+
+	// 4. Consume the rparen token.
+	if !p.expectPeek(token.RPAREN) {
+		return nil
+	}
+
+	// 5. Advance; parse the consequence statement.
+	p.nextToken()
+	ifStmt.Consequence = p.parseStatement()
+
+	// 6. If there is an else, parse it.
+	if p.peekTkn.Type == token.ELSE {
+		p.nextToken() // Sitting on ELSE
+		p.nextToken() // Sitting on the <<statement>> after ELSE
+		ifStmt.Alternative = p.parseStatement()
+	}
+
+	return ifStmt
+}
+
 // expectPeek checks if the next token is of the expected type.
 // If it is it advances the tokens.
 func (p *Parser) expectPeek(t token.TokenType) bool {