format: Bracketing keyword ref elements in formatter output

ast/policy.go

@@ -100,7 +100,7 @@ var Wildcard = &Term{Value: Var("_")}
 var WildcardPrefix = "$"
 
 // Keywords contains strings that map to language keywords.
-var Keywords = KeywordsV0
+var Keywords = KeywordsForRegoVersion(DefaultRegoVersion)
 
 var KeywordsV0 = [...]string{
 	"not",
@@ -134,9 +134,23 @@ var KeywordsV1 = [...]string{
 	"every",
 }
 
+func KeywordsForRegoVersion(v RegoVersion) []string {
+	switch v {
+	case RegoV0:
+		return KeywordsV0[:]
+	case RegoV1, RegoV0CompatV1:
+		return KeywordsV1[:]
+	}
+	return nil
+}
+
 // IsKeyword returns true if s is a language keyword.
 func IsKeyword(s string) bool {
-	for _, x := range Keywords {
+	return IsInKeywords(s, Keywords)
+}
+
+func IsInKeywords(s string, keywords []string) bool {
+	for _, x := range keywords {
 		if x == s {
 			return true
 		}

compile/compile_test.go

@@ -1741,7 +1741,7 @@ contains := 2 {
 import rego.v1
 
 p if {
-	data.foo.contains = input.x
+	data.foo["contains"] = input.x
 }
 `,
 				`package foo

format/format.go

@@ -126,7 +126,16 @@ type fmtOpts struct {
 	// than if they don't.
 	refHeads bool
 
-	regoV1 bool
+	regoV1         bool
+	futureKeywords []string
+}
+
+func (o fmtOpts) keywords() []string {
+	if o.regoV1 {
+		return ast.KeywordsV1[:]
+	}
+	kws := ast.KeywordsV0[:]
+	return append(kws, o.futureKeywords...)
 }
 
 func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
@@ -171,6 +180,10 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 			}
 
 		case *ast.Import:
+			if kw, ok := future.WhichFutureKeyword(n); ok {
+				o.futureKeywords = append(o.futureKeywords, kw)
+			}
+
 			switch {
 			case isRegoV1Compatible(n):
 				o.contains = true
@@ -200,8 +213,9 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 	})
 
 	w := &writer{
-		indent: "\t",
-		errs:   make([]*ast.Error, 0),
+		indent:  "\t",
+		errs:    make([]*ast.Error, 0),
+		fmtOpts: o,
 	}
 
 	switch x := x.(type) {
@@ -219,18 +233,17 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 				x.Imports = ensureFutureKeywordImport(x.Imports, kw)
 			}
 		}
-		w.writeModule(x, o)
+		w.writeModule(x)
 	case *ast.Package:
 		w.writePackage(x, nil)
 	case *ast.Import:
 		w.writeImports([]*ast.Import{x}, nil)
 	case *ast.Rule:
-		w.writeRule(x, false /* isElse */, o, nil)
+		w.writeRule(x, false /* isElse */, nil)
 	case *ast.Head:
 		w.writeHead(x,
 			false, // isDefault
 			false, // isExpandedConst
-			o,
 			nil)
 	case ast.Body:
 		w.writeBody(x, nil)
@@ -302,9 +315,10 @@ type writer struct {
 	beforeEnd *ast.Comment
 	delay     bool
 	errs      ast.Errors
+	fmtOpts   fmtOpts
 }
 
-func (w *writer) writeModule(module *ast.Module, o fmtOpts) {
+func (w *writer) writeModule(module *ast.Module) {
 	var pkg *ast.Package
 	var others []interface{}
 	var comments []*ast.Comment
@@ -342,7 +356,7 @@ func (w *writer) writeModule(module *ast.Module, o fmtOpts) {
 		imports, others = gatherImports(others)
 		comments = w.writeImports(imports, comments)
 		rules, others = gatherRules(others)
-		comments = w.writeRules(rules, o, comments)
+		comments = w.writeRules(rules, comments)
 	}
 
 	for i, c := range comments {
@@ -365,7 +379,15 @@ func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) []*ast.
 	comments = w.insertComments(comments, pkg.Location)
 
 	w.startLine()
-	w.write(pkg.String())
+
+	// Omit head as all packages have the DefaultRootDocument prepended at parse time.
+	path := make(ast.Ref, len(pkg.Path)-1)
+	path[0] = ast.VarTerm(string(pkg.Path[1].Value.(ast.String)))
+	copy(path[1:], pkg.Path[2:])
+
+	w.write("package ")
+	w.writeRef(path)
+
 	w.blankLine()
 
 	return comments
@@ -380,16 +402,16 @@ func (w *writer) writeComments(comments []*ast.Comment) {
 	}
 }
 
-func (w *writer) writeRules(rules []*ast.Rule, o fmtOpts, comments []*ast.Comment) []*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, false, o, comments)
+		comments = w.writeRule(rule, false, comments)
 		w.blankLine()
 	}
 	return comments
 }
 
-func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
 	if rule == nil {
 		return comments
 	}
@@ -408,25 +430,25 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*a
 	// pretend that the rule has no body in this case.
 	isExpandedConst := rule.Body.Equal(ast.NewBody(ast.NewExpr(ast.BooleanTerm(true)))) && rule.Else == nil
 
-	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, o, comments)
+	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
 
 	// this excludes partial sets UNLESS `contains` is used
-	partialSetException := o.contains || rule.Head.Value != nil
+	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
 
 	if len(rule.Body) == 0 || isExpandedConst {
 		w.endLine()
 		return comments
 	}
 
-	if (o.regoV1 || o.ifs) && partialSetException {
+	if (w.fmtOpts.regoV1 || w.fmtOpts.ifs) && partialSetException {
 		w.write(" if")
 		if len(rule.Body) == 1 {
 			if rule.Body[0].Location.Row == rule.Head.Location.Row {
 				w.write(" ")
 				comments = w.writeExpr(rule.Body[0], comments)
 				w.endLine()
 				if rule.Else != nil {
-					comments = w.writeElse(rule, o, comments)
+					comments = w.writeElse(rule, comments)
 				}
 				return comments
 			}
@@ -454,12 +476,12 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, o fmtOpts, comments []*a
 	w.startLine()
 	w.write("}")
 	if rule.Else != nil {
-		comments = w.writeElse(rule, o, comments)
+		comments = w.writeElse(rule, comments)
 	}
 	return comments
 }
 
-func (w *writer) writeElse(rule *ast.Rule, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) []*ast.Comment {
 	// If there was nothing else on the line before the "else" starts
 	// then preserve this style of else block, otherwise it will be
 	// started as an "inline" else eg:
@@ -521,16 +543,16 @@ func (w *writer) writeElse(rule *ast.Rule, o fmtOpts, comments []*ast.Comment) [
 		rule.Else.Head.Value.Location = rule.Else.Head.Location
 	}
 
-	return w.writeRule(rule.Else, true, o, comments)
+	return w.writeRule(rule.Else, true, comments)
 }
 
-func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fmtOpts, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comments []*ast.Comment) []*ast.Comment {
 	ref := head.Ref()
 	if head.Key != nil && head.Value == nil && !head.HasDynamicRef() {
 		ref = ref.GroundPrefix()
 	}
-	if o.refHeads || len(ref) == 1 {
-		w.write(ref.String())
+	if w.fmtOpts.refHeads || len(ref) == 1 {
+		w.writeRef(ref)
 	} else {
 		w.write(ref[0].String())
 		w.write("[")
@@ -548,7 +570,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 		w.write(")")
 	}
 	if head.Key != nil {
-		if o.contains && head.Value == nil {
+		if w.fmtOpts.contains && head.Value == nil {
 			w.write(" contains ")
 			comments = w.writeTerm(head.Key, comments)
 		} else if head.Value == nil { // no `if` for p[x] notation
@@ -566,7 +588,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 		// * a.b -> a contains "b"
 		// * a.b.c -> a.b.c := true
 		// * a.b.c.d -> a.b.c.d := true
-		isRegoV1RefConst := o.regoV1 && isExpandedConst && head.Key == nil && len(head.Args) == 0
+		isRegoV1RefConst := w.fmtOpts.regoV1 && isExpandedConst && head.Key == nil && len(head.Args) == 0
 
 		if head.Location == head.Value.Location &&
 			head.Name != "else" &&
@@ -578,7 +600,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, o fm
 			return comments
 		}
 
-		if head.Assign || o.regoV1 {
+		if head.Assign || w.fmtOpts.regoV1 {
 			// preserve assignment operator, and enforce it if formatting for Rego v1
 			w.write(" := ")
 		} else {
@@ -856,7 +878,7 @@ var varRegexp = regexp.MustCompile("^[[:alpha:]_][[:alpha:][:digit:]_]*$")
 
 func (w *writer) writeRefStringPath(s ast.String) {
 	str := string(s)
-	if varRegexp.MatchString(str) && !ast.IsKeyword(str) {
+	if varRegexp.MatchString(str) && !ast.IsInKeywords(str, w.fmtOpts.keywords()) {
 		w.write("." + str)
 	} else {
 		w.writeBracketed(s.String())
@@ -1067,7 +1089,7 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 		})
 		for _, i := range group {
 			w.startLine()
-			w.write(i.String())
+			w.writeImport(i)
 			if c, ok := m[i]; ok {
 				w.write(" " + c.String())
 			}
@@ -1079,6 +1101,28 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 	return comments
 }
 
+func (w *writer) writeImport(imp *ast.Import) {
+	path := imp.Path.Value.(ast.Ref)
+
+	buf := []string{"import"}
+
+	if _, ok := future.WhichFutureKeyword(imp); ok {
+		// We don't want to wrap future.keywords imports in parens, so we create a new writer that doesn't
+		w2 := writer{
+			buf: bytes.Buffer{},
+		}
+		w2.writeRef(path)
+		buf = append(buf, w2.buf.String())
+	} else {
+		buf = append(buf, path.String())
+	}
+
+	if len(imp.Alias) > 0 {
+		buf = append(buf, "as "+imp.Alias.String())
+	}
+	w.write(strings.Join(buf, " "))
+}
+
 type entryWriter func(interface{}, []*ast.Comment) []*ast.Comment
 
 func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close *ast.Location, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
@@ -1505,7 +1549,12 @@ func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import {
 		}
 	}
 	imp := &ast.Import{
-		Path: ast.MustParseTerm("future.keywords." + kw),
+		// NOTE: This is a hack to not error on the ref containing a keyword already present in v1.
+		// A cleaner solution would be to instead allow refs to contain keyword terms.
+		// E.g. in v1, `import future.keywords["in"]` is valid, but `import future.keywords.in` is not
+		// as it contains a reserved keyword.
+		Path: ast.MustParseTerm("future.keywords[\"" + kw + "\"]"),
+		//Path: ast.MustParseTerm("future.keywords." + kw),
 	}
 	imp.Location = defaultLocation(imp)
 	return append(imps, imp)