From 3d9db07da2a5047dbb94278be4e189da7e4437dc Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Fri, 8 Dec 2017 22:34:11 -0600
Subject: [PATCH] Fix formatting of else statements with functions

The formatter was not handling else statements for functions properly.
The function args were being written after the else keyword which
resulted in rules that did not parse.
---
 format/format.go                     | 15 ++++++++-------
 format/testfiles/test.rego           | 14 ++++++++++++++
 format/testfiles/test.rego.formatted | 15 +++++++++++++++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/format/format.go b/format/format.go
index 50c2186c1d..01015b8b5d 100644
--- a/format/format.go
+++ b/format/format.go
@@ -91,7 +91,7 @@ func Ast(x interface{}) (formatted []byte, err error) {
 	case *ast.Import:
 		w.writeImports([]*ast.Import{x}, nil)
 	case *ast.Rule:
-		w.writeRule(x, nil)
+		w.writeRule(x, false, nil)
 	case *ast.Head:
 		w.writeHead(x, nil)
 	case ast.Body:
@@ -191,13 +191,13 @@ func (w *writer) writeComments(comments []*ast.Comment) {
 func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.Comment {
 	for _, rule := range rules {
 		comments = w.insertComments(comments, rule.Location)
-		comments = w.writeRule(rule, comments)
+		comments = w.writeRule(rule, false, comments)
 		w.blankLine()
 	}
 	return comments
 }
 
-func (w *writer) writeRule(rule *ast.Rule, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
 	if rule == nil {
 		return comments
 	}
@@ -212,9 +212,9 @@ func (w *writer) writeRule(rule *ast.Rule, comments []*ast.Comment) []*ast.Comme
 	// `foo = {"a": "b"} { true }` in the AST. We want to preserve that notation
 	// in the formatted code instead of expanding the bodies into rules, so we
 	// pretend that the rule has no body in this case.
-	isExpandedConst := rule.Body.Equal(ast.NewBody(ast.NewExpr(ast.BooleanTerm(true))))
+	isExpandedConst := rule.Body.Equal(ast.NewBody(ast.NewExpr(ast.BooleanTerm(true)))) && rule.Else == nil
 
-	if len(rule.Body) == 0 || isExpandedConst {
+	if (len(rule.Body) == 0 || isExpandedConst) && !isElse {
 		w.endLine()
 		return comments
 	}
@@ -241,8 +241,9 @@ func (w *writer) writeRule(rule *ast.Rule, comments []*ast.Comment) []*ast.Comme
 	if rule.Else != nil {
 		w.blankLine()
 		rule.Else.Head.Name = ast.Var("else")
-		comments = w.insertComments(comments, rule.Else.Loc())
-		comments = w.writeRule(rule.Else, comments)
+		rule.Else.Head.Args = nil
+		comments = w.insertComments(comments, rule.Else.Head.Location)
+		comments = w.writeRule(rule.Else, true, comments)
 	}
 	return comments
 }
diff --git a/format/testfiles/test.rego b/format/testfiles/test.rego
index 1275fc7aca..be69f241fe 100644
--- a/format/testfiles/test.rego
+++ b/format/testfiles/test.rego
@@ -56,6 +56,20 @@ fn(x) = y {
 		y = x
 }
 
+fn_else(x) = 1 {
+    true
+} # foo
+else  =
+# bar
+2
+{
+    true
+} else = 3
+# baz
+{
+    false
+}
+
 long(x) = true {
     x = "foo %host"
     }
diff --git a/format/testfiles/test.rego.formatted b/format/testfiles/test.rego.formatted
index f7e63752a8..1a9ea91d86 100644
--- a/format/testfiles/test.rego.formatted
+++ b/format/testfiles/test.rego.formatted
@@ -59,6 +59,21 @@ fn(x) = y {
 	y = x
 }
 
+fn_else(x) = 1 {
+	true
+} # foo
+
+# bar
+else = 2 {
+	true
+}
+
+else = 3 {
+	# baz
+
+	false
+}
+
 long(x) {
 	x = "foo %host"
 }