github.com/k14s/starlark-go@v0.0.0-20200720175618-3a5c849cc368/syntax/walk.go (about)

     1  // Copyright 2017 The Bazel 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 syntax
     6  
     7  // Walk traverses a syntax tree in depth-first order.
     8  // It starts by calling f(n); n must not be nil.
     9  // If f returns true, Walk calls itself
    10  // recursively for each non-nil child of n.
    11  // Walk then calls f(nil).
    12  func Walk(n Node, f func(Node) bool) {
    13  	if n == nil {
    14  		panic("nil")
    15  	}
    16  	if !f(n) {
    17  		return
    18  	}
    19  
    20  	// TODO(adonovan): opt: order cases using profile data.
    21  	switch n := n.(type) {
    22  	case *File:
    23  		walkStmts(n.Stmts, f)
    24  
    25  	case *ExprStmt:
    26  		Walk(n.X, f)
    27  
    28  	case *BranchStmt:
    29  		// no-op
    30  
    31  	case *IfStmt:
    32  		Walk(n.Cond, f)
    33  		walkStmts(n.True, f)
    34  		walkStmts(n.False, f)
    35  
    36  	case *AssignStmt:
    37  		Walk(n.LHS, f)
    38  		Walk(n.RHS, f)
    39  
    40  	case *DefStmt:
    41  		Walk(n.Name, f)
    42  		for _, param := range n.Params {
    43  			Walk(param, f)
    44  		}
    45  		walkStmts(n.Body, f)
    46  
    47  	case *ForStmt:
    48  		Walk(n.Vars, f)
    49  		Walk(n.X, f)
    50  		walkStmts(n.Body, f)
    51  
    52  	case *ReturnStmt:
    53  		if n.Result != nil {
    54  			Walk(n.Result, f)
    55  		}
    56  
    57  	case *LoadStmt:
    58  		Walk(n.Module, f)
    59  		for _, from := range n.From {
    60  			Walk(from, f)
    61  		}
    62  		for _, to := range n.To {
    63  			Walk(to, f)
    64  		}
    65  
    66  	case *Ident, *Literal:
    67  		// no-op
    68  
    69  	case *ListExpr:
    70  		for _, x := range n.List {
    71  			Walk(x, f)
    72  		}
    73  
    74  	case *ParenExpr:
    75  		Walk(n.X, f)
    76  
    77  	case *CondExpr:
    78  		Walk(n.Cond, f)
    79  		Walk(n.True, f)
    80  		Walk(n.False, f)
    81  
    82  	case *IndexExpr:
    83  		Walk(n.X, f)
    84  		Walk(n.Y, f)
    85  
    86  	case *DictEntry:
    87  		Walk(n.Key, f)
    88  		Walk(n.Value, f)
    89  
    90  	case *SliceExpr:
    91  		Walk(n.X, f)
    92  		if n.Lo != nil {
    93  			Walk(n.Lo, f)
    94  		}
    95  		if n.Hi != nil {
    96  			Walk(n.Hi, f)
    97  		}
    98  		if n.Step != nil {
    99  			Walk(n.Step, f)
   100  		}
   101  
   102  	case *Comprehension:
   103  		Walk(n.Body, f)
   104  		for _, clause := range n.Clauses {
   105  			Walk(clause, f)
   106  		}
   107  
   108  	case *IfClause:
   109  		Walk(n.Cond, f)
   110  
   111  	case *ForClause:
   112  		Walk(n.Vars, f)
   113  		Walk(n.X, f)
   114  
   115  	case *TupleExpr:
   116  		for _, x := range n.List {
   117  			Walk(x, f)
   118  		}
   119  
   120  	case *DictExpr:
   121  		for _, entry := range n.List {
   122  			entry := entry.(*DictEntry)
   123  			Walk(entry.Key, f)
   124  			Walk(entry.Value, f)
   125  		}
   126  
   127  	case *UnaryExpr:
   128  		if n.X != nil {
   129  			Walk(n.X, f)
   130  		}
   131  
   132  	case *BinaryExpr:
   133  		Walk(n.X, f)
   134  		Walk(n.Y, f)
   135  
   136  	case *DotExpr:
   137  		Walk(n.X, f)
   138  		Walk(n.Name, f)
   139  
   140  	case *CallExpr:
   141  		Walk(n.Fn, f)
   142  		for _, arg := range n.Args {
   143  			Walk(arg, f)
   144  		}
   145  
   146  	case *LambdaExpr:
   147  		for _, param := range n.Params {
   148  			Walk(param, f)
   149  		}
   150  		Walk(n.Body, f)
   151  
   152  	default:
   153  		panic(n)
   154  	}
   155  
   156  	f(nil)
   157  }
   158  
   159  func walkStmts(stmts []Stmt, f func(Node) bool) {
   160  	for _, stmt := range stmts {
   161  		Walk(stmt, f)
   162  	}
   163  }