tree - walk.go

     1  package tree
     2  
     3  import (
     4  	"errors"
     5  	"iter"
     6  	"slices"
     7  )
     8  
     9  // WalkFunc is the type of the function called by Walk to visit each Node.
    10  //
    11  // The path argument contains the names of the Nodes leading to the current
    12  // Node. The WalkFunc should not retain the path slice, instead making a copy
    13  // with slices.Clone or similar.
    14  //
    15  // The returned error controls how the Walk continues.
    16  //
    17  // The SkipNode error will prevent walking the children of the current Node.
    18  //
    19  // The SkipAll error will cancel the remaining walk, but return a nil error
    20  // from Walk.
    21  //
    22  // Any other error will be returned via the Walk function.
    23  type WalkFunc func(path []string, n Node) error
    24  
    25  // Walk recursively walks the tree hierarchy, calling the supplied function for
    26  // each Node visited.
    27  //
    28  // See the WalkFunc type for information on the arguments and how the returned
    29  // error is handled.
    30  func Walk(n Node, fn WalkFunc) error {
    31  	if err := walk(n, fn, nil); err != SkipAll {
    32  		return err
    33  	}
    34  
    35  	return nil
    36  }
    37  
    38  func walk(n Node, fn WalkFunc, path []string) error {
    39  	for name, child := range n.Children() {
    40  		cp := append(path, name)
    41  
    42  		switch err := fn(cp, child); err {
    43  		default:
    44  			return err
    45  		case nil:
    46  			if err := walk(child, fn, cp); err != nil {
    47  				return err
    48  			}
    49  		case SkipNode:
    50  		}
    51  	}
    52  
    53  	return nil
    54  }
    55  
    56  // Flatten iterates through the tree returning each path and node in lexical
    57  // order.
    58  func Flatten(n Node) iter.Seq2[[]string, Node] {
    59  	return func(yield func([]string, Node) bool) {
    60  		Walk(n, func(path []string, n Node) error {
    61  			if !yield(slices.Clone(path), n) {
    62  				return SkipAll
    63  			}
    64  
    65  			return nil
    66  		})
    67  	}
    68  }
    69  
    70  // Filter iterates through the tree in lexical order, returning the path and
    71  // node for each Node that passes the given test function.
    72  //
    73  // The function should return < 0 to skip the current node and all of its
    74  // children, 0 to continue recursing down the tree, but not to yield the current
    75  // node, and > 0 to yield the current node and continue recursing.
    76  func Filter(n Node, f func([]string, Node) int) iter.Seq2[[]string, Node] {
    77  	return func(yield func([]string, Node) bool) {
    78  		Walk(n, func(path []string, n Node) error {
    79  			if r := f(path, n); r < 0 {
    80  				return SkipNode
    81  			} else if r == 0 || yield(slices.Clone(path), n) {
    82  				return nil
    83  			}
    84  
    85  			return SkipAll
    86  		})
    87  	}
    88  }
    89  
    90  var (
    91  	// SkipNode tells the Walk to not iterate through the current Nodes children.
    92  	SkipNode = errors.New("skip this node")
    93  
    94  	// SkipAll tells the Walk to stop and return a nil error.
    95  	SkipAll = errors.New("stop the walk")
    96  )
    97