walk.go

Documentation: github.com/bazelbuild/buildtools/build

     1  /*
     2  Copyright 2020 Google LLC
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      https://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package build
    18  
    19  // StopTraversalError is a special error that tells the walker to not traverse
    20  // further and visit child nodes of the current node.
    21  type StopTraversalError struct{}
    22  
    23  func (m *StopTraversalError) Error() string {
    24  	return "Stop traversal"
    25  }
    26  
    27  // Walk walks the expression tree v, calling f on all subexpressions
    28  // in a preorder traversal.
    29  //
    30  // The stk argument is the stack of expressions in the recursion above x,
    31  // from outermost to innermost.
    32  //
    33  func Walk(v Expr, f func(x Expr, stk []Expr)) {
    34  	var stack []Expr
    35  	walk1(&v, &stack, func(x *Expr, stk []Expr) (Expr, error) {
    36  		f(*x, stk)
    37  		return nil, nil
    38  	})
    39  }
    40  
    41  // WalkPointers is the same as Walk but calls the callback function with pointers to nodes.
    42  func WalkPointers(v Expr, f func(x *Expr, stk []Expr)) {
    43  	var stack []Expr
    44  	walk1(&v, &stack, func(x *Expr, stk []Expr) (Expr, error) {
    45  		f(x, stk)
    46  		return nil, nil
    47  	})
    48  }
    49  
    50  // WalkInterruptable is the same as Walk but allows traversal to be interrupted.
    51  func WalkInterruptable(v Expr, f func(x Expr, stk []Expr) error) {
    52  	var stack []Expr
    53  	walk1(&v, &stack, func(x *Expr, stk []Expr) (Expr, error) {
    54  		err := f(*x, stk)
    55  		return nil, err
    56  	})
    57  }
    58  
    59  // Edit walks the expression tree v, calling f on all subexpressions
    60  // in a preorder traversal. If f returns a non-nil value, the tree is mutated.
    61  // The new value replaces the old one.
    62  //
    63  // The stk argument is the stack of expressions in the recursion above x,
    64  // from outermost to innermost.
    65  //
    66  func Edit(v Expr, f func(x Expr, stk []Expr) Expr) Expr {
    67  	var stack []Expr
    68  	return walk1(&v, &stack, func(x *Expr, stk []Expr) (Expr, error) {
    69  		return f(*x, stk), nil
    70  	})
    71  }
    72  
    73  // EditChildren is similar to Edit but doesn't visit the initial node, instead goes
    74  // directly to its children.
    75  func EditChildren(v Expr, f func(x Expr, stk []Expr) Expr) {
    76  	stack := []Expr{v}
    77  	WalkOnce(v, func(x *Expr) {
    78  		walk1(x, &stack, func(x *Expr, stk []Expr) (Expr, error) {
    79  			return f(*x, stk), nil
    80  		})
    81  	})
    82  }
    83  
    84  // walk1 is a helper function for Walk, WalkWithPostfix, and Edit.
    85  func walk1(v *Expr, stack *[]Expr, f func(x *Expr, stk []Expr) (Expr, error)) Expr {
    86  	if v == nil || *v == nil {
    87  		return nil
    88  	}
    89  
    90  	res, err := f(v, *stack)
    91  	if res != nil {
    92  		*v = res
    93  	}
    94  	if err != nil {
    95  		if _, ok := err.(*StopTraversalError); ok {
    96  			// Don't traverse inside
    97  			return nil
    98  		}
    99  	}
   100  	*stack = append(*stack, *v)
   101  
   102  	WalkOnce(*v, func(x *Expr) {
   103  		walk1(x, stack, f)
   104  	})
   105  
   106  	*stack = (*stack)[:len(*stack)-1]
   107  	return *v
   108  }
   109  
   110  // WalkOnce calls f on every child of v.
   111  func WalkOnce(v Expr, f func(x *Expr)) {
   112  	switch v := v.(type) {
   113  	case *File:
   114  		for i := range v.Stmt {
   115  			f(&v.Stmt[i])
   116  		}
   117  	case *DotExpr:
   118  		f(&v.X)
   119  	case *IndexExpr:
   120  		f(&v.X)
   121  		f(&v.Y)
   122  	case *KeyValueExpr:
   123  		f(&v.Key)
   124  		f(&v.Value)
   125  	case *SliceExpr:
   126  		f(&v.X)
   127  		if v.From != nil {
   128  			f(&v.From)
   129  		}
   130  		if v.To != nil {
   131  			f(&v.To)
   132  		}
   133  		if v.Step != nil {
   134  			f(&v.Step)
   135  		}
   136  	case *ParenExpr:
   137  		f(&v.X)
   138  	case *UnaryExpr:
   139  		f(&v.X)
   140  	case *BinaryExpr:
   141  		f(&v.X)
   142  		f(&v.Y)
   143  	case *AssignExpr:
   144  		f(&v.LHS)
   145  		f(&v.RHS)
   146  	case *LambdaExpr:
   147  		for i := range v.Params {
   148  			f(&v.Params[i])
   149  		}
   150  		for i := range v.Body {
   151  			f(&v.Body[i])
   152  		}
   153  	case *CallExpr:
   154  		f(&v.X)
   155  		for i := range v.List {
   156  			f(&v.List[i])
   157  		}
   158  	case *ListExpr:
   159  		for i := range v.List {
   160  			f(&v.List[i])
   161  		}
   162  	case *SetExpr:
   163  		for i := range v.List {
   164  			f(&v.List[i])
   165  		}
   166  	case *TupleExpr:
   167  		for i := range v.List {
   168  			f(&v.List[i])
   169  		}
   170  	case *DictExpr:
   171  		for i := range v.List {
   172  			e := Expr(v.List[i])
   173  			f(&e)
   174  		}
   175  	case *Comprehension:
   176  		f(&v.Body)
   177  		for _, c := range v.Clauses {
   178  			f(&c)
   179  		}
   180  	case *IfClause:
   181  		f(&v.Cond)
   182  	case *ForClause:
   183  		f(&v.Vars)
   184  		f(&v.X)
   185  	case *ConditionalExpr:
   186  		f(&v.Then)
   187  		f(&v.Test)
   188  		f(&v.Else)
   189  	case *LoadStmt:
   190  		module := (Expr)(v.Module)
   191  		f(&module)
   192  		v.Module = module.(*StringExpr)
   193  		for i := range v.From {
   194  			from := (Expr)(v.From[i])
   195  			f(&from)
   196  			v.From[i] = from.(*Ident)
   197  			to := (Expr)(v.To[i])
   198  			f(&to)
   199  			v.To[i] = to.(*Ident)
   200  		}
   201  	case *DefStmt:
   202  		for i := range v.Params {
   203  			f(&v.Params[i])
   204  		}
   205  		for i := range v.Body {
   206  			f(&v.Body[i])
   207  		}
   208  	case *IfStmt:
   209  		f(&v.Cond)
   210  		for i := range v.True {
   211  			f(&v.True[i])
   212  		}
   213  		for i := range v.False {
   214  			f(&v.False[i])
   215  		}
   216  	case *ForStmt:
   217  		f(&v.Vars)
   218  		f(&v.X)
   219  		for i := range v.Body {
   220  			f(&v.Body[i])
   221  		}
   222  	case *ReturnStmt:
   223  		if v.Result != nil {
   224  			f(&v.Result)
   225  		}
   226  	}
   227  }
   228  
   229  // walkStatements is a helper function for WalkStatements
   230  func walkStatements(v Expr, stack *[]Expr, f func(x Expr, stk []Expr) error) {
   231  	if v == nil {
   232  		return
   233  	}
   234  
   235  	err := f(v, *stack)
   236  	if err != nil {
   237  		if _, ok := err.(*StopTraversalError); ok {
   238  			return
   239  		}
   240  	}
   241  
   242  	*stack = append(*stack, v)
   243  
   244  	traverse := func(x Expr) {
   245  		walkStatements(x, stack, f)
   246  	}
   247  
   248  	switch expr := v.(type) {
   249  	case *File:
   250  		for _, s := range expr.Stmt {
   251  			traverse(s)
   252  		}
   253  	case *DefStmt:
   254  		for _, s := range expr.Body {
   255  			traverse(s)
   256  		}
   257  	case *IfStmt:
   258  		for _, s := range expr.True {
   259  			traverse(s)
   260  		}
   261  		for _, s := range expr.False {
   262  			traverse(s)
   263  		}
   264  	case *ForStmt:
   265  		for _, s := range expr.Body {
   266  			traverse(s)
   267  		}
   268  	}
   269  
   270  	*stack = (*stack)[:len(*stack)-1]
   271  }
   272  
   273  // WalkStatements traverses sub statements (not all nodes)
   274  func WalkStatements(v Expr, f func(x Expr, stk []Expr) error) {
   275  	var stack []Expr
   276  	walkStatements(v, &stack, f)
   277  }
   278
View as plain text