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