walk_test.go

     1  // Copyright 2024 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_test
     6  
     7  import (
     8  	"go/ast"
     9  	"go/parser"
    10  	"go/token"
    11  	"reflect"
    12  	"slices"
    13  	"strings"
    14  	"testing"
    15  )
    16  
    17  func TestPreorder_Break(t *testing.T) {
    18  	// This test checks that Preorder correctly handles a break statement while
    19  	// in the middle of walking a node. Previously, incorrect handling of the
    20  	// boolean returned by the yield function resulted in the iterator calling
    21  	// yield for sibling nodes even after yield had returned false. With that
    22  	// bug, this test failed with a runtime panic.
    23  	src := "package p\ntype T struct {\n\tF int `json:\"f\"` // a field\n}\n"
    24  
    25  	fset := token.NewFileSet()
    26  	f, err := parser.ParseFile(fset, "", src, 0)
    27  	if err != nil {
    28  		panic(err)
    29  	}
    30  
    31  	for n := range ast.Preorder(f) {
    32  		if id, ok := n.(*ast.Ident); ok && id.Name == "F" {
    33  			break
    34  		}
    35  	}
    36  }
    37  
    38  func TestPreorderStack(t *testing.T) {
    39  	const src = `package a
    40  func f() {
    41  	print("hello")
    42  }
    43  func g() {
    44  	print("goodbye")
    45  	panic("oops")
    46  }
    47  `
    48  	fset := token.NewFileSet()
    49  	f, _ := parser.ParseFile(fset, "a.go", src, 0)
    50  
    51  	str := func(n ast.Node) string {
    52  		return strings.TrimPrefix(reflect.TypeOf(n).String(), "*ast.")
    53  	}
    54  
    55  	var events []string
    56  	var gotStack []string
    57  	ast.PreorderStack(f, nil, func(n ast.Node, stack []ast.Node) bool {
    58  		events = append(events, str(n))
    59  		if decl, ok := n.(*ast.FuncDecl); ok && decl.Name.Name == "f" {
    60  			return false // skip subtree of f()
    61  		}
    62  		if lit, ok := n.(*ast.BasicLit); ok && lit.Value == `"oops"` {
    63  			for _, n := range stack {
    64  				gotStack = append(gotStack, str(n))
    65  			}
    66  		}
    67  		return true
    68  	})
    69  
    70  	// Check sequence of events.
    71  	wantEvents := []string{
    72  		"File", "Ident", // package a
    73  		"FuncDecl",                                                // func f()  [pruned]
    74  		"FuncDecl", "Ident", "FuncType", "FieldList", "BlockStmt", // func g()
    75  		"ExprStmt", "CallExpr", "Ident", "BasicLit", // print...
    76  		"ExprStmt", "CallExpr", "Ident", "BasicLit", // panic...
    77  	}
    78  	if !slices.Equal(events, wantEvents) {
    79  		t.Errorf("PreorderStack events:\ngot:  %s\nwant: %s", events, wantEvents)
    80  	}
    81  
    82  	// Check captured stack.
    83  	wantStack := []string{"File", "FuncDecl", "BlockStmt", "ExprStmt", "CallExpr"}
    84  	if !slices.Equal(gotStack, wantStack) {
    85  		t.Errorf("PreorderStack stack:\ngot:  %s\nwant: %s", gotStack, wantStack)
    86  	}
    87  }
    88
View as plain text