From 2ec029f415ed9a8a1448a5bb389d79084d606035 Mon Sep 17 00:00:00 2001
From: Ruud Kamphuis <ruudk@users.noreply.github.com>
Date: Wed, 20 Jun 2018 00:04:12 +0200
Subject: [PATCH] Support `else if` (#58)

---
 README.md            |  4 +++-
 ast/if_expression.go | 15 +++++++++++++++
 compiler.go          | 26 ++++++++++++++++++++++----
 if_test.go           | 26 ++++++++++++++++++++++++++
 parser/parser.go     | 44 ++++++++++++++++++++++++++++++++++++++++----
 5 files changed, 106 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index f1a94cf..bc8441e 100644
--- a/README.md
+++ b/README.md
@@ -93,12 +93,14 @@ You can add comments like this:
 
 ## If/Else Statements
 
-The basic syntax of `if/else` statements is as follows:
+The basic syntax of `if/else if/else` statements is as follows:
 
 ```erb
 <%
 if (true) {
   # do something
+} else if (false) {
+  # do something
 } else {
   # do something else
 }
diff --git a/ast/if_expression.go b/ast/if_expression.go
index eb309f4..6f7e9cf 100644
--- a/ast/if_expression.go
+++ b/ast/if_expression.go
@@ -8,9 +8,16 @@ type IfExpression struct {
 	TokenAble
 	Condition Expression
 	Block     *BlockStatement
+	ElseIf    []*ElseIfExpression
 	ElseBlock *BlockStatement
 }
 
+type ElseIfExpression struct {
+	TokenAble
+	Condition Expression
+	Block     *BlockStatement
+}
+
 func (ie *IfExpression) expressionNode() {}
 
 func (ie *IfExpression) String() string {
@@ -22,6 +29,14 @@ func (ie *IfExpression) String() string {
 	out.WriteString(ie.Block.String())
 	out.WriteString(" }")
 
+	for _, elseIf := range ie.ElseIf {
+		out.WriteString(" } else if (")
+		out.WriteString(elseIf.Condition.String())
+		out.WriteString(") { ")
+		out.WriteString(elseIf.Block.String())
+		out.WriteString(" }")
+	}
+
 	if ie.ElseBlock != nil {
 		out.WriteString(" } else { ")
 		out.WriteString(ie.ElseBlock.String())
diff --git a/compiler.go b/compiler.go
index ddbf4c4..273e975 100644
--- a/compiler.go
+++ b/compiler.go
@@ -182,13 +182,31 @@ func (c *compiler) evalIfExpression(node *ast.IfExpression) (interface{}, error)
 		}
 	}
 
-	var r interface{}
 	if c.isTruthy(con) {
 		return c.evalBlockStatement(node.Block)
-	} else {
-		if node.ElseBlock != nil {
-			return c.evalBlockStatement(node.ElseBlock)
+	}
+
+	return c.evalElseAndElseIfExpressions(node)
+}
+
+func (c *compiler) evalElseAndElseIfExpressions(node *ast.IfExpression) (interface{}, error) {
+	// fmt.Println("evalElseIfExpression")
+	var r interface{}
+	for _, eiNode := range node.ElseIf {
+		eiCon, err := c.evalExpression(eiNode.Condition)
+		if err != nil {
+			if errors.Cause(err) != ErrUnknownIdentifier {
+				return nil, errors.WithStack(err)
+			}
 		}
+
+		if c.isTruthy(eiCon) {
+			return c.evalBlockStatement(eiNode.Block)
+		}
+	}
+
+	if node.ElseBlock != nil {
+		return c.evalBlockStatement(node.ElseBlock)
 	}
 
 	return r, nil
diff --git a/if_test.go b/if_test.go
index 60004b1..5754aa0 100644
--- a/if_test.go
+++ b/if_test.go
@@ -134,6 +134,32 @@ func Test_Render_If_Else_True(t *testing.T) {
 	r.Equal("<p>hi</p>", s)
 }
 
+func Test_Render_If_Else_If_Else_True(t *testing.T) {
+	r := require.New(t)
+	ctx := NewContext()
+	input := `<p><%= if (state == "foo") { %>hi foo<% } else if (state == "bar") { %>hi bar<% } else if (state == "fizz") { %>hi fizz<% } else { %>hi buzz<% } %></p>`
+
+	ctx.Set("state", "foo")
+	s, err := Render(input, ctx)
+	r.NoError(err)
+	r.Equal("<p>hi foo</p>", s)
+
+	ctx.Set("state", "bar")
+	s, err = Render(input, ctx)
+	r.NoError(err)
+	r.Equal("<p>hi bar</p>", s)
+
+	ctx.Set("state", "fizz")
+	s, err = Render(input, ctx)
+	r.NoError(err)
+	r.Equal("<p>hi fizz</p>", s)
+
+	ctx.Set("state", "buzz")
+	s, err = Render(input, ctx)
+	r.NoError(err)
+	r.Equal("<p>hi buzz</p>", s)
+}
+
 func Test_Render_If_Matches(t *testing.T) {
 	r := require.New(t)
 	input := `<p><%= if ("foo" ~= "bar") { return "hi" } else { return "bye" } %></p>`
diff --git a/parser/parser.go b/parser/parser.go
index 8674dd5..a03caac 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -473,16 +473,52 @@ func (p *parser) parseIfExpression() ast.Expression {
 
 	expression.Block = p.parseBlockStatement()
 
-	if p.peekTokenIs(token.ELSE) {
+	for p.peekTokenIs(token.ELSE) {
 		p.nextToken()
 
-		if !p.expectPeek(token.LBRACE) {
-			return nil
+		if p.peekTokenIs(token.IF) {
+			p.nextToken()
+
+			ifElseExp := p.parseElseIfExpression()
+
+			if ifElseExp == nil {
+				return nil
+			}
+
+			expression.ElseIf = append(expression.ElseIf, ifElseExp)
+		} else {
+			if !p.expectPeek(token.LBRACE) {
+				return nil
+			}
+
+			expression.ElseBlock = p.parseBlockStatement()
 		}
+	}
+
+	return expression
+}
+
+func (p *parser) parseElseIfExpression() *ast.ElseIfExpression {
+	// fmt.Println("parseElseIfExpression")
+	expression := &ast.ElseIfExpression{TokenAble: ast.TokenAble{p.curToken}}
+
+	if !p.expectPeek(token.LPAREN) {
+		return nil
+	}
+
+	p.nextToken()
+	expression.Condition = p.parseExpression(LOWEST)
+
+	if !p.expectPeek(token.RPAREN) {
+		return nil
+	}
 
-		expression.ElseBlock = p.parseBlockStatement()
+	if !p.expectPeek(token.LBRACE) {
+		return nil
 	}
 
+	expression.Block = p.parseBlockStatement()
+
 	return expression
 }