|
| 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 typeindex provides an [Index] of type information for a |
| 6 | +// package, allowing efficient lookup of, say, whether a given symbol |
| 7 | +// is referenced and, if so, where from; or of the [inspector.Cursor] for |
| 8 | +// the declaration of a particular [types.Object] symbol. |
| 9 | +package typeindex |
| 10 | + |
| 11 | +import ( |
| 12 | + "encoding/binary" |
| 13 | + "go/ast" |
| 14 | + "go/types" |
| 15 | + "iter" |
| 16 | + |
| 17 | + "golang.org/x/tools/go/ast/edge" |
| 18 | + "golang.org/x/tools/go/ast/inspector" |
| 19 | + "golang.org/x/tools/go/types/typeutil" |
| 20 | +) |
| 21 | + |
| 22 | +// IsPackageLevel reports whether obj is a package-level symbol. |
| 23 | +func IsPackageLevel(obj types.Object) bool { |
| 24 | + return obj.Pkg() != nil && obj.Parent() == obj.Pkg().Scope() |
| 25 | +} |
| 26 | + |
| 27 | +// New constructs an Index for the package of type-annotated syntax |
| 28 | +// |
| 29 | +// TODO(adonovan): accept a FileSet too? |
| 30 | +// We regret not requiring one in inspector.New. |
| 31 | +func New(inspect *inspector.Inspector, pkg *types.Package, info *types.Info) *Index { |
| 32 | + ix := &Index{ |
| 33 | + inspect: inspect, |
| 34 | + info: info, |
| 35 | + packages: make(map[string]*types.Package), |
| 36 | + def: make(map[types.Object]inspector.Cursor), |
| 37 | + uses: make(map[types.Object]*uses), |
| 38 | + } |
| 39 | + |
| 40 | + addPackage := func(pkg2 *types.Package) { |
| 41 | + if pkg2 != nil && pkg2 != pkg { |
| 42 | + ix.packages[pkg2.Path()] = pkg2 |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + for cur := range inspect.Root().Preorder((*ast.ImportSpec)(nil), (*ast.Ident)(nil)) { |
| 47 | + switch n := cur.Node().(type) { |
| 48 | + case *ast.ImportSpec: |
| 49 | + // Index direct imports, including blank ones. |
| 50 | + if pkgname := info.PkgNameOf(n); pkgname != nil { |
| 51 | + addPackage(pkgname.Imported()) |
| 52 | + } |
| 53 | + |
| 54 | + case *ast.Ident: |
| 55 | + // Index all defining and using identifiers. |
| 56 | + if obj := info.Defs[n]; obj != nil { |
| 57 | + ix.def[obj] = cur |
| 58 | + } |
| 59 | + |
| 60 | + if obj := info.Uses[n]; obj != nil { |
| 61 | + // Index indirect dependencies (via fields and methods). |
| 62 | + if !IsPackageLevel(obj) { |
| 63 | + addPackage(obj.Pkg()) |
| 64 | + } |
| 65 | + |
| 66 | + us, ok := ix.uses[obj] |
| 67 | + if !ok { |
| 68 | + us = &uses{} |
| 69 | + us.code = us.initial[:0] |
| 70 | + ix.uses[obj] = us |
| 71 | + } |
| 72 | + delta := cur.Index() - us.last |
| 73 | + if delta < 0 { |
| 74 | + panic("non-monotonic") |
| 75 | + } |
| 76 | + us.code = binary.AppendUvarint(us.code, uint64(delta)) |
| 77 | + us.last = cur.Index() |
| 78 | + } |
| 79 | + } |
| 80 | + } |
| 81 | + return ix |
| 82 | +} |
| 83 | + |
| 84 | +// An Index holds an index mapping [types.Object] symbols to their syntax. |
| 85 | +// In effect, it is the inverse of [types.Info]. |
| 86 | +type Index struct { |
| 87 | + inspect *inspector.Inspector |
| 88 | + info *types.Info |
| 89 | + packages map[string]*types.Package // packages of all symbols referenced from this package |
| 90 | + def map[types.Object]inspector.Cursor // Cursor of *ast.Ident that defines the Object |
| 91 | + uses map[types.Object]*uses // Cursors of *ast.Idents that use the Object |
| 92 | +} |
| 93 | + |
| 94 | +// A uses holds the list of Cursors of Idents that use a given symbol. |
| 95 | +// |
| 96 | +// The Uses map of [types.Info] is substantial, so it pays to compress |
| 97 | +// its inverse mapping here, both in space and in CPU due to reduced |
| 98 | +// allocation. A Cursor is 2 words; a Cursor.Index is 4 bytes; but |
| 99 | +// since Cursors are naturally delivered in ascending order, we can |
| 100 | +// use varint-encoded deltas at a cost of only ~1.7-2.2 bytes per use. |
| 101 | +// |
| 102 | +// Many variables have only one or two uses, so their encoded uses may |
| 103 | +// fit in the 4 bytes of initial, saving further CPU and space |
| 104 | +// essentially for free since the struct's size class is 4 words. |
| 105 | +type uses struct { |
| 106 | + code []byte // varint-encoded deltas of successive Cursor.Index values |
| 107 | + last int32 // most recent Cursor.Index value; used during encoding |
| 108 | + initial [4]byte // use slack in size class as initial space for code |
| 109 | +} |
| 110 | + |
| 111 | +// Uses returns the sequence of Cursors of [*ast.Ident]s in this package |
| 112 | +// that refer to obj. If obj is nil, the sequence is empty. |
| 113 | +func (ix *Index) Uses(obj types.Object) iter.Seq[inspector.Cursor] { |
| 114 | + return func(yield func(inspector.Cursor) bool) { |
| 115 | + if uses := ix.uses[obj]; uses != nil { |
| 116 | + var last int32 |
| 117 | + for code := uses.code; len(code) > 0; { |
| 118 | + delta, n := binary.Uvarint(code) |
| 119 | + last += int32(delta) |
| 120 | + if !yield(ix.inspect.At(last)) { |
| 121 | + return |
| 122 | + } |
| 123 | + code = code[n:] |
| 124 | + } |
| 125 | + } |
| 126 | + } |
| 127 | +} |
| 128 | + |
| 129 | +// Used reports whether any of the specified objects are used, in |
| 130 | +// other words, obj != nil && Uses(obj) is non-empty for some obj in objs. |
| 131 | +// |
| 132 | +// (This treatment of nil allows Used to be called directly on the |
| 133 | +// result of [Index.Object] so that analyzers can conveniently skip |
| 134 | +// packages that don't use a symbol of interest.) |
| 135 | +func (ix *Index) Used(objs ...types.Object) bool { |
| 136 | + for _, obj := range objs { |
| 137 | + if obj != nil && ix.uses[obj] != nil { |
| 138 | + return true |
| 139 | + } |
| 140 | + } |
| 141 | + return false |
| 142 | +} |
| 143 | + |
| 144 | +// Def returns the Cursor of the [*ast.Ident] in this package |
| 145 | +// that declares the specified object, if any. |
| 146 | +func (ix *Index) Def(obj types.Object) (inspector.Cursor, bool) { |
| 147 | + cur, ok := ix.def[obj] |
| 148 | + return cur, ok |
| 149 | +} |
| 150 | + |
| 151 | +// Package returns the package of the specified path, |
| 152 | +// or nil if it is not referenced from this package. |
| 153 | +func (ix *Index) Package(path string) *types.Package { |
| 154 | + return ix.packages[path] |
| 155 | +} |
| 156 | + |
| 157 | +// Object returns the package-level symbol name within the package of |
| 158 | +// the specified path, or nil if the package or symbol does not exist |
| 159 | +// or is not visible from this package. |
| 160 | +func (ix *Index) Object(path, name string) types.Object { |
| 161 | + if pkg := ix.Package(path); pkg != nil { |
| 162 | + return pkg.Scope().Lookup(name) |
| 163 | + } |
| 164 | + return nil |
| 165 | +} |
| 166 | + |
| 167 | +// Selection returns the named method or field belonging to the |
| 168 | +// package-level type returned by Object(path, typename). |
| 169 | +func (ix *Index) Selection(path, typename, name string) types.Object { |
| 170 | + if obj := ix.Object(path, typename); obj != nil { |
| 171 | + if tname, ok := obj.(*types.TypeName); ok { |
| 172 | + obj, _, _ := types.LookupFieldOrMethod(tname.Type(), true, obj.Pkg(), name) |
| 173 | + return obj |
| 174 | + } |
| 175 | + } |
| 176 | + return nil |
| 177 | +} |
| 178 | + |
| 179 | +// Calls returns the sequence of cursors for *ast.CallExpr nodes that |
| 180 | +// call the specified callee, as defined by [typeutil.Callee]. |
| 181 | +// If callee is nil, the sequence is empty. |
| 182 | +func (ix *Index) Calls(callee types.Object) iter.Seq[inspector.Cursor] { |
| 183 | + return func(yield func(inspector.Cursor) bool) { |
| 184 | + for cur := range ix.Uses(callee) { |
| 185 | + ek, _ := cur.ParentEdge() |
| 186 | + |
| 187 | + // The call may be of the form f() or x.f(), |
| 188 | + // optionally with parens; ascend from f to call. |
| 189 | + // |
| 190 | + // It is tempting but wrong to use the first |
| 191 | + // CallExpr ancestor: we have to make sure the |
| 192 | + // ident is in the CallExpr.Fun position, otherwise |
| 193 | + // f(f, f) would have two spurious matches. |
| 194 | + // Avoiding Enclosing is also significantly faster. |
| 195 | + |
| 196 | + // inverse unparen: f -> (f) |
| 197 | + for ek == edge.ParenExpr_X { |
| 198 | + cur = cur.Parent() |
| 199 | + ek, _ = cur.ParentEdge() |
| 200 | + } |
| 201 | + |
| 202 | + // ascend selector: f -> x.f |
| 203 | + if ek == edge.SelectorExpr_Sel { |
| 204 | + cur = cur.Parent() |
| 205 | + ek, _ = cur.ParentEdge() |
| 206 | + } |
| 207 | + |
| 208 | + // inverse unparen again |
| 209 | + for ek == edge.ParenExpr_X { |
| 210 | + cur = cur.Parent() |
| 211 | + ek, _ = cur.ParentEdge() |
| 212 | + } |
| 213 | + |
| 214 | + // ascend from f or x.f to call |
| 215 | + if ek == edge.CallExpr_Fun { |
| 216 | + curCall := cur.Parent() |
| 217 | + call := curCall.Node().(*ast.CallExpr) |
| 218 | + if typeutil.Callee(ix.info, call) == callee { |
| 219 | + if !yield(curCall) { |
| 220 | + return |
| 221 | + } |
| 222 | + } |
| 223 | + } |
| 224 | + } |
| 225 | + } |
| 226 | +} |
0 commit comments