mirror of
https://github.com/open-policy-agent/opa.git
synced 2026-08-14 20:32:27 -06:00
f638d624b1
The formatter was not starting a new line after writing comments that preceed an expression. As a result the first part of the expression written (e.g., "not", function name, "some", etc.) would be written without any indentation. Fixes #1560 Signed-off-by: Torin Sandall <torinsandall@gmail.com>
982 lines
23 KiB
Go
982 lines
23 KiB
Go
// Copyright 2017 The OPA Authors. All rights reserved.
|
|
// Use of this source code is governed by an Apache2
|
|
// license that can be found in the LICENSE file.
|
|
|
|
// Package format implements formatting of Rego source files.
|
|
package format
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"sort"
|
|
|
|
"github.com/open-policy-agent/opa/ast"
|
|
)
|
|
|
|
// Bytes formats Rego source code. The bytes provided do not have to be an entire
|
|
// source file, but they must be parse-able. If the bytes are not parse-able, Bytes
|
|
// will return an error resulting from the attempt to parse them.
|
|
func Bytes(src []byte) ([]byte, error) {
|
|
astElem, err := ast.Parse("", src, ast.CommentsOption())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return Ast(astElem)
|
|
}
|
|
|
|
// Source formats a Rego source file. The bytes provided must describe a complete
|
|
// Rego module. If they don't, Source will return an error resulting from the attempt
|
|
// to parse the bytes.
|
|
func Source(filename string, src []byte) ([]byte, error) {
|
|
module, err := ast.ParseModule(filename, string(src))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
formatted, err := Ast(module)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("%s: %v", filename, err)
|
|
}
|
|
return formatted, nil
|
|
}
|
|
|
|
// MustAst is a helper function to format a Rego AST element. If any errors
|
|
// occurs this function will panic. This is mostly used for test
|
|
func MustAst(x interface{}) []byte {
|
|
bs, err := Ast(x)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return bs
|
|
}
|
|
|
|
// Ast formats a Rego AST element. If the passed value is not a valid AST
|
|
// element, Ast returns nil and an error. Ast relies on all AST elements having
|
|
// non-nil Location values. If an AST element with a nil Location value is
|
|
// encountered, a default location will be set on the AST node.
|
|
func Ast(x interface{}) (formatted []byte, err error) {
|
|
|
|
ast.WalkNodes(x, func(x ast.Node) bool {
|
|
if b, ok := x.(ast.Body); ok {
|
|
if len(b) == 0 {
|
|
return false
|
|
}
|
|
}
|
|
if x.Loc() == nil {
|
|
x.SetLoc(defaultLocation(x))
|
|
}
|
|
return false
|
|
})
|
|
|
|
w := &writer{indent: "\t"}
|
|
switch x := x.(type) {
|
|
case *ast.Module:
|
|
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, nil)
|
|
case *ast.Head:
|
|
w.writeHead(x, false, false, nil)
|
|
case ast.Body:
|
|
w.writeBody(x, nil)
|
|
case *ast.Expr:
|
|
w.writeExpr(x, nil)
|
|
case *ast.With:
|
|
w.writeWith(x, nil)
|
|
case *ast.Term:
|
|
w.writeTerm(x, nil)
|
|
case ast.Value:
|
|
w.writeTerm(&ast.Term{Value: x, Location: &ast.Location{}}, nil)
|
|
case *ast.Comment:
|
|
w.writeComments([]*ast.Comment{x})
|
|
default:
|
|
return nil, fmt.Errorf("not an ast element: %v", x)
|
|
}
|
|
|
|
return squashTrailingNewlines(w.buf.Bytes()), nil
|
|
}
|
|
|
|
func squashTrailingNewlines(bs []byte) []byte {
|
|
if bytes.HasSuffix(bs, []byte("\n")) {
|
|
return append(bytes.TrimRight(bs, "\n"), '\n')
|
|
}
|
|
return bs
|
|
}
|
|
|
|
func defaultLocation(x ast.Node) *ast.Location {
|
|
return ast.NewLocation([]byte(x.String()), "", 1, 1)
|
|
}
|
|
|
|
type writer struct {
|
|
buf bytes.Buffer
|
|
|
|
indent string
|
|
level int
|
|
inline bool
|
|
beforeEnd *ast.Comment
|
|
delay bool
|
|
}
|
|
|
|
func (w *writer) writeModule(module *ast.Module) {
|
|
var pkg *ast.Package
|
|
var others []interface{}
|
|
var comments []*ast.Comment
|
|
visitor := ast.NewGenericVisitor(func(x interface{}) bool {
|
|
switch x := x.(type) {
|
|
case *ast.Comment:
|
|
comments = append(comments, x)
|
|
return true
|
|
case *ast.Import, *ast.Rule:
|
|
others = append(others, x)
|
|
return true
|
|
case *ast.Package:
|
|
pkg = x
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
})
|
|
ast.Walk(visitor, module)
|
|
|
|
sort.Slice(comments, func(i, j int) bool {
|
|
return locLess(comments[i], comments[j])
|
|
})
|
|
|
|
// XXX: The parser currently duplicates comments for some reason, so we need
|
|
// to remove duplicates here.
|
|
comments = dedupComments(comments)
|
|
sort.Slice(others, func(i, j int) bool {
|
|
return locLess(others[i], others[j])
|
|
})
|
|
|
|
comments = w.writePackage(pkg, comments)
|
|
var imports []*ast.Import
|
|
var rules []*ast.Rule
|
|
for len(others) > 0 {
|
|
imports, others = gatherImports(others)
|
|
comments = w.writeImports(imports, comments)
|
|
rules, others = gatherRules(others)
|
|
comments = w.writeRules(rules, comments)
|
|
}
|
|
|
|
for i, c := range comments {
|
|
w.writeLine(c.String())
|
|
if i == len(comments)-1 {
|
|
w.write("\n")
|
|
}
|
|
}
|
|
}
|
|
|
|
func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, pkg.Location)
|
|
|
|
w.startLine()
|
|
w.write(pkg.String())
|
|
w.blankLine()
|
|
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeComments(comments []*ast.Comment) {
|
|
for i := 0; i < len(comments); i++ {
|
|
if i > 0 && locCmp(comments[i], comments[i-1]) > 1 {
|
|
w.blankLine()
|
|
}
|
|
w.writeLine(comments[i].String())
|
|
}
|
|
}
|
|
|
|
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, comments)
|
|
w.blankLine()
|
|
}
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
|
|
if rule == nil {
|
|
return comments
|
|
}
|
|
|
|
w.startLine()
|
|
if rule.Default {
|
|
w.write("default ")
|
|
}
|
|
|
|
// OPA transforms lone bodies like `foo = {"a": "b"}` into rules of the form
|
|
// `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)))) && rule.Else == nil
|
|
|
|
comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
|
|
|
|
if (len(rule.Body) == 0 || isExpandedConst) && !isElse {
|
|
w.endLine()
|
|
return comments
|
|
}
|
|
|
|
w.write(" {")
|
|
w.endLine()
|
|
w.up()
|
|
|
|
comments = w.writeBody(rule.Body, comments)
|
|
|
|
var close *ast.Location
|
|
|
|
if len(rule.Head.Args) > 0 {
|
|
close = closingLoc('(', ')', '{', '}', rule.Location)
|
|
} else {
|
|
close = closingLoc('[', ']', '{', '}', rule.Location)
|
|
}
|
|
|
|
comments = w.insertComments(comments, close)
|
|
|
|
w.down()
|
|
w.startLine()
|
|
w.write("}")
|
|
if rule.Else != nil {
|
|
w.blankLine()
|
|
rule.Else.Head.Name = ast.Var("else")
|
|
rule.Else.Head.Args = nil
|
|
comments = w.insertComments(comments, rule.Else.Head.Location)
|
|
comments = w.writeRule(rule.Else, true, comments)
|
|
}
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeHead(head *ast.Head, isDefault bool, isExpandedConst bool, comments []*ast.Comment) []*ast.Comment {
|
|
w.write(head.Name.String())
|
|
if len(head.Args) > 0 {
|
|
w.write("(")
|
|
var args []interface{}
|
|
for _, arg := range head.Args {
|
|
args = append(args, arg)
|
|
}
|
|
comments = w.writeIterable(args, head.Location, closingLoc(0, 0, '(', ')', head.Location), comments, w.listWriter())
|
|
w.write(")")
|
|
}
|
|
if head.Key != nil {
|
|
w.write("[")
|
|
comments = w.writeTerm(head.Key, comments)
|
|
w.write("]")
|
|
}
|
|
if head.Value != nil && (head.Key != nil || ast.Compare(head.Value, ast.BooleanTerm(true)) != 0 || isExpandedConst || isDefault) {
|
|
if head.Assign {
|
|
w.write(" := ")
|
|
} else {
|
|
w.write(" = ")
|
|
}
|
|
comments = w.writeTerm(head.Value, comments)
|
|
}
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) insertComments(comments []*ast.Comment, loc *ast.Location) []*ast.Comment {
|
|
before, at, comments := partitionComments(comments, loc)
|
|
w.writeComments(before)
|
|
if len(before) > 0 && loc.Row-before[len(before)-1].Location.Row > 1 {
|
|
w.blankLine()
|
|
}
|
|
|
|
w.beforeLineEnd(at)
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeBody(body ast.Body, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, body.Loc())
|
|
offset := 0
|
|
for i, expr := range body {
|
|
if i > 0 && expr.Location.Row-body[i-1].Location.Row-offset > 1 {
|
|
w.blankLine()
|
|
}
|
|
w.startLine()
|
|
|
|
comments = w.writeExpr(expr, comments)
|
|
w.endLine()
|
|
}
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, expr.Location)
|
|
if !w.inline {
|
|
w.startLine()
|
|
}
|
|
|
|
if expr.Negated {
|
|
w.write("not ")
|
|
}
|
|
|
|
switch t := expr.Terms.(type) {
|
|
case *ast.SomeDecl:
|
|
comments = w.writeSomeDecl(t, comments)
|
|
case []*ast.Term:
|
|
comments = w.writeFunctionCall(expr, comments)
|
|
case *ast.Term:
|
|
comments = w.writeTerm(t, comments)
|
|
}
|
|
|
|
var indented bool
|
|
for i, with := range expr.With {
|
|
if i > 0 && with.Location.Row-expr.With[i-1].Location.Row > 0 {
|
|
if !indented {
|
|
indented = true
|
|
|
|
w.up()
|
|
defer w.down()
|
|
}
|
|
w.endLine()
|
|
w.startLine()
|
|
}
|
|
comments = w.writeWith(with, comments)
|
|
}
|
|
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeSomeDecl(decl *ast.SomeDecl, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, decl.Location)
|
|
w.write("some ")
|
|
|
|
row := decl.Location.Row
|
|
|
|
for i, term := range decl.Symbols {
|
|
|
|
if term.Location.Row > row {
|
|
w.endLine()
|
|
w.startLine()
|
|
w.write(w.indent)
|
|
row = term.Location.Row
|
|
} else if i > 0 {
|
|
w.write(" ")
|
|
}
|
|
|
|
comments = w.writeTerm(term, comments)
|
|
|
|
if i < len(decl.Symbols)-1 {
|
|
w.write(",")
|
|
}
|
|
}
|
|
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*ast.Comment {
|
|
|
|
terms := expr.Terms.([]*ast.Term)
|
|
|
|
bi, ok := ast.BuiltinMap[terms[0].Value.String()]
|
|
if !ok || bi.Infix == "" {
|
|
return w.writeFunctionCallPlain(terms, comments)
|
|
}
|
|
|
|
numDeclArgs := len(bi.Decl.Args())
|
|
numCallArgs := len(terms) - 1
|
|
|
|
if numCallArgs == numDeclArgs {
|
|
// Print infix where result is unassigned (e.g., x != y)
|
|
comments = w.writeTerm(terms[1], comments)
|
|
w.write(" " + string(bi.Infix) + " ")
|
|
return w.writeTerm(terms[2], comments)
|
|
} else if numCallArgs == numDeclArgs+1 {
|
|
// Print infix where result is assigned (e.g., z = x + y)
|
|
comments = w.writeTerm(terms[3], comments)
|
|
w.write(" " + ast.Equality.Infix + " ")
|
|
comments = w.writeTerm(terms[1], comments)
|
|
w.write(" " + bi.Infix + " ")
|
|
comments = w.writeTerm(terms[2], comments)
|
|
return comments
|
|
}
|
|
|
|
return w.writeFunctionCallPlain(terms, comments)
|
|
}
|
|
|
|
func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comment) []*ast.Comment {
|
|
w.write(string(terms[0].String()) + "(")
|
|
if len(terms) > 1 {
|
|
for _, v := range terms[1 : len(terms)-1] {
|
|
comments = w.writeTerm(v, comments)
|
|
w.write(", ")
|
|
}
|
|
comments = w.writeTerm(terms[len(terms)-1], comments)
|
|
}
|
|
w.write(")")
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeWith(with *ast.With, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, with.Location)
|
|
w.write(" with ")
|
|
comments = w.writeTerm(with.Target, comments)
|
|
w.write(" as ")
|
|
return w.writeTerm(with.Value, comments)
|
|
}
|
|
|
|
func (w *writer) writeTerm(term *ast.Term, comments []*ast.Comment) []*ast.Comment {
|
|
return w.writeTermParens(false, term, comments)
|
|
}
|
|
|
|
func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Comment) []*ast.Comment {
|
|
comments = w.insertComments(comments, term.Location)
|
|
if !w.inline {
|
|
w.startLine()
|
|
}
|
|
|
|
switch x := term.Value.(type) {
|
|
case ast.Ref:
|
|
w.write(x.String())
|
|
case ast.Object:
|
|
comments = w.writeObject(x, term.Location, comments)
|
|
case ast.Array:
|
|
comments = w.writeArray(x, term.Location, comments)
|
|
case ast.Set:
|
|
comments = w.writeSet(x, term.Location, comments)
|
|
case *ast.ArrayComprehension:
|
|
comments = w.writeArrayComprehension(x, term.Location, comments)
|
|
case *ast.ObjectComprehension:
|
|
comments = w.writeObjectComprehension(x, term.Location, comments)
|
|
case *ast.SetComprehension:
|
|
comments = w.writeSetComprehension(x, term.Location, comments)
|
|
case ast.String:
|
|
if term.Location.Text[0] == '.' {
|
|
// This string was parsed from a ref, so preserve the value.
|
|
w.write(`"` + string(x) + `"`)
|
|
} else {
|
|
// To preserve raw strings, we need to output the original text,
|
|
// not what x.String() would give us.
|
|
w.write(string(term.Location.Text))
|
|
}
|
|
case ast.Call:
|
|
comments = w.writeCall(parens, x, term.Location, comments)
|
|
case fmt.Stringer:
|
|
w.write(x.String())
|
|
}
|
|
|
|
if !w.inline {
|
|
w.startLine()
|
|
}
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeCall(parens bool, x ast.Call, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
|
|
bi, ok := ast.BuiltinMap[x[0].String()]
|
|
if !ok || bi.Infix == "" {
|
|
return w.writeFunctionCallPlain([]*ast.Term(x), comments)
|
|
}
|
|
|
|
// TODO(tsandall): improve to consider precedence?
|
|
if parens {
|
|
w.write("(")
|
|
}
|
|
comments = w.writeTermParens(true, x[1], comments)
|
|
w.write(" " + bi.Infix + " ")
|
|
comments = w.writeTermParens(true, x[2], comments)
|
|
if parens {
|
|
w.write(")")
|
|
}
|
|
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeObject(obj ast.Object, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
w.write("{")
|
|
defer w.write("}")
|
|
|
|
var s []interface{}
|
|
obj.Foreach(func(k, v *ast.Term) {
|
|
s = append(s, ast.Item(k, v))
|
|
})
|
|
return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.objectWriter())
|
|
}
|
|
|
|
func (w *writer) writeArray(arr ast.Array, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
w.write("[")
|
|
defer w.write("]")
|
|
|
|
var s []interface{}
|
|
for _, t := range arr {
|
|
s = append(s, t)
|
|
}
|
|
return w.writeIterable(s, loc, closingLoc(0, 0, '[', ']', loc), comments, w.listWriter())
|
|
}
|
|
|
|
func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
|
|
if set.Len() == 0 {
|
|
w.write("set()")
|
|
return w.insertComments(comments, closingLoc(0, 0, '(', ')', loc))
|
|
}
|
|
|
|
w.write("{")
|
|
defer w.write("}")
|
|
|
|
var s []interface{}
|
|
set.Foreach(func(t *ast.Term) {
|
|
s = append(s, t)
|
|
})
|
|
return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.listWriter())
|
|
}
|
|
|
|
func (w *writer) writeArrayComprehension(arr *ast.ArrayComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
w.write("[")
|
|
defer w.write("]")
|
|
|
|
return w.writeComprehension('[', ']', arr.Term, arr.Body, loc, comments)
|
|
}
|
|
|
|
func (w *writer) writeSetComprehension(set *ast.SetComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
w.write("{")
|
|
defer w.write("}")
|
|
|
|
return w.writeComprehension('{', '}', set.Term, set.Body, loc, comments)
|
|
}
|
|
|
|
func (w *writer) writeObjectComprehension(object *ast.ObjectComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
w.write("{")
|
|
defer w.write("}")
|
|
|
|
object.Value.Location = object.Key.Location // Ensure the value is not written on the next line.
|
|
if object.Key.Location.Row-loc.Row > 1 {
|
|
w.endLine()
|
|
w.startLine()
|
|
}
|
|
|
|
comments = w.writeTerm(object.Key, comments)
|
|
w.write(": ")
|
|
return w.writeComprehension('{', '}', object.Value, object.Body, loc, comments)
|
|
}
|
|
|
|
func (w *writer) writeComprehension(open, close byte, term *ast.Term, body ast.Body, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
if term.Location.Row-loc.Row > 1 {
|
|
w.endLine()
|
|
w.startLine()
|
|
}
|
|
|
|
comments = w.writeTerm(term, comments)
|
|
w.write(" |")
|
|
|
|
return w.writeComprehensionBody(open, close, body, term.Location, loc, comments)
|
|
}
|
|
|
|
func (w *writer) writeComprehensionBody(open, close byte, body ast.Body, term, compr *ast.Location, comments []*ast.Comment) []*ast.Comment {
|
|
var exprs []interface{}
|
|
for _, expr := range body {
|
|
exprs = append(exprs, expr)
|
|
}
|
|
lines := groupIterable(exprs, term)
|
|
|
|
if body.Loc().Row-term.Row > 0 || len(lines) > 1 {
|
|
w.endLine()
|
|
w.up()
|
|
defer w.startLine()
|
|
defer w.down()
|
|
|
|
comments = w.writeBody(body, comments)
|
|
} else {
|
|
w.write(" ")
|
|
i := 0
|
|
for ; i < len(body)-1; i++ {
|
|
comments = w.writeExpr(body[i], comments)
|
|
w.write("; ")
|
|
}
|
|
comments = w.writeExpr(body[i], comments)
|
|
}
|
|
|
|
return w.insertComments(comments, closingLoc(0, 0, open, close, compr))
|
|
}
|
|
|
|
func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []*ast.Comment {
|
|
m, comments := mapImportsToComments(imports, comments)
|
|
|
|
groups := groupImports(imports)
|
|
for _, group := range groups {
|
|
comments = w.insertComments(comments, group[0].Loc())
|
|
|
|
// Sort imports within a newline grouping.
|
|
sort.Slice(group, func(i, j int) bool {
|
|
a := group[i]
|
|
b := group[j]
|
|
return a.Compare(b) < 0
|
|
})
|
|
for _, i := range group {
|
|
w.startLine()
|
|
w.write(i.String())
|
|
if c, ok := m[i]; ok {
|
|
w.write(" " + c.String())
|
|
}
|
|
w.endLine()
|
|
}
|
|
w.blankLine()
|
|
}
|
|
|
|
return comments
|
|
}
|
|
|
|
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 {
|
|
lines := groupIterable(elements, last)
|
|
if len(lines) > 1 {
|
|
w.delayBeforeEnd()
|
|
w.startMultilineSeq()
|
|
}
|
|
|
|
i := 0
|
|
for ; i < len(lines)-1; i++ {
|
|
comments = w.writeIterableLine(lines[i], comments, fn)
|
|
w.write(",")
|
|
|
|
w.endLine()
|
|
w.startLine()
|
|
}
|
|
|
|
comments = w.writeIterableLine(lines[i], comments, fn)
|
|
|
|
if len(lines) > 1 {
|
|
w.write(",")
|
|
w.endLine()
|
|
comments = w.insertComments(comments, close)
|
|
w.down()
|
|
w.startLine()
|
|
}
|
|
|
|
return comments
|
|
}
|
|
|
|
func (w *writer) writeIterableLine(elements []interface{}, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
|
|
if len(elements) == 0 {
|
|
return comments
|
|
}
|
|
|
|
i := 0
|
|
for ; i < len(elements)-1; i++ {
|
|
comments = fn(elements[i], comments)
|
|
w.write(", ")
|
|
}
|
|
|
|
return fn(elements[i], comments)
|
|
}
|
|
|
|
func (w *writer) objectWriter() entryWriter {
|
|
return func(x interface{}, comments []*ast.Comment) []*ast.Comment {
|
|
entry := x.([2]*ast.Term)
|
|
comments = w.writeTerm(entry[0], comments)
|
|
w.write(": ")
|
|
return w.writeTerm(entry[1], comments)
|
|
}
|
|
}
|
|
|
|
func (w *writer) listWriter() entryWriter {
|
|
return func(x interface{}, comments []*ast.Comment) []*ast.Comment {
|
|
return w.writeTerm(x.(*ast.Term), comments)
|
|
}
|
|
}
|
|
|
|
func groupIterable(elements []interface{}, last *ast.Location) (lines [][]interface{}) {
|
|
var cur []interface{}
|
|
for i, t := range elements {
|
|
loc := getLoc(t)
|
|
lineDiff := loc.Row - last.Row
|
|
if lineDiff > 0 && i > 0 {
|
|
lines = append(lines, cur)
|
|
cur = nil
|
|
}
|
|
|
|
last = loc
|
|
cur = append(cur, t)
|
|
}
|
|
return append(lines, cur)
|
|
}
|
|
|
|
func mapImportsToComments(imports []*ast.Import, comments []*ast.Comment) (map[*ast.Import]*ast.Comment, []*ast.Comment) {
|
|
var leftovers []*ast.Comment
|
|
m := map[*ast.Import]*ast.Comment{}
|
|
|
|
for _, c := range comments {
|
|
matched := false
|
|
for _, i := range imports {
|
|
if c.Loc().Row == i.Loc().Row {
|
|
m[i] = c
|
|
matched = true
|
|
break
|
|
}
|
|
}
|
|
if !matched {
|
|
leftovers = append(leftovers, c)
|
|
}
|
|
}
|
|
|
|
return m, leftovers
|
|
}
|
|
|
|
func groupImports(imports []*ast.Import) (groups [][]*ast.Import) {
|
|
if len(imports) == 0 {
|
|
return nil
|
|
}
|
|
|
|
last := imports[0]
|
|
var group []*ast.Import
|
|
for _, i := range imports {
|
|
if i.Loc().Row-last.Loc().Row > 1 {
|
|
groups = append(groups, group)
|
|
group = []*ast.Import{}
|
|
}
|
|
group = append(group, i)
|
|
last = i
|
|
}
|
|
if len(group) > 0 {
|
|
groups = append(groups, group)
|
|
}
|
|
|
|
return groups
|
|
}
|
|
|
|
func partitionComments(comments []*ast.Comment, l *ast.Location) (before []*ast.Comment, at *ast.Comment, after []*ast.Comment) {
|
|
for _, c := range comments {
|
|
switch cmp := c.Location.Row - l.Row; {
|
|
case cmp < 0:
|
|
before = append(before, c)
|
|
case cmp > 0:
|
|
after = append(after, c)
|
|
case cmp == 0:
|
|
at = c
|
|
}
|
|
}
|
|
|
|
return before, at, after
|
|
}
|
|
|
|
func gatherImports(others []interface{}) (imports []*ast.Import, rest []interface{}) {
|
|
i := 0
|
|
loop:
|
|
for ; i < len(others); i++ {
|
|
switch x := others[i].(type) {
|
|
case *ast.Import:
|
|
imports = append(imports, x)
|
|
case *ast.Rule:
|
|
break loop
|
|
}
|
|
}
|
|
return imports, others[i:]
|
|
}
|
|
|
|
func gatherRules(others []interface{}) (rules []*ast.Rule, rest []interface{}) {
|
|
i := 0
|
|
loop:
|
|
for ; i < len(others); i++ {
|
|
switch x := others[i].(type) {
|
|
case *ast.Rule:
|
|
rules = append(rules, x)
|
|
case *ast.Import:
|
|
break loop
|
|
}
|
|
}
|
|
return rules, others[i:]
|
|
}
|
|
|
|
func locLess(a, b interface{}) bool {
|
|
return locCmp(a, b) < 0
|
|
}
|
|
|
|
func locCmp(a, b interface{}) int {
|
|
al := getLoc(a)
|
|
bl := getLoc(b)
|
|
if cmp := al.Row - bl.Row; cmp != 0 {
|
|
return cmp
|
|
}
|
|
return al.Col - bl.Col
|
|
}
|
|
|
|
func getLoc(x interface{}) *ast.Location {
|
|
switch x := x.(type) {
|
|
case ast.Statement:
|
|
return x.Loc()
|
|
case *ast.Head:
|
|
return x.Location
|
|
case *ast.Expr:
|
|
return x.Location
|
|
case *ast.With:
|
|
return x.Location
|
|
case *ast.Term:
|
|
return x.Location
|
|
case *ast.Location:
|
|
return x
|
|
case [2]*ast.Term:
|
|
// Special case to allow for easy printing of objects.
|
|
return x[0].Location
|
|
default:
|
|
panic("Not reached")
|
|
}
|
|
}
|
|
|
|
func closingLoc(skipOpen, skipClose, open, close byte, loc *ast.Location) *ast.Location {
|
|
i, offset := 0, 0
|
|
|
|
// Skip past parens/brackets/braces in rule heads.
|
|
if skipOpen > 0 {
|
|
i, offset = skipPast(skipOpen, skipClose, loc)
|
|
}
|
|
|
|
for ; i < len(loc.Text) && loc.Text[i] != open; i++ {
|
|
}
|
|
|
|
if i >= len(loc.Text) {
|
|
return &ast.Location{Row: -1}
|
|
}
|
|
|
|
state := 1
|
|
for state > 0 {
|
|
i++
|
|
if i >= len(loc.Text) {
|
|
return &ast.Location{Row: -1}
|
|
}
|
|
|
|
switch loc.Text[i] {
|
|
case open:
|
|
state++
|
|
case close:
|
|
state--
|
|
case '\n':
|
|
offset++
|
|
}
|
|
}
|
|
|
|
return &ast.Location{Row: loc.Row + offset}
|
|
}
|
|
|
|
func skipPast(open, close byte, loc *ast.Location) (int, int) {
|
|
i := 0
|
|
for ; i < len(loc.Text) && loc.Text[i] != open; i++ {
|
|
}
|
|
|
|
state := 1
|
|
offset := 0
|
|
for state > 0 {
|
|
i++
|
|
if i >= len(loc.Text) {
|
|
return i, offset
|
|
}
|
|
|
|
switch loc.Text[i] {
|
|
case open:
|
|
state++
|
|
case close:
|
|
state--
|
|
case '\n':
|
|
offset++
|
|
}
|
|
}
|
|
|
|
return i, offset
|
|
}
|
|
|
|
func dedupComments(comments []*ast.Comment) []*ast.Comment {
|
|
if len(comments) == 0 {
|
|
return nil
|
|
}
|
|
|
|
filtered := []*ast.Comment{comments[0]}
|
|
for i := 1; i < len(comments); i++ {
|
|
if comments[i].Location.Equal(comments[i-1].Location) {
|
|
continue
|
|
}
|
|
filtered = append(filtered, comments[i])
|
|
}
|
|
return filtered
|
|
}
|
|
|
|
// startLine begins a line with the current indentation level.
|
|
func (w *writer) startLine() {
|
|
if w.inline {
|
|
panic("currently in a line")
|
|
}
|
|
w.inline = true
|
|
for i := 0; i < w.level; i++ {
|
|
w.write(w.indent)
|
|
}
|
|
}
|
|
|
|
// endLine ends a line with a newline.
|
|
func (w *writer) endLine() {
|
|
if !w.inline {
|
|
panic("not in a line")
|
|
}
|
|
w.inline = false
|
|
if w.beforeEnd != nil && !w.delay {
|
|
w.write(" " + w.beforeEnd.String())
|
|
w.beforeEnd = nil
|
|
}
|
|
w.delay = false
|
|
w.write("\n")
|
|
}
|
|
|
|
// beforeLineEnd registers a comment to be printed at the end of the current line.
|
|
func (w *writer) beforeLineEnd(c *ast.Comment) {
|
|
if w.beforeEnd != nil {
|
|
if c == nil {
|
|
return
|
|
}
|
|
panic("overwriting non-nil beforeEnd")
|
|
}
|
|
w.beforeEnd = c
|
|
}
|
|
|
|
func (w *writer) delayBeforeEnd() {
|
|
w.delay = true
|
|
}
|
|
|
|
// line prints a blank line. If the writer is currently in the middle of a line,
|
|
// line ends it and then prints a blank one.
|
|
func (w *writer) blankLine() {
|
|
if w.inline {
|
|
w.endLine()
|
|
}
|
|
w.write("\n")
|
|
}
|
|
|
|
// write the input string and writes it to the buffer.
|
|
func (w *writer) write(s string) {
|
|
w.buf.WriteString(s)
|
|
}
|
|
|
|
// writeLine writes the string on a newly started line, then terminate the line.
|
|
func (w *writer) writeLine(s string) {
|
|
if !w.inline {
|
|
w.startLine()
|
|
}
|
|
w.write(s)
|
|
w.endLine()
|
|
}
|
|
|
|
func (w *writer) startMultilineSeq() {
|
|
w.endLine()
|
|
w.up()
|
|
w.startLine()
|
|
}
|
|
|
|
func (w *writer) endMultilineSeq() {
|
|
w.write(",")
|
|
w.endLine()
|
|
w.down()
|
|
w.startLine()
|
|
}
|
|
|
|
// up increases the indentation level
|
|
func (w *writer) up() {
|
|
w.level++
|
|
}
|
|
|
|
// down decreases the indentation level
|
|
func (w *writer) down() {
|
|
if w.level == 0 {
|
|
panic("negative indentation level")
|
|
}
|
|
w.level--
|
|
}
|