walk.go

     1  // Copyright 2009 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package ast
     6  
     7  import (
     8  	"fmt"
     9  	"iter"
    10  )
    11  
    12  // A Visitor's Visit method is invoked for each node encountered by [Walk].
    13  // If the result visitor w is not nil, [Walk] visits each of the children
    14  // of node with the visitor w, followed by a call of w.Visit(nil).
    15  type Visitor interface {
    16  	Visit(node Node) (w Visitor)
    17  }
    18  
    19  func walkList[N Node](v Visitor, list []N) {
    20  	for _, node := range list {
    21  		Walk(v, node)
    22  	}
    23  }
    24  
    25  // TODO(gri): Investigate if providing a closure to Walk leads to
    26  // simpler use (and may help eliminate Inspect in turn).
    27  
    28  // Walk traverses an AST in depth-first order: It starts by calling
    29  // v.Visit(node); node must not be nil. If the visitor w returned by
    30  // v.Visit(node) is not nil, Walk is invoked recursively with visitor
    31  // w for each of the non-nil children of node, followed by a call of
    32  // w.Visit(nil).
    33  func Walk(v Visitor, node Node) {
    34  	if v = v.Visit(node); v == nil {
    35  		return
    36  	}
    37  
    38  	// walk children
    39  	// (the order of the cases matches the order
    40  	// of the corresponding node types in ast.go)
    41  	switch n := node.(type) {
    42  	// Comments and fields
    43  	case *Comment:
    44  		// nothing to do
    45  
    46  	case *CommentGroup:
    47  		walkList(v, n.List)
    48  
    49  	case *Field:
    50  		if n.Doc != nil {
    51  			Walk(v, n.Doc)
    52  		}
    53  		walkList(v, n.Names)
    54  		if n.Type != nil {
    55  			Walk(v, n.Type)
    56  		}
    57  		if n.Tag != nil {
    58  			Walk(v, n.Tag)
    59  		}
    60  		if n.Comment != nil {
    61  			Walk(v, n.Comment)
    62  		}
    63  
    64  	case *FieldList:
    65  		walkList(v, n.List)
    66  
    67  	// Expressions
    68  	case *BadExpr, *Ident, *BasicLit:
    69  		// nothing to do
    70  
    71  	case *Ellipsis:
    72  		if n.Elt != nil {
    73  			Walk(v, n.Elt)
    74  		}
    75  
    76  	case *FuncLit:
    77  		Walk(v, n.Type)
    78  		Walk(v, n.Body)
    79  
    80  	case *CompositeLit:
    81  		if n.Type != nil {
    82  			Walk(v, n.Type)
    83  		}
    84  		walkList(v, n.Elts)
    85  
    86  	case *ParenExpr:
    87  		Walk(v, n.X)
    88  
    89  	case *SelectorExpr:
    90  		Walk(v, n.X)
    91  		Walk(v, n.Sel)
    92  
    93  	case *IndexExpr:
    94  		Walk(v, n.X)
    95  		Walk(v, n.Index)
    96  
    97  	case *IndexListExpr:
    98  		Walk(v, n.X)
    99  		walkList(v, n.Indices)
   100  
   101  	case *SliceExpr:
   102  		Walk(v, n.X)
   103  		if n.Low != nil {
   104  			Walk(v, n.Low)
   105  		}
   106  		if n.High != nil {
   107  			Walk(v, n.High)
   108  		}
   109  		if n.Max != nil {
   110  			Walk(v, n.Max)
   111  		}
   112  
   113  	case *TypeAssertExpr:
   114  		Walk(v, n.X)
   115  		if n.Type != nil {
   116  			Walk(v, n.Type)
   117  		}
   118  
   119  	case *CallExpr:
   120  		Walk(v, n.Fun)
   121  		walkList(v, n.Args)
   122  
   123  	case *StarExpr:
   124  		Walk(v, n.X)
   125  
   126  	case *UnaryExpr:
   127  		Walk(v, n.X)
   128  
   129  	case *BinaryExpr:
   130  		Walk(v, n.X)
   131  		Walk(v, n.Y)
   132  
   133  	case *KeyValueExpr:
   134  		Walk(v, n.Key)
   135  		Walk(v, n.Value)
   136  
   137  	// Types
   138  	case *ArrayType:
   139  		if n.Len != nil {
   140  			Walk(v, n.Len)
   141  		}
   142  		Walk(v, n.Elt)
   143  
   144  	case *StructType:
   145  		Walk(v, n.Fields)
   146  
   147  	case *FuncType:
   148  		if n.TypeParams != nil {
   149  			Walk(v, n.TypeParams)
   150  		}
   151  		if n.Params != nil {
   152  			Walk(v, n.Params)
   153  		}
   154  		if n.Results != nil {
   155  			Walk(v, n.Results)
   156  		}
   157  
   158  	case *InterfaceType:
   159  		Walk(v, n.Methods)
   160  
   161  	case *MapType:
   162  		Walk(v, n.Key)
   163  		Walk(v, n.Value)
   164  
   165  	case *ChanType:
   166  		Walk(v, n.Value)
   167  
   168  	// Statements
   169  	case *BadStmt:
   170  		// nothing to do
   171  
   172  	case *DeclStmt:
   173  		Walk(v, n.Decl)
   174  
   175  	case *EmptyStmt:
   176  		// nothing to do
   177  
   178  	case *LabeledStmt:
   179  		Walk(v, n.Label)
   180  		Walk(v, n.Stmt)
   181  
   182  	case *ExprStmt:
   183  		Walk(v, n.X)
   184  
   185  	case *SendStmt:
   186  		Walk(v, n.Chan)
   187  		Walk(v, n.Value)
   188  
   189  	case *IncDecStmt:
   190  		Walk(v, n.X)
   191  
   192  	case *AssignStmt:
   193  		walkList(v, n.Lhs)
   194  		walkList(v, n.Rhs)
   195  
   196  	case *GoStmt:
   197  		Walk(v, n.Call)
   198  
   199  	case *DeferStmt:
   200  		Walk(v, n.Call)
   201  
   202  	case *ReturnStmt:
   203  		walkList(v, n.Results)
   204  
   205  	case *BranchStmt:
   206  		if n.Label != nil {
   207  			Walk(v, n.Label)
   208  		}
   209  
   210  	case *BlockStmt:
   211  		walkList(v, n.List)
   212  
   213  	case *IfStmt:
   214  		if n.Init != nil {
   215  			Walk(v, n.Init)
   216  		}
   217  		Walk(v, n.Cond)
   218  		Walk(v, n.Body)
   219  		if n.Else != nil {
   220  			Walk(v, n.Else)
   221  		}
   222  
   223  	case *CaseClause:
   224  		walkList(v, n.List)
   225  		walkList(v, n.Body)
   226  
   227  	case *SwitchStmt:
   228  		if n.Init != nil {
   229  			Walk(v, n.Init)
   230  		}
   231  		if n.Tag != nil {
   232  			Walk(v, n.Tag)
   233  		}
   234  		Walk(v, n.Body)
   235  
   236  	case *TypeSwitchStmt:
   237  		if n.Init != nil {
   238  			Walk(v, n.Init)
   239  		}
   240  		Walk(v, n.Assign)
   241  		Walk(v, n.Body)
   242  
   243  	case *CommClause:
   244  		if n.Comm != nil {
   245  			Walk(v, n.Comm)
   246  		}
   247  		walkList(v, n.Body)
   248  
   249  	case *SelectStmt:
   250  		Walk(v, n.Body)
   251  
   252  	case *ForStmt:
   253  		if n.Init != nil {
   254  			Walk(v, n.Init)
   255  		}
   256  		if n.Cond != nil {
   257  			Walk(v, n.Cond)
   258  		}
   259  		if n.Post != nil {
   260  			Walk(v, n.Post)
   261  		}
   262  		Walk(v, n.Body)
   263  
   264  	case *RangeStmt:
   265  		if n.Key != nil {
   266  			Walk(v, n.Key)
   267  		}
   268  		if n.Value != nil {
   269  			Walk(v, n.Value)
   270  		}
   271  		Walk(v, n.X)
   272  		Walk(v, n.Body)
   273  
   274  	// Declarations
   275  	case *ImportSpec:
   276  		if n.Doc != nil {
   277  			Walk(v, n.Doc)
   278  		}
   279  		if n.Name != nil {
   280  			Walk(v, n.Name)
   281  		}
   282  		Walk(v, n.Path)
   283  		if n.Comment != nil {
   284  			Walk(v, n.Comment)
   285  		}
   286  
   287  	case *ValueSpec:
   288  		if n.Doc != nil {
   289  			Walk(v, n.Doc)
   290  		}
   291  		walkList(v, n.Names)
   292  		if n.Type != nil {
   293  			Walk(v, n.Type)
   294  		}
   295  		walkList(v, n.Values)
   296  		if n.Comment != nil {
   297  			Walk(v, n.Comment)
   298  		}
   299  
   300  	case *TypeSpec:
   301  		if n.Doc != nil {
   302  			Walk(v, n.Doc)
   303  		}
   304  		Walk(v, n.Name)
   305  		if n.TypeParams != nil {
   306  			Walk(v, n.TypeParams)
   307  		}
   308  		Walk(v, n.Type)
   309  		if n.Comment != nil {
   310  			Walk(v, n.Comment)
   311  		}
   312  
   313  	case *BadDecl:
   314  		// nothing to do
   315  
   316  	case *GenDecl:
   317  		if n.Doc != nil {
   318  			Walk(v, n.Doc)
   319  		}
   320  		walkList(v, n.Specs)
   321  
   322  	case *FuncDecl:
   323  		if n.Doc != nil {
   324  			Walk(v, n.Doc)
   325  		}
   326  		if n.Recv != nil {
   327  			Walk(v, n.Recv)
   328  		}
   329  		Walk(v, n.Name)
   330  		Walk(v, n.Type)
   331  		if n.Body != nil {
   332  			Walk(v, n.Body)
   333  		}
   334  
   335  	// Files and packages
   336  	case *File:
   337  		if n.Doc != nil {
   338  			Walk(v, n.Doc)
   339  		}
   340  		Walk(v, n.Name)
   341  		walkList(v, n.Decls)
   342  		// don't walk n.Comments - they have been
   343  		// visited already through the individual
   344  		// nodes
   345  
   346  	case *Package:
   347  		for _, f := range n.Files {
   348  			Walk(v, f)
   349  		}
   350  
   351  	default:
   352  		panic(fmt.Sprintf("ast.Walk: unexpected node type %T", n))
   353  	}
   354  
   355  	v.Visit(nil)
   356  }
   357  
   358  type inspector func(Node) bool
   359  
   360  func (f inspector) Visit(node Node) Visitor {
   361  	if f(node) {
   362  		return f
   363  	}
   364  	return nil
   365  }
   366  
   367  // Inspect traverses an AST in depth-first order: It starts by calling
   368  // f(node); node must not be nil. If f returns true, Inspect invokes f
   369  // recursively for each of the non-nil children of node, followed by a
   370  // call of f(nil).
   371  func Inspect(node Node, f func(Node) bool) {
   372  	Walk(inspector(f), node)
   373  }
   374  
   375  // Preorder returns an iterator over all the nodes of the syntax tree
   376  // beneath (and including) the specified root, in depth-first
   377  // preorder.
   378  //
   379  // For greater control over the traversal of each subtree, use [Inspect].
   380  func Preorder(root Node) iter.Seq[Node] {
   381  	return func(yield func(Node) bool) {
   382  		ok := true
   383  		Inspect(root, func(n Node) bool {
   384  			if n != nil {
   385  				// yield must not be called once ok is false.
   386  				ok = ok && yield(n)
   387  			}
   388  			return ok
   389  		})
   390  	}
   391  }
   392
View as plain text