blob: 76d79070ded6c0423d99b2449f4e932d32d19b8f [file] [log] [blame]
// Copyright 2018 The CUE Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package astutil
import (
"encoding/hex"
"fmt"
"hash/fnv"
"path"
"reflect"
"strconv"
"strings"
"cuelang.org/go/cue/ast"
"cuelang.org/go/cue/token"
)
// A Cursor describes a node encountered during Apply.
// Information about the node and its parent is available
// from the Node, Parent, and Index methods.
//
// The methods Replace, Delete, InsertBefore, and InsertAfter
// can be used to change the AST without disrupting Apply.
// Delete, InsertBefore, and InsertAfter are only defined for modifying
// a StructLit and will panic in any other context.
type Cursor interface {
// Node returns the current Node.
Node() ast.Node
// Parent returns the parent of the current Node.
Parent() Cursor
// Index reports the index >= 0 of the current Node in the slice of Nodes
// that contains it, or a value < 0 if the current Node is not part of a
// list.
Index() int
// Import reports an opaque identifier that refers to the given package. It
// may only be called if the input to apply was an ast.File. If the import
// does not exist, it will be added.
Import(path string) *ast.Ident
// Replace replaces the current Node with n.
// The replacement node is not walked by Apply. Comments of the old node
// are copied to the new node if it has not yet an comments associated
// with it.
Replace(n ast.Node)
// Delete deletes the current Node from its containing struct.
// If the current Node is not part of a struct, Delete panics.
Delete()
// InsertAfter inserts n after the current Node in its containing struct.
// If the current Node is not part of a struct, InsertAfter panics.
// Unless n is wrapped by ApplyRecursively, Apply does not walk n.
InsertAfter(n ast.Node)
// InsertBefore inserts n before the current Node in its containing struct.
// If the current Node is not part of a struct, InsertBefore panics.
// Unless n is wrapped by ApplyRecursively, Apply does not walk n.
InsertBefore(n ast.Node)
self() *cursor
}
// ApplyRecursively indicates that a node inserted with InsertBefore,
// or InsertAfter should be processed recursively.
func ApplyRecursively(n ast.Node) ast.Node {
return recursive{n}
}
type recursive struct {
ast.Node
}
type info struct {
f *ast.File
current *declsCursor
importPatch []*ast.Ident
}
type cursor struct {
file *info
parent Cursor
node ast.Node
typ interface{} // the type of the node
index int // position of any of the sub types.
replaced bool
}
func newCursor(parent Cursor, n ast.Node, typ interface{}) *cursor {
return &cursor{
parent: parent,
typ: typ,
node: n,
index: -1,
}
}
func fileInfo(c Cursor) (info *info) {
for ; c != nil; c = c.Parent() {
if i := c.self().file; i != nil {
return i
}
}
return nil
}
func (c *cursor) self() *cursor { return c }
func (c *cursor) Parent() Cursor { return c.parent }
func (c *cursor) Index() int { return c.index }
func (c *cursor) Node() ast.Node { return c.node }
func (c *cursor) Import(importPath string) *ast.Ident {
info := fileInfo(c)
if info == nil {
return nil
}
name := path.Base(importPath)
if p := strings.LastIndexByte(name, ':'); p > 0 {
name = name[p+1:]
}
// TODO: come up with something much better.
// For instance, hoist the uniquer form cue/export.go to
// here and make export.go use this.
hash := fnv.New32()
name += hex.EncodeToString(hash.Sum([]byte(importPath)))[:6]
quoted := strconv.Quote(importPath)
var imports *ast.ImportDecl
var spec *ast.ImportSpec
decls := info.current.decls
i := 0
outer:
for ; i < len(decls); i++ {
d := decls[i]
switch t := d.(type) {
default:
break outer
case *ast.Package:
case *ast.CommentGroup:
case *ast.ImportDecl:
imports = t
for _, s := range t.Specs {
if s.Path.Value != quoted ||
s.Name == nil ||
s.Name.Name != name {
continue
}
spec = s
break
}
}
}
if spec == nil {
// Import not found, add one.
if imports == nil {
imports = &ast.ImportDecl{}
a := append(append(decls[:i], imports), decls[i:]...)
decls = a
info.current.decls = decls
}
path := &ast.BasicLit{Kind: token.STRING, Value: quoted}
spec = &ast.ImportSpec{
Name: ast.NewIdent(name),
Path: path,
}
imports.Specs = append(imports.Specs, spec)
ast.SetRelPos(imports.Specs[0], token.NoRelPos)
}
ident := &ast.Ident{Node: spec} // Name is set later.
info.importPatch = append(info.importPatch, ident)
ident.Name = name
return ident
}
func (c *cursor) Replace(n ast.Node) {
// panic if the value cannot convert to the original type.
reflect.ValueOf(n).Convert(reflect.TypeOf(c.typ).Elem())
if ast.Comments(n) != nil {
CopyComments(n, c.node)
}
if r, ok := n.(recursive); ok {
n = r.Node
} else {
c.replaced = true
}
c.node = n
}
func (c *cursor) InsertAfter(n ast.Node) { panic("unsupported") }
func (c *cursor) InsertBefore(n ast.Node) { panic("unsupported") }
func (c *cursor) Delete() { panic("unsupported") }
// Apply traverses a syntax tree recursively, starting with root,
// and calling pre and post for each node as described below.
// Apply returns the syntax tree, possibly modified.
//
// If pre is not nil, it is called for each node before the node's
// children are traversed (pre-order). If pre returns false, no
// children are traversed, and post is not called for that node.
//
// If post is not nil, and a prior call of pre didn't return false,
// post is called for each node after its children are traversed
// (post-order). If post returns false, traversal is terminated and
// Apply returns immediately.
//
// Only fields that refer to AST nodes are considered children;
// i.e., token.Pos, Scopes, Objects, and fields of basic types
// (strings, etc.) are ignored.
//
// Children are traversed in the order in which they appear in the
// respective node's struct definition.
//
func Apply(node ast.Node, before, after func(Cursor) bool) ast.Node {
apply(&applier{before: before, after: after}, nil, &node)
return node
}
// A applyVisitor's before method is invoked for each node encountered by Walk.
// If the result applyVisitor w is true, Walk visits each of the children
// of node with the applyVisitor w, followed by a call of w.After.
type applyVisitor interface {
Before(Cursor) applyVisitor
After(Cursor) bool
}
// Helper functions for common node lists. They may be empty.
func applyExprList(v applyVisitor, parent Cursor, ptr interface{}, list []ast.Expr) {
c := newCursor(parent, nil, nil)
for i, x := range list {
c.index = i
c.node = x
c.typ = &list[i]
applyCursor(v, c)
if x != c.node {
list[i] = c.node.(ast.Expr)
}
}
}
type declsCursor struct {
*cursor
decls, after, process []ast.Decl
delete bool
}
func (c *declsCursor) InsertAfter(n ast.Node) {
if r, ok := n.(recursive); ok {
n = r.Node
c.process = append(c.process, n.(ast.Decl))
}
c.after = append(c.after, n.(ast.Decl))
}
func (c *declsCursor) InsertBefore(n ast.Node) {
if r, ok := n.(recursive); ok {
n = r.Node
c.process = append(c.process, n.(ast.Decl))
}
c.decls = append(c.decls, n.(ast.Decl))
}
func (c *declsCursor) Delete() { c.delete = true }
func applyDeclList(v applyVisitor, parent Cursor, list []ast.Decl) []ast.Decl {
c := &declsCursor{
cursor: newCursor(parent, nil, nil),
decls: make([]ast.Decl, 0, len(list)),
}
if file, ok := parent.Node().(*ast.File); ok {
c.cursor.file = &info{f: file, current: c}
}
for i, x := range list {
c.node = x
c.typ = &list[i]
applyCursor(v, c)
if !c.delete {
c.decls = append(c.decls, c.node.(ast.Decl))
}
c.delete = false
for i := 0; i < len(c.process); i++ {
x := c.process[i]
c.node = x
c.typ = &c.process[i]
applyCursor(v, c)
if c.delete {
panic("cannot delete a node that was added with InsertBefore or InsertAfter")
}
}
c.decls = append(c.decls, c.after...)
c.after = c.after[:0]
c.process = c.process[:0]
}
// TODO: ultimately, programmatically linked nodes have to be resolved
// at the end.
// if info := c.cursor.file; info != nil {
// done := map[*ast.ImportSpec]bool{}
// for _, ident := range info.importPatch {
// spec := ident.Node.(*ast.ImportSpec)
// if done[spec] {
// continue
// }
// done[spec] = true
// path, _ := strconv.Unquote(spec.Path)
// ident.Name =
// }
// }
return c.decls
}
func apply(v applyVisitor, parent Cursor, nodePtr interface{}) {
res := reflect.Indirect(reflect.ValueOf(nodePtr))
n := res.Interface()
node := n.(ast.Node)
c := newCursor(parent, node, nodePtr)
applyCursor(v, c)
if node != c.node {
res.Set(reflect.ValueOf(c.node))
}
}
// applyCursor traverses an AST in depth-first order: It starts by calling
// v.Visit(node); node must not be nil. If the visitor w returned by
// v.Visit(node) is not nil, apply is invoked recursively with visitor
// w for each of the non-nil children of node, followed by a call of
// w.Visit(nil).
//
func applyCursor(v applyVisitor, c Cursor) {
if v = v.Before(c); v == nil {
return
}
node := c.Node()
// TODO: record the comment groups and interleave with the values like for
// parsing and printing?
comments := node.Comments()
for _, cm := range comments {
apply(v, c, &cm)
}
// apply children
// (the order of the cases matches the order
// of the corresponding node types in go)
switch n := node.(type) {
// Comments and fields
case *ast.Comment:
// nothing to do
case *ast.CommentGroup:
for _, cg := range n.List {
apply(v, c, &cg)
}
case *ast.Attribute:
// nothing to do
case *ast.Field:
apply(v, c, &n.Label)
if n.Value != nil {
apply(v, c, &n.Value)
}
for _, a := range n.Attrs {
apply(v, c, &a)
}
case *ast.StructLit:
n.Elts = applyDeclList(v, c, n.Elts)
// Expressions
case *ast.BottomLit, *ast.BadExpr, *ast.Ident, *ast.BasicLit:
// nothing to do
case *ast.TemplateLabel:
apply(v, c, &n.Ident)
case *ast.Interpolation:
applyExprList(v, c, &n, n.Elts)
case *ast.ListLit:
applyExprList(v, c, &n, n.Elts)
case *ast.Ellipsis:
if n.Type != nil {
apply(v, c, &n.Type)
}
case *ast.ParenExpr:
apply(v, c, &n.X)
case *ast.SelectorExpr:
apply(v, c, &n.X)
apply(v, c, &n.Sel)
case *ast.IndexExpr:
apply(v, c, &n.X)
apply(v, c, &n.Index)
case *ast.SliceExpr:
apply(v, c, &n.X)
if n.Low != nil {
apply(v, c, &n.Low)
}
if n.High != nil {
apply(v, c, &n.High)
}
case *ast.CallExpr:
apply(v, c, &n.Fun)
applyExprList(v, c, &n, n.Args)
case *ast.UnaryExpr:
apply(v, c, &n.X)
case *ast.BinaryExpr:
apply(v, c, &n.X)
apply(v, c, &n.Y)
// Declarations
case *ast.ImportSpec:
if n.Name != nil {
apply(v, c, &n.Name)
}
apply(v, c, &n.Path)
case *ast.BadDecl:
// nothing to do
case *ast.ImportDecl:
for _, s := range n.Specs {
apply(v, c, &s)
}
case *ast.EmbedDecl:
apply(v, c, &n.Expr)
case *ast.Alias:
apply(v, c, &n.Ident)
apply(v, c, &n.Expr)
case *ast.Comprehension:
clauses := n.Clauses
for i := range n.Clauses {
apply(v, c, &clauses[i])
}
apply(v, c, &n.Value)
// Files and packages
case *ast.File:
n.Decls = applyDeclList(v, c, n.Decls)
case *ast.Package:
apply(v, c, &n.Name)
case *ast.ListComprehension:
apply(v, c, &n.Expr)
clauses := n.Clauses
for i := range clauses {
apply(v, c, &clauses[i])
}
case *ast.ForClause:
if n.Key != nil {
apply(v, c, &n.Key)
}
apply(v, c, &n.Value)
apply(v, c, &n.Source)
case *ast.IfClause:
apply(v, c, &n.Condition)
default:
panic(fmt.Sprintf("Walk: unexpected node type %T", n))
}
v.After(c)
}
type applier struct {
before func(Cursor) bool
after func(Cursor) bool
commentStack []commentFrame
current commentFrame
}
type commentFrame struct {
cg []*ast.CommentGroup
pos int8
}
func (f *applier) Before(c Cursor) applyVisitor {
node := c.Node()
if f.before == nil || (f.before(c) && node == c.Node()) {
f.commentStack = append(f.commentStack, f.current)
f.current = commentFrame{cg: node.Comments()}
f.visitComments(c, f.current.pos)
return f
}
return nil
}
func (f *applier) After(c Cursor) bool {
f.visitComments(c, 127)
p := len(f.commentStack) - 1
f.current = f.commentStack[p]
f.commentStack = f.commentStack[:p]
f.current.pos++
if f.after != nil {
f.after(c)
}
return true
}
func (f *applier) visitComments(p Cursor, pos int8) {
c := &f.current
for i := 0; i < len(c.cg); i++ {
cg := c.cg[i]
if cg.Position == pos {
continue
}
cursor := newCursor(p, cg, cg)
if f.before == nil || (f.before(cursor) && !cursor.replaced) {
for j, c := range cg.List {
cursor := newCursor(p, c, &c)
if f.before == nil || (f.before(cursor) && !cursor.replaced) {
if f.after != nil {
f.after(cursor)
}
}
cg.List[j] = cursor.node.(*ast.Comment)
}
if f.after != nil {
f.after(cursor)
}
}
c.cg[i] = cursor.node.(*ast.CommentGroup)
}
}