From 0c7d0f7d422aadbe5b78366e926823afaa5e7bd0 Mon Sep 17 00:00:00 2001 From: Colby Mchenry Date: Mon, 5 Oct 2026 16:09:30 +0000 Subject: [PATCH] =?UTF-8?q?perf(resolution):=20remove=20the=20Python=20ind?= =?UTF-8?q?exing=20regressions=20since=201.6.2=20=E2=80=94=20graph=20byte-?= =?UTF-8?q?identical=20(#2332)=20(#2344)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Indexing CPython took about three times as long on 1.6.2 as at 290e03f (193 s vs 59 s wall, 8.3 GB vs 3.5 GB peak on a 16-thread Windows box). A CPU profile put ~290 s of the resolver workers' 412 s on two Python checks that first shipped in 1.6.2: - pythonExternalWrites (#2291) read every .py file once per module global its own module types. CPython has 2,375 Python files against the resolver's 1,000-entry content cache, so every read went to disk (184 s). - isPythonLocallyBound (#2198, #2219) re-stripped the calling file for every (function, name), and, asked once per same-named candidate by fitsPythonCallShape, re-derived the enclosing function each time: 3.55M calls for 53K distinct answers (104 s). Now: - A global's own module is read first; when its bindings already leave the type unknown (319 of 379 globals on CPython), no other module is read. - One pass per resolution indexes which Python files spell a name after a dot or a quote, and which write through .__dict__; the write scan visits only those files, in the same path order. Comment stripping only blanks text, so any write the scan accepts is spelled that way in the raw text; a non-ASCII name still reads every file. Index keys are flat copies (a regex capture would pin its whole file). - The repo files an import names are resolved once per mapping, in a WeakMap that lives as long as the resolver's import cache keeps it. - isPythonLocallyBound keeps the current file's stripped lines and per-name module-level answer, and answers a ref once. Verification: nodes, edges and unresolved_refs dumps of CPython and pretix are byte-identical to 6560052a. Medians of 3 interleaved CPython runs on a loaded machine: 290e03f 100 s / 271 s CPU / 3.3 GB, 6560052a 322 s / 801 s / 7.9 GB, this change 123 s / 361 s / 5.5 GB. In the CPU profile pythonExternalWrites drops 184 s -> 16 s and isPythonLocallyBound 104 s -> 13 s. Most of the remaining gap to 290e03f on CPython is JS checks running on its bundled minified d3 (matchDestructuredCallResult, #2334; jsFunctionLocalScope from #2226). __tests__/python-resolution-work.test.ts counts the work instead of timing it: reads of untouched files, strips of the calling file, and node lookups with 16 same-named candidates (4 / 2 / 8 with the fix, 19 / 18 / 128 on main). Each count fails when only its own part of the fix is reverted. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- __tests__/python-resolution-work.test.ts | 102 +++++++++++++++ src/resolution/name-matcher.ts | 155 +++++++++++++++++++---- 3 files changed, 231 insertions(+), 28 deletions(-) create mode 100644 __tests__/python-resolution-work.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index f280890b4..68d7cc062 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,7 +20,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - Upgrading CodeGraph while an agent session is open no longer leaves the old version's background server in charge of your project: the first session started from the new install stops it and starts a current one in its place, even while sessions opened before the upgrade are still running. That old server could no longer load the language parsers the upgrade removed, so it saved every file it re-indexed with no symbols, while new sessions could only read the index beside it without keeping it up to date. A background server from a newer install is never stopped, and sessions opened before the upgrade keep the old version until you restart them. Thanks @lipchey for the report. (#2335) - A file is no longer saved with no symbols when its language parser can't be loaded, which is what happened to every file a background server re-indexed after an upgrade removed its install: the file keeps what it had and is indexed again once the parser loads, and files an earlier version emptied this way are re-indexed by the next sync. A background server also exits on its own once its install is upgraded or removed, so the next session starts one from the current install. Thanks @lipchey for the report. (#2335) - `codegraph status` no longer says the index is up to date while indexed files are missing their symbols: it now names files the parser couldn't read and files stored without their symbols (which `codegraph sync` repairs), `status --json` counts both, and `codegraph files --json` lists each file's recorded errors. Thanks @lipchey for the report. (#2336) - +- Indexing large Python projects is much faster again and needs less memory: since 1.6.2, resolving Python references re-read source files over and over, so a project the size of CPython took several times as long to index. The graph it builds is unchanged. Thanks @bompus for the report. (#2332) ## [1.6.2] - 2026-10-03 diff --git a/__tests__/python-resolution-work.test.ts b/__tests__/python-resolution-work.test.ts new file mode 100644 index 000000000..f4faf445b --- /dev/null +++ b/__tests__/python-resolution-work.test.ts @@ -0,0 +1,102 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import * as fs from 'fs'; +import * as os from 'os'; +import * as path from 'path'; +import { CodeGraph } from '../src'; +import type { ReferenceResolver } from '../src/resolution'; +import { stripCommentsForRegex } from '../src/resolution/strip-comments'; + +// Pass-through, so a test can count how often a file's text is stripped. +vi.mock('../src/resolution/strip-comments', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, stripCommentsForRegex: vi.fn(actual.stripCommentsForRegex) }; +}); + +/** + * #2332: resolving Python does a bounded amount of work per file and per call. + * The cross-module write scan for a module global read every Python file once + * per global; the local-binding check stripped the calling file once per + * function, and looked up the function around a call once per same-named + * candidate. On CPython that was minutes of CPU. Counted, never timed. + */ +describe('Python resolution work (#2332)', () => { + let tmpDir: string | undefined; + let cg: CodeGraph | undefined; + + afterEach(() => { + vi.restoreAllMocks(); + vi.mocked(stripCommentsForRegex).mockClear(); + cg?.close(); + cg = undefined; + if (tmpDir) fs.rmSync(tmpDir, { recursive: true, force: true, maxRetries: 5 }); + tmpDir = undefined; + }); + + /** Index `files`, counting per file what resolution reads, strips and looks up. */ + async function indexCounting(files: Record) { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'cg-py-work-')); + for (const [file, content] of Object.entries(files)) { + fs.mkdirSync(path.dirname(path.join(tmpDir, file)), { recursive: true }); + fs.writeFileSync(path.join(tmpDir, file), content); + } + cg = CodeGraph.initSync(tmpDir); + const context = (cg as unknown as { resolver: ReferenceResolver }).resolver.getResolutionContext(); + const reads = vi.spyOn(context, 'readFile'); + const lookups = vi.spyOn(context, 'getNodesInFile'); + await cg.indexAll(); + return { + graph: cg, + reads: (file: string) => reads.mock.calls.filter(([f]) => f === file).length, + lookups: (file: string) => lookups.mock.calls.filter(([f]) => f === file).length, + strips: (file: string) => vi.mocked(stripCommentsForRegex).mock.calls.filter(([text]) => text === files[file]).length, + }; + } + + it('reads no file once per module global, and strips no file once per function', async () => { + const SCALE = 16; + const conns = Array.from({ length: SCALE }, (_, i) => `conn${i}`); + const files: Record = { + // Globals typed by their own module's writes: each one's type also + // depends on what every other module writes to it. + 'store.py': 'class Store:\n def fetch(self, ids):\n return ids\n', + 'settings.py': `from store import Store\n${conns.map(c => `${c} = None\n`).join('')}\n` + + `def init():\n global ${conns.join(', ')}\n${conns.map(c => ` ${c} = Store()\n`).join('')}`, + 'consumer.py': `import settings\n${conns.map((c, i) => `\ndef cb${i}(pool):\n pool.submit(settings.${c}.fetch)\n`).join('')}`, + // A bare call in every function: each asks whether its function binds the name. + 'helpers.py': 'def helper():\n return 1\n', + 'callers.py': Array.from({ length: SCALE }, (_, i) => `def f${i}():\n return helper()\n\n`).join('') + + 'def shadowed(make):\n helper = make()\n return helper()\n', + }; + // Files no reference touches. + const fillers = Array.from({ length: 20 }, (_, i) => `pkg/filler${i}.py`); + fillers.forEach((file, i) => { files[file] = `def filler${i}(x):\n return x + ${i}\n`; }); + const { graph, reads, strips } = await indexCounting(files); + + // Both paths ran: every global's method value reached Store.fetch, and the + // function that binds `helper` itself calls its own value. + const fetch = graph.getNodesByName('fetch').find(n => n.qualifiedName === 'Store::fetch')!; + expect(graph.getIncomingEdges(fetch.id).filter(e => e.metadata?.fnRef === true) + .map(e => graph.getNode(e.source)?.name).sort()).toEqual(conns.map((_, i) => `cb${i}`).sort()); + const helper = graph.getNodesByName('helper').find(n => n.kind === 'function')!; + expect(graph.getIncomingEdges(helper.id).filter(e => e.kind === 'calls') + .map(e => graph.getNode(e.source)?.name).sort()).toEqual(Array.from({ length: SCALE }, (_, i) => `f${i}`).sort()); + // A few whole-project passes read each file once; the write scan read every + // file once more per global. + expect(Math.max(...fillers.map(reads))).toBeLessThan(SCALE / 2); + // Stripped once for the pass, not once per function. + expect(strips('callers.py')).toBeLessThan(SCALE / 2); + }); + + it('finds the function around a call once, however many functions share its name', async () => { + const CANDIDATES = 16; + const CALLS = 8; + const files: Record = { + 'callers.py': Array.from({ length: CALLS }, (_, i) => `def f${i}():\n return helper()\n\n`).join(''), + }; + for (let k = 0; k < CANDIDATES; k++) files[`lib${k}/helpers.py`] = 'def helper():\n return 1\n'; + const { lookups } = await indexCounting(files); + // Each same-named candidate asks whether the caller binds `helper`; the + // answer is the call's, so the caller's nodes are read per call. + expect(lookups('callers.py')).toBeLessThanOrEqual(2 * CALLS); + }); +}); diff --git a/src/resolution/name-matcher.ts b/src/resolution/name-matcher.ts index e0b4ad28b..2984a4346 100644 --- a/src/resolution/name-matcher.ts +++ b/src/resolution/name-matcher.ts @@ -7,7 +7,7 @@ import * as fs from 'fs'; import * as path from 'path'; import { Language, Node } from '../types'; -import { UnresolvedRef, ResolvedRef, ResolutionContext, isSupertypeTarget, CPP_DEFINE_SIGNATURE, isInheritanceRef, isImportableKind } from './types'; +import { UnresolvedRef, ResolvedRef, ResolutionContext, ImportMapping, isSupertypeTarget, CPP_DEFINE_SIGNATURE, isInheritanceRef, isImportableKind } from './types'; import { blankStringContents, stripCommentsForRegex } from './strip-comments'; import { JS_BUILT_INS, JS_BUILTIN_METHODS, TS_PRIMITIVE_TYPES } from './js-builtins'; import { SWIFT_TYPE_PATH_CALL, resolveSwiftTypePathCall } from './swift-type-visibility'; @@ -659,26 +659,13 @@ function pythonGlobalClasses(global: Node, ref: UnresolvedRef, context: Resoluti let memo = PYTHON_GLOBAL_CLASSES.get(context); if (!memo) { memo = new Map(); PYTHON_GLOBAL_CLASSES.set(context, memo); } if (memo.has(global.id)) return memo.get(global.id)!; - const file = global.filePath; - const external = pythonExternalWrites(global, context); + const own = pythonOwnGlobalWrites(global, context); + const external = own && pythonExternalWrites(global, context); // Each type is resolved in the file that wrote it (its imports name the class). - const writes: Array<{ type: string; file: string }> = [...external.writes]; - let classes: Node[] | null = external.unknown || pythonDynamicGlobalWrite(global.name, file, context) ? null : []; - for (const b of classes ? pythonGlobalBindings(global.name, file, context) : []) { - if (b.kind !== 'assign') { classes = null; break; } - const constructor = b.value && b.value !== 'None' ? pythonConstructorCall(b.value) : null; - if (b.type) { - // `conn: Base = make()` trusts the annotation; `conn: A = B()` contradicts it. - if (constructor && constructor.split('.').pop() !== b.type.split('.').pop()) { classes = null; break; } - writes.push({ type: b.type, file }); - continue; - } - if (b.value === 'None') continue; - if (!constructor) { classes = null; break; } - writes.push({ type: constructor, file }); - } + const writes = external && !external.unknown ? [...external.writes, ...own!] : null; + let classes: Node[] | null = writes && []; const seen = new Set(); - for (const write of classes ? writes : []) { + for (const write of writes ?? []) { const cls = pythonRefClass(write.type, { ...ref, filePath: write.file }, context); if (!cls) { classes = null; break; } if (!seen.has(cls.id)) { seen.add(cls.id); classes!.push(cls); } @@ -687,6 +674,33 @@ function pythonGlobalClasses(global: Node, ref: UnresolvedRef, context: Resoluti return classes; } +/** + * The types the global's own module writes to it, or null when a binding + * there leaves its type unknown — whatever other modules write, so they are + * not read (#2332). + */ +function pythonOwnGlobalWrites(global: Node, context: ResolutionContext): Array<{ type: string; file: string }> | null { + return pythonNameScan(context, `own\0${global.id}`, () => { + const file = global.filePath; + if (pythonDynamicGlobalWrite(global.name, file, context)) return null; + const writes: Array<{ type: string; file: string }> = []; + for (const b of pythonGlobalBindings(global.name, file, context)) { + if (b.kind !== 'assign') return null; + const constructor = b.value && b.value !== 'None' ? pythonConstructorCall(b.value) : null; + if (b.type) { + // `conn: Base = make()` trusts the annotation; `conn: A = B()` contradicts it. + if (constructor && constructor.split('.').pop() !== b.type.split('.').pop()) return null; + writes.push({ type: b.type, file }); + continue; + } + if (b.value === 'None') continue; + if (!constructor) return null; + writes.push({ type: constructor, file }); + } + return writes; + }); +} + /** * The repo files a Python module path can name from `fromFile`. Relative * paths (`..settings`) resolve exactly; absolute ones match a file path @@ -722,9 +736,7 @@ function pythonModuleAliases(filePath: string, moduleFile: string, context: Reso const aliases = new Set(); const ambiguous = new Set(); for (const m of context.getImportMappings(filePath, 'python')) { - const dotted = m.isNamespace ? m.source - : /^\.+$/.test(m.source) ? `${m.source}${m.exportedName}` : `${m.source}.${m.exportedName}`; - const files = pythonModuleFiles(dotted, filePath, context); + const files = pythonImportedFiles(m, filePath, context); if (!files.includes(moduleFile)) continue; // The mapping cannot tell `import a.b` from `import a.b as b`; the source line can. // Exactly this module (not `other.a.b`), outside string literals. @@ -736,6 +748,24 @@ function pythonModuleAliases(filePath: string, moduleFile: string, context: Reso return { aliases: [...aliases], ambiguous: [...ambiguous] }; } +const PYTHON_IMPORTED_FILES = new WeakMap>(); +/** + * The repo files an import of `filePath` can name. Every global's write scan + * asks again of the same files (#2332); kept for as long as the resolver keeps + * the mapping itself, so the memo never outlives its import cache. + */ +function pythonImportedFiles(m: ImportMapping, filePath: string, context: ResolutionContext): string[] { + let memo = PYTHON_IMPORTED_FILES.get(context); + if (!memo) PYTHON_IMPORTED_FILES.set(context, (memo = new WeakMap())); + let files = memo.get(m); + if (!files) { + const dotted = m.isNamespace ? m.source + : /^\.+$/.test(m.source) ? `${m.source}${m.exportedName}` : `${m.source}.${m.exportedName}`; + memo.set(m, (files = pythonModuleFiles(dotted, filePath, context))); + } + return files; +} + const PYTHON_MAIN_GUARD = /^if\s+(?:__name__\s*==\s*(['"])__main__\1|(['"])__main__\2\s*==\s*__name__)\s*:/; /** Line indexes inside a top-level `if __name__ == "__main__":` block — script code, not module state. */ function pythonMainBlockLines(filePath: string, context: ResolutionContext): Set { @@ -775,8 +805,8 @@ function pythonExternalWrites(global: Node, context: ResolutionContext): { write const out = { writes: [] as Array<{ type: string; file: string }>, unknown: false, writers: new Set() }; const name = global.name; const escape = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); - for (const file of context.getAllFiles()) { - if (file === global.filePath || !/\.pyi?$/.test(file) || !context.readFile(file)?.includes(name)) continue; + for (const file of pythonWriteCandidates(name, context)) { + if (file === global.filePath) continue; const test = isPythonTestFile(file); const { aliases, ambiguous } = pythonModuleAliases(file, global.filePath, context); if (!test && !aliases.length && !ambiguous.length) continue; @@ -812,6 +842,41 @@ function pythonExternalWrites(global: Node, context: ResolutionContext): { write }); } +/** + * The Python files that can write a module global named `name`, in path + * order: those that spell it after a dot (`settings.conn = X`) or a quote + * (`setattr(settings, "conn", X)`), and those that write through `.__dict__`, + * whose statement may spell it anywhere. Comment stripping only blanks text, + * so whatever a stripped statement spells, the file's text spells too. The + * files are indexed once per pass, instead of every global reading every + * Python file (#2332). + */ +function pythonWriteCandidates(name: string, context: ResolutionContext): string[] { + const index = pythonNameScan(context, 'write-candidates', () => { + const files = context.getAllFiles().filter(f => /\.pyi?$/.test(f)); + const spelled = new Map(); + const dict: Array<{ at: number; words: string }> = []; + files.forEach((file, at) => { + const source = context.readFile(file) ?? ''; + const names = new Set(); + for (const m of source.matchAll(/[.'"](\w+)/g)) names.add(m[1]!); + for (const n of names) { + const list = spelled.get(n); + // A capture is a sliced view that would pin the file's whole text: key a flat copy. + if (list) list.push(at); else spelled.set(Buffer.from(n).toString(), [at]); + } + // A word-only name is in the text exactly when it is in one of its words. + if (names.has('__dict__')) dict.push({ at, words: [...new Set(source.match(/\w+/g))].join('\n') }); + }); + return { files, spelled, dict }; + }); + // The index holds ASCII words; any other name is looked for in every file's text. + if (!/^\w+$/.test(name)) return index.files.filter(f => context.readFile(f)?.includes(name)); + const hits = new Set(index.spelled.get(name)); + for (const { at, words } of index.dict) if (words.includes(name)) hits.add(at); + return [...hits].sort((a, b) => a - b).map(at => index.files[at]!); +} + /** * Whether the global's own module can write it through its namespace dict: * `globals()` / `vars()` / `sys.modules[__name__]` used as anything but a @@ -853,6 +918,8 @@ function pythonDynamicGlobalWrite(name: string, filePath: string, context: Resol * as a base-typed receiver does. Otherwise, no edge. */ function pythonGlobalMembers(global: Node, member: string, ref: UnresolvedRef, context: ResolutionContext): Node[] { + // Unknown from its own module alone: no edge, and no other module to read. + if (!pythonOwnGlobalWrites(global, context)) return []; // A test that installs its own double sees the double, not the production type. if (pythonExternalWrites(global, context).writers.has(ref.filePath)) return []; const classes = pythonGlobalClasses(global, ref, context); @@ -2538,6 +2605,14 @@ function isDecoratedFixture(n: Node, context: ResolutionContext): boolean { } const PY_LOCAL_BINDS = new WeakMap>(); +/** + * The file isPythonLocallyBound last read: its code lines, and per name + * whether its module binds it. Refs arrive grouped by file, so one file per + * context spares re-stripping the file for every function and name (#2332). + */ +const PY_LOCAL_FILE = new WeakMap }>(); +/** The ref isPythonLocallyBound last answered, and the answer. */ +const PY_LOCAL_LAST = new WeakMap(); /** * Whether the function around a Python call — or its module, at top level — @@ -2547,6 +2622,17 @@ const PY_LOCAL_BINDS = new WeakMap>(); * every such call went to one test file's `def view`. */ function isPythonLocallyBound(name: string, ref: UnresolvedRef, context: ResolutionContext): boolean { + // fitsPythonCallShape asks once per same-named candidate, and finding the + // function around the call reads every node in the file: answer a ref once (#2332). + const last = PY_LOCAL_LAST.get(context); + if (last?.ref === ref && last.name === name) return last.bound; + const bound = pythonLocalBinding(name, ref, context); + PY_LOCAL_LAST.set(context, { ref, name, bound }); + return bound; +} + +/** isPythonLocallyBound's answer, kept per calling function and name. */ +function pythonLocalBinding(name: string, ref: UnresolvedRef, context: ResolutionContext): boolean { const fn = context.getNodesInFile(ref.filePath) .filter((f) => (f.kind === 'function' || f.kind === 'method') && f.startLine <= ref.line && f.endLine >= ref.line) .sort((a, b) => (a.endLine - a.startLine) - (b.endLine - b.startLine))[0]; @@ -2561,7 +2647,12 @@ function isPythonLocallyBound(name: string, ref: UnresolvedRef, context: Resolut return false; } // Code only: `{% user_display user as user_display %}` in a docstring binds nothing. - const lines = stripCommentsForRegex(context.readFile(ref.filePath) ?? '', 'python').split(/\r?\n/); + let file = PY_LOCAL_FILE.get(context); + if (file?.filePath !== ref.filePath) { + const lines = stripCommentsForRegex(context.readFile(ref.filePath) ?? '', 'python').split(/\r?\n/); + PY_LOCAL_FILE.set(context, (file = { filePath: ref.filePath, lines, module: new Map() })); + } + const lines = file.lines; const n = name; const assigns = new RegExp(`^\\s*(?:[\\w\\s,*()\\[\\]]*,\\s*)?\\(?\\*?${n}\\)?\\s*(?:,[\\w\\s,*()\\[\\]]*)?(?::[^=]+)?=(?!=)`); const targets = new RegExp(`\\bfor\\s+[\\w\\s,()]*\\b${n}\\b[\\w\\s,()]*\\s+in\\b|\\bas\\s+${n}\\b`); @@ -2582,8 +2673,15 @@ function isPythonLocallyBound(name: string, ref: UnresolvedRef, context: Resolut } } // A module-level binding (`view = api_view(['GET'])(handler)`). - const top = new RegExp(`^(?:[\\w,\\s]*,\\s*)?${n}\\s*(?:,[\\w\\s,]*)?(?::[^=]+)?=(?!=)`); - for (let line = 0; !bound && line < lines.length; line++) bound = top.test(lines[line] ?? ''); + if (!bound) { + let module = file.module.get(n); + if (module === undefined) { + const top = new RegExp(`^(?:[\\w,\\s]*,\\s*)?${n}\\s*(?:,[\\w\\s,]*)?(?::[^=]+)?=(?!=)`); + module = lines.some(line => top.test(line)); + file.module.set(n, module); + } + bound = module; + } memo.set(key, bound); return bound; } @@ -6606,6 +6704,7 @@ export function clearNameMatcherMemos(context: ResolutionContext): void { PYTHON_STATEMENT_STARTS.delete(context); PYTHON_GLOBAL_CLASSES.delete(context); PYTHON_NAME_SCANS.delete(context); + PYTHON_IMPORTED_FILES.delete(context); AWAITED_TYPE_MEMO.delete(context); AWAITED_FILES.delete(context); C_STATIC_MEMO.delete(context); @@ -6664,6 +6763,8 @@ export function clearNameMatcherMemos(context: ResolutionContext): void { JVM_PACKAGES.delete(context); MINIFIED_SCRIPTS.delete(context); PY_LOCAL_BINDS.delete(context); + PY_LOCAL_FILE.delete(context); + PY_LOCAL_LAST.delete(context); OVERLOAD_SETS.delete(context); PHP_FILE_SCOPES.delete(context); JAVA_STATIC_IMPORTS.delete(context);