hclwrite: add space between "in" keyword and expr in for expression

Our normal ruleset thinks that the "in" keyword here is a variable
reference and so writes it as "in[y]". Since there's never any reason for
a variable to appear immediately after another variable, we can check
for a preceding identifier as a heuristic to recognize whether in is
probably being used as a keyword rather than as a variable.

This is not exact, but the only time this should be a false positive is
if there were a syntax error in the input, and we don't make any
guarantees about the result in that case anyway.

This fixes #52.

3 files changed, 42 insertions, 1 deletions

diff --git a/hclwrite/format.go b/hclwrite/format.go
index a084935..c699195 100644
--- a/hclwrite/format.go
+++ b/hclwrite/format.go
@@ -4,6 +4,8 @@ import (
 	"github.com/hashicorp/hcl2/hcl/hclsyntax"
 )
 
+var inKeyword = hclsyntax.Keyword([]byte{'i', 'n'})
+
 // placeholder token used when we don't have a token but we don't want
 // to pass a real "nil" and complicate things with nil pointer checks
 var nilToken = &Token{
@@ -247,6 +249,15 @@ func spaceAfterToken(subject, before, after *Token) bool {
 		// No extra spaces within templates
 		return false
 
+	case inKeyword.TokenMatches(subject.asHCLSyntax()) && before.Type == hclsyntax.TokenIdent:
+		// This is a special case for inside for expressions where a user
+		// might want to use a literal tuple constructor:
+		// [for x in [foo]: x]
+		// ... in that case, we would normally produce in[foo] thinking that
+		// in is a reference, but we'll recognize it as a keyword here instead
+		// to make the result less confusing.
+		return true
+
 	case after.Type == hclsyntax.TokenOBrack && (subject.Type == hclsyntax.TokenIdent || subject.Type == hclsyntax.TokenNumberLit || tokenBracketChange(subject) < 0):
 		return false
 
@@ -283,7 +294,7 @@ func spaceAfterToken(subject, before, after *Token) bool {
 			return true
 		}
 
-	case subject.Type == hclsyntax.TokenOBrace || (after != nil && after.Type == hclsyntax.TokenCBrace):
+	case subject.Type == hclsyntax.TokenOBrace || after.Type == hclsyntax.TokenCBrace:
 		// Unlike other bracket types, braces have spaces on both sides of them,
 		// both in single-line nested blocks foo { bar = baz } and in object
 		// constructor expressions foo = { bar = baz }.
@@ -294,6 +305,10 @@ func spaceAfterToken(subject, before, after *Token) bool {
 		}
 		return true
 
+	case after.Type == hclsyntax.TokenColon:
+		// Never spaces before colons
+		return false
+
 	case tokenBracketChange(subject) > 0:
 		// No spaces after open brackets
 		return false
diff --git a/hclwrite/format_test.go b/hclwrite/format_test.go
index 50fee30..1c06682 100644
--- a/hclwrite/format_test.go
+++ b/hclwrite/format_test.go
@@ -112,6 +112,14 @@ foo(
 			`[[]]`,
 		},
 		{
+			`[for x in y: x]`,
+			`[for x in y: x]`,
+		},
+		{
+			`[for x in [y]: x]`,
+			`[for x in [y]: x]`,
+		},
+		{
 			`
 [
 [
diff --git a/hclwrite/tokens.go b/hclwrite/tokens.go
index d37b1c6..d87f818 100644
--- a/hclwrite/tokens.go
+++ b/hclwrite/tokens.go
@@ -5,6 +5,7 @@ import (
 	"io"
 
 	"github.com/apparentlymart/go-textseg/textseg"
+	"github.com/hashicorp/hcl2/hcl"
 	"github.com/hashicorp/hcl2/hcl/hclsyntax"
 )
 
@@ -22,6 +23,23 @@ type Token struct {
 	SpacesBefore int
 }
 
+// asHCLSyntax returns the receiver expressed as an incomplete hclsyntax.Token.
+// A complete token is not possible since we don't have source location
+// information here, and so this method is unexported so we can be sure it will
+// only be used for internal purposes where we know the range isn't important.
+//
+// This is primarily intended to allow us to re-use certain functionality from
+// hclsyntax rather than re-implementing it against our own token type here.
+func (t *Token) asHCLSyntax() hclsyntax.Token {
+	return hclsyntax.Token{
+		Type:  t.Type,
+		Bytes: t.Bytes,
+		Range: hcl.Range{
+			Filename: "<invalid>",
+		},
+	}
+}
+
 // Tokens is a flat list of tokens.
 type Tokens []*Token