github.com/joomcode/cue@v0.4.4-0.20221111115225-539fe3512047/cue/ast/walk.go (about)

     1  // Copyright 2018 The CUE Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package ast
    16  
    17  import (
    18  	"fmt"
    19  
    20  	"github.com/joomcode/cue/cue/token"
    21  )
    22  
    23  // Walk traverses an AST in depth-first order: It starts by calling f(node);
    24  // node must not be nil. If before returns true, Walk invokes f recursively for
    25  // each of the non-nil children of node, followed by a call of after. Both
    26  // functions may be nil. If before is nil, it is assumed to always return true.
    27  //
    28  func Walk(node Node, before func(Node) bool, after func(Node)) {
    29  	walk(&inspector{before: before, after: after}, node)
    30  }
    31  
    32  // A visitor's before method is invoked for each node encountered by Walk.
    33  // If the result visitor w is true, Walk visits each of the children
    34  // of node with the visitor w, followed by a call of w.After.
    35  type visitor interface {
    36  	Before(node Node) (w visitor)
    37  	After(node Node)
    38  }
    39  
    40  // Helper functions for common node lists. They may be empty.
    41  
    42  func walkExprList(v visitor, list []Expr) {
    43  	for _, x := range list {
    44  		walk(v, x)
    45  	}
    46  }
    47  
    48  func walkDeclList(v visitor, list []Decl) {
    49  	for _, x := range list {
    50  		walk(v, x)
    51  	}
    52  }
    53  
    54  // walk traverses an AST in depth-first order: It starts by calling
    55  // v.Visit(node); node must not be nil. If the visitor w returned by
    56  // v.Visit(node) is not nil, walk is invoked recursively with visitor
    57  // w for each of the non-nil children of node, followed by a call of
    58  // w.Visit(nil).
    59  //
    60  func walk(v visitor, node Node) {
    61  	if v = v.Before(node); v == nil {
    62  		return
    63  	}
    64  
    65  	// TODO: record the comment groups and interleave with the values like for
    66  	// parsing and printing?
    67  	for _, c := range Comments(node) {
    68  		walk(v, c)
    69  	}
    70  
    71  	// walk children
    72  	// (the order of the cases matches the order
    73  	// of the corresponding node types in go)
    74  	switch n := node.(type) {
    75  	// Comments and fields
    76  	case *Comment:
    77  		// nothing to do
    78  
    79  	case *CommentGroup:
    80  		for _, c := range n.List {
    81  			walk(v, c)
    82  		}
    83  
    84  	case *Attribute:
    85  		// nothing to do
    86  
    87  	case *Field:
    88  		walk(v, n.Label)
    89  		if n.Value != nil {
    90  			walk(v, n.Value)
    91  		}
    92  		for _, a := range n.Attrs {
    93  			walk(v, a)
    94  		}
    95  
    96  	case *StructLit:
    97  		walkDeclList(v, n.Elts)
    98  
    99  	// Expressions
   100  	case *BottomLit, *BadExpr, *Ident, *BasicLit:
   101  		// nothing to do
   102  
   103  	case *Interpolation:
   104  		for _, e := range n.Elts {
   105  			walk(v, e)
   106  		}
   107  
   108  	case *ListLit:
   109  		walkExprList(v, n.Elts)
   110  
   111  	case *Ellipsis:
   112  		if n.Type != nil {
   113  			walk(v, n.Type)
   114  		}
   115  
   116  	case *ParenExpr:
   117  		walk(v, n.X)
   118  
   119  	case *SelectorExpr:
   120  		walk(v, n.X)
   121  		walk(v, n.Sel)
   122  
   123  	case *IndexExpr:
   124  		walk(v, n.X)
   125  		walk(v, n.Index)
   126  
   127  	case *SliceExpr:
   128  		walk(v, n.X)
   129  		if n.Low != nil {
   130  			walk(v, n.Low)
   131  		}
   132  		if n.High != nil {
   133  			walk(v, n.High)
   134  		}
   135  
   136  	case *CallExpr:
   137  		walk(v, n.Fun)
   138  		walkExprList(v, n.Args)
   139  
   140  	case *UnaryExpr:
   141  		walk(v, n.X)
   142  
   143  	case *BinaryExpr:
   144  		walk(v, n.X)
   145  		walk(v, n.Y)
   146  
   147  	// Declarations
   148  	case *ImportSpec:
   149  		if n.Name != nil {
   150  			walk(v, n.Name)
   151  		}
   152  		walk(v, n.Path)
   153  
   154  	case *BadDecl:
   155  		// nothing to do
   156  
   157  	case *ImportDecl:
   158  		for _, s := range n.Specs {
   159  			walk(v, s)
   160  		}
   161  
   162  	case *EmbedDecl:
   163  		walk(v, n.Expr)
   164  
   165  	case *LetClause:
   166  		walk(v, n.Ident)
   167  		walk(v, n.Expr)
   168  
   169  	case *Alias:
   170  		walk(v, n.Ident)
   171  		walk(v, n.Expr)
   172  
   173  	case *Comprehension:
   174  		for _, c := range n.Clauses {
   175  			walk(v, c)
   176  		}
   177  		walk(v, n.Value)
   178  
   179  	// Files and packages
   180  	case *File:
   181  		walkDeclList(v, n.Decls)
   182  
   183  	case *Package:
   184  		walk(v, n.Name)
   185  
   186  	case *ForClause:
   187  		if n.Key != nil {
   188  			walk(v, n.Key)
   189  		}
   190  		walk(v, n.Value)
   191  		walk(v, n.Source)
   192  
   193  	case *IfClause:
   194  		walk(v, n.Condition)
   195  
   196  	default:
   197  		panic(fmt.Sprintf("Walk: unexpected node type %T", n))
   198  	}
   199  
   200  	v.After(node)
   201  }
   202  
   203  type inspector struct {
   204  	before func(Node) bool
   205  	after  func(Node)
   206  
   207  	commentStack []commentFrame
   208  	current      commentFrame
   209  }
   210  
   211  type commentFrame struct {
   212  	cg  []*CommentGroup
   213  	pos int8
   214  }
   215  
   216  func (f *inspector) Before(node Node) visitor {
   217  	if f.before == nil || f.before(node) {
   218  		f.commentStack = append(f.commentStack, f.current)
   219  		f.current = commentFrame{cg: Comments(node)}
   220  		f.visitComments(f.current.pos)
   221  		return f
   222  	}
   223  	return nil
   224  }
   225  
   226  func (f *inspector) After(node Node) {
   227  	f.visitComments(127)
   228  	p := len(f.commentStack) - 1
   229  	f.current = f.commentStack[p]
   230  	f.commentStack = f.commentStack[:p]
   231  	f.current.pos++
   232  	if f.after != nil {
   233  		f.after(node)
   234  	}
   235  }
   236  
   237  func (f *inspector) Token(t token.Token) {
   238  	f.current.pos++
   239  }
   240  
   241  func (f *inspector) setPos(i int8) {
   242  	f.current.pos = i
   243  }
   244  
   245  func (f *inspector) visitComments(pos int8) {
   246  	c := &f.current
   247  	for ; len(c.cg) > 0; c.cg = c.cg[1:] {
   248  		cg := c.cg[0]
   249  		if cg.Position == pos {
   250  			continue
   251  		}
   252  		if f.before == nil || f.before(cg) {
   253  			for _, c := range cg.List {
   254  				if f.before == nil || f.before(c) {
   255  					if f.after != nil {
   256  						f.after(c)
   257  					}
   258  				}
   259  			}
   260  			if f.after != nil {
   261  				f.after(cg)
   262  			}
   263  		}
   264  	}
   265  }