util.go

     1  // Copyright 2025 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 astutil
     6  
     7  import (
     8  	"fmt"
     9  	"go/ast"
    10  	"go/token"
    11  	"strconv"
    12  	"unicode/utf8"
    13  )
    14  
    15  // RangeInStringLiteral calculates the positional range within a string literal
    16  // corresponding to the specified start and end byte offsets within the logical string.
    17  func RangeInStringLiteral(lit *ast.BasicLit, start, end int) (token.Pos, token.Pos, error) {
    18  	startPos, err := PosInStringLiteral(lit, start)
    19  	if err != nil {
    20  		return 0, 0, fmt.Errorf("start: %v", err)
    21  	}
    22  	endPos, err := PosInStringLiteral(lit, end)
    23  	if err != nil {
    24  		return 0, 0, fmt.Errorf("end: %v", err)
    25  	}
    26  	return startPos, endPos, nil
    27  }
    28  
    29  // PosInStringLiteral returns the position within a string literal
    30  // corresponding to the specified byte offset within the logical
    31  // string that it denotes.
    32  func PosInStringLiteral(lit *ast.BasicLit, offset int) (token.Pos, error) {
    33  	raw := lit.Value
    34  
    35  	value, err := strconv.Unquote(raw)
    36  	if err != nil {
    37  		return 0, err
    38  	}
    39  	if !(0 <= offset && offset <= len(value)) {
    40  		return 0, fmt.Errorf("invalid offset")
    41  	}
    42  
    43  	// remove quotes
    44  	quote := raw[0] // '"' or '`'
    45  	raw = raw[1 : len(raw)-1]
    46  
    47  	var (
    48  		i   = 0                // byte index within logical value
    49  		pos = lit.ValuePos + 1 // position within literal
    50  	)
    51  	for raw != "" && i < offset {
    52  		r, _, rest, _ := strconv.UnquoteChar(raw, quote) // can't fail
    53  		sz := len(raw) - len(rest)                       // length of literal char in raw bytes
    54  		pos += token.Pos(sz)
    55  		raw = raw[sz:]
    56  		i += utf8.RuneLen(r)
    57  	}
    58  	return pos, nil
    59  }
    60  
    61  // PreorderStack traverses the tree rooted at root,
    62  // calling f before visiting each node.
    63  //
    64  // Each call to f provides the current node and traversal stack,
    65  // consisting of the original value of stack appended with all nodes
    66  // from root to n, excluding n itself. (This design allows calls
    67  // to PreorderStack to be nested without double counting.)
    68  //
    69  // If f returns false, the traversal skips over that subtree. Unlike
    70  // [ast.Inspect], no second call to f is made after visiting node n.
    71  // In practice, the second call is nearly always used only to pop the
    72  // stack, and it is surprisingly tricky to do this correctly; see
    73  // https://go.dev/issue/73319.
    74  func PreorderStack(root ast.Node, stack []ast.Node, f func(n ast.Node, stack []ast.Node) bool) {
    75  	before := len(stack)
    76  	ast.Inspect(root, func(n ast.Node) bool {
    77  		if n != nil {
    78  			if !f(n, stack) {
    79  				// Do not push, as there will be no corresponding pop.
    80  				return false
    81  			}
    82  			stack = append(stack, n) // push
    83  		} else {
    84  			stack = stack[:len(stack)-1] // pop
    85  		}
    86  		return true
    87  	})
    88  	if len(stack) != before {
    89  		panic("push/pop mismatch")
    90  	}
    91  }
    92
View as plain text