⚡ perf(vscode): 语言插件给词定词性时，在不在字符串 / 注解 / 注释里按语法树一遍收齐的范围查，不再逐词往上找父节点

vscode/src/codeReader/languages/bash.ts

@@ -2,7 +2,7 @@ import type { BlockShape, CodeBlock, CommandFacts, CommandWord, DeclaratorFacts,
 import type { CommentRules, ListBlocks } from './shared'
 import * as path from 'node:path'
 import { bashNameReader } from './bashNames'
-import { attachComments, children, compoundShape, leafBlock, leafShape, listBlocks, named } from './shared'
+import { attachComments, children, compoundShape, insideOf, leafBlock, leafShape, listBlocks, named } from './shared'
 
 /** 按习惯当常量用的名字：全大写（可带数字、下划线）。 */
 const CONSTANT_NAME = /^[A-Z][A-Z0-9_]*$/
@@ -543,15 +543,8 @@ function bashSymbols(root: Node): LocalSymbol[] {
 
 // ─── 词性 ────────────────────────────────────────────────────────────────────
 
-function insideString(leaf: Node): boolean {
-  for (let node = leaf.parent; node; node = node.parent) {
-    if (EXPANSIONS.has(node.type))
-      return false
-    if (node.type === 'string' || node.type === 'heredoc_body')
-      return true
-  }
-  return false
-}
+/** 字符串、here document 里的词：里面的展开（`$x`、`$(…)`）不算。 */
+const insideString = insideOf({ ...Object.fromEntries([...EXPANSIONS].map(type => [type, false])), string: true, heredoc_body: true })
 
 function bashTokenRole(leaf: Node): TokenRole {
   const parent = leaf.parent

vscode/src/codeReader/languages/python.ts

@@ -2,7 +2,7 @@ import type { Access, BlockShape, CallArgument, CallFacts, ClassFacts, CodeBlock
 import type { CommentRules, ExceptionRules } from './shared'
 import * as path from 'node:path'
 import { pythonNameReader } from './pythonNames'
-import { attachComments, children, compoundShape, leafBlock, leafShape, listBlocks, named, sameNode, syntaxChildren, withExceptions } from './shared'
+import { attachComments, children, compoundShape, insideOf, leafBlock, leafShape, listBlocks, named, sameNode, syntaxChildren, withExceptions } from './shared'
 
 /** 常量、类字段只认直接写在模块 / 类主体里的赋值；这两处的声明也才有访问级别。 */
 type Scope = 'module' | 'class' | 'other'
@@ -587,29 +587,14 @@ const OPERATOR_KEYWORDS = new Set(['and', 'or', 'not', 'is'])
 const COMPARISON = new Set(['comparison_operator', 'not in', 'is not'])
 const LITERAL_KEYWORDS = new Set(['none', 'true', 'false'])
 
-function insideString(leaf: Node): boolean {
-  for (let node = leaf.parent; node; node = node.parent) {
-    if (node.type === 'interpolation')
-      return false
-    if (node.type === 'string')
-      return true
-  }
-  return false
-}
+/** 字符串里的词：f-string 插值里的不算。 */
+const insideString = insideOf({ string: true, interpolation: false })
 
-/** 注解里的名字：typed_parameter / 返回值 / 带注解赋值的 type 节点之内。 */
-function insideAnnotation(leaf: Node): boolean {
-  for (let node = leaf.parent; node; node = node.parent) {
-    if (node.type === 'type')
-      return true
-    if (node.type === 'block' || node.type === 'module')
-      return false
-  }
-  return false
-}
+/** 注解里的名字：typed_parameter / 返回值 / 带注解赋值的 type 节点之内，往外找到语句块为止。 */
+const insideAnnotation = insideOf({ type: true, block: false, module: false })
 
-function isCallee(node: Node): boolean {
-  const parent = node.parent
+/** node 是不是被调用的那个；手上有它的父节点就传进来（找父节点要从根往下走一趟）。 */
+function isCallee(node: Node, parent: Node | null = node.parent): boolean {
   return !!parent && parent.type === 'call' && sameNode(parent.childForFieldName('function'), node)
 }
 
@@ -626,7 +611,7 @@ function identifierRole(leaf: Node): TokenRole {
   // 父类：类名后面括号里直接写的名字
   if (parent.type === 'argument_list' && parent.parent?.type === 'class_definition')
     return 'type'
-  if (parent.type === 'decorator' || isCallee(leaf))
+  if (parent.type === 'decorator' || isCallee(leaf, parent))
     return 'function'
   if (parent.type === 'attribute' && sameNode(parent.childForFieldName('attribute'), leaf) && isCallee(parent))
     return 'function'

vscode/src/codeReader/languages/rust.ts

@@ -2,7 +2,7 @@ import type { Access, AttributeFacts, AttributeItem, CfgPredicate, ClassFacts, C
 import type { CommentRules } from './shared'
 import * as path from 'node:path'
 import { leadingAttributes, visibilityFacts } from '../rules/rustTree'
-import { attachComments, children, compoundShape, extendHeader, headerSpan, isCommentNode, leafBlock, leafShape, listBlocks, named, syntaxChildren } from './shared'
+import { attachComments, children, compoundShape, extendHeader, headerSpan, insideOf, isCommentNode, leafBlock, leafShape, listBlocks, named, syntaxChildren } from './shared'
 
 const STORAGE = new Set(['pub', 'async', 'const', 'static', 'mut', 'let'])
 
@@ -393,10 +393,13 @@ function isAsync(node: Node): boolean {
   return modifiers?.children.some(child => child?.type === 'async') ?? false
 }
 
+const insideComment = insideOf({ line_comment: true, block_comment: true })
+const insideString = insideOf({ string_literal: true })
+
 function rustTokenRole(leaf: Node): TokenRole {
-  if (leaf.type === 'line_comment' || leaf.type === 'block_comment' || inside(leaf, 'line_comment') || inside(leaf, 'block_comment'))
+  if (leaf.type === 'line_comment' || leaf.type === 'block_comment' || insideComment(leaf))
     return 'comment'
-  if (inside(leaf, 'string_literal') || leaf.type === 'string_content')
+  if (insideString(leaf) || leaf.type === 'string_content')
     return 'string'
   if (leaf.type === 'integer_literal' || leaf.type === 'float_literal' || leaf.type === 'char_literal')
     return 'number'
@@ -412,13 +415,3 @@ function rustTokenRole(leaf: Node): TokenRole {
   }
   return 'plain'
 }
-
-function inside(node: Node, type: string): boolean {
-  let current = node.parent
-  while (current) {
-    if (current.type === type)
-      return true
-    current = current.parent
-  }
-  return false
-}

vscode/src/codeReader/languages/shared.test.ts

@@ -0,0 +1,106 @@
+import type { LanguageAdapter, Node } from '../contract'
+import { Node as TreeNode } from '@vscode/tree-sitter-wasm'
+import { describe, expect, it, vi } from 'vitest'
+import { readFixture, runtime, WASM_DIR } from '../testSupport'
+import { createBashAdapter } from './bash'
+import { createPythonAdapter } from './python'
+import { createRustAdapter } from './rust'
+import { insideOf } from './shared'
+import { createTypeScriptAdapter } from './typescript'
+
+function leavesOf(node: Node, out: Node[] = []): Node[] {
+  if (node.childCount === 0)
+    out.push(node)
+  for (const child of node.children) {
+    if (child)
+      leavesOf(child, out)
+  }
+  return out
+}
+
+/** 从父节点起一层层往上找，碰到的第一个列了的种类说了算。 */
+function walkUp(verdicts: Record<string, boolean>, leaf: Node): boolean {
+  for (let node = leaf.parent; node; node = node.parent) {
+    if (Object.hasOwn(verdicts, node.type))
+      return verdicts[node.type]
+  }
+  return false
+}
+
+async function verdictsOf(adapter: LanguageAdapter, languageId: string, source: string, verdicts: Record<string, boolean>): Promise<{ fast: boolean[], slow: boolean[] }> {
+  const tree = await runtime.parse(adapter, languageId, source)
+  try {
+    const leaves = leavesOf(tree.rootNode)
+    const inside = insideOf(verdicts)
+    return { fast: leaves.map(inside), slow: leaves.map(leaf => walkUp(verdicts, leaf)) }
+  }
+  finally {
+    tree.delete()
+  }
+}
+
+const PYTHON_STRINGS = [
+  'name: Annotated[str, Comment("名字")] = f"{user.name!r:>{width}} {f\'{inner["k"]}\'}"',
+  'def f(a: int, b: "Fwd" = None) -> dict[str, list[int]]:',
+  '    x: Optional[int] = None',
+  '    return f"{a}" "plain" rb\'\\x00\' + g(lambda y: f"{y}")',
+].join('\n')
+
+/* eslint-disable no-template-curly-in-string -- 源码样例里本来就有模板字符串、shell 展开 */
+const TYPESCRIPT_STRINGS = [
+  'const s = `a ${b + `c ${d} e`} f` + "g" + /h[i]/.source',
+  'const t: string = tag`x ${y} z`',
+].join('\n')
+
+const BASH_STRINGS = [
+  'echo "a $b ${c:-"d"} $(e "f") g" \'h\'',
+  'cat <<EOF',
+  'line $x $(y z)',
+  'EOF',
+].join('\n')
+/* eslint-enable no-template-curly-in-string */
+
+describe('insideOf：叶子在不在某种节点里面', () => {
+  it('python：和一层层往上找父节点一样（f-string 插值里的不算字符串、套着的 f-string 又算；注解到语句块为止）', async () => {
+    const py = createPythonAdapter(WASM_DIR)
+    for (const source of [PYTHON_STRINGS, readFixture('sample.py'), readFixture('routes.py')]) {
+      for (const verdicts of [{ string: true, interpolation: false }, { type: true, block: false, module: false }] as Record<string, boolean>[]) {
+        const { fast, slow } = await verdictsOf(py, 'python', source, verdicts)
+        expect(fast).toEqual(slow)
+      }
+    }
+    const { fast } = await verdictsOf(py, 'python', PYTHON_STRINGS, { string: true, interpolation: false })
+    expect(fast.filter(Boolean).length).toBeGreaterThan(0)
+  })
+
+  it('typescript、bash、rust：同样和一层层往上找一样', async () => {
+    const cases: [LanguageAdapter, string, string, Record<string, boolean>][] = [
+      [createTypeScriptAdapter(WASM_DIR), 'typescript', TYPESCRIPT_STRINGS, { string: true, template_string: true, regex: true, template_substitution: false }],
+      [createTypeScriptAdapter(WASM_DIR), 'typescriptreact', readFixture('posterDialog.tsx.txt'), { string: true, template_string: true, regex: true, template_substitution: false }],
+      [createBashAdapter(WASM_DIR), 'shellscript', BASH_STRINGS, { simple_expansion: false, expansion: false, command_substitution: false, string: true, heredoc_body: true }],
+      [createBashAdapter(WASM_DIR), 'shellscript', readFixture('fetchRuntime.sh'), { simple_expansion: false, expansion: false, command_substitution: false, string: true, heredoc_body: true }],
+      [createRustAdapter(WASM_DIR), 'rust', readFixture('rustComments.rs'), { line_comment: true, block_comment: true }],
+      [createRustAdapter(WASM_DIR), 'rust', readFixture('rustLang.rs'), { string_literal: true }],
+    ]
+    for (const [adapter, languageId, source, verdicts] of cases) {
+      const { fast, slow } = await verdictsOf(adapter, languageId, source, verdicts)
+      expect(fast).toEqual(slow)
+    }
+  })
+
+  it('一棵语法树一遍收齐：判定时不找父节点', async () => {
+    const py = createPythonAdapter(WASM_DIR)
+    const tree = await runtime.parse(py, 'python', readFixture('routes.py'))
+    try {
+      const leaves = leavesOf(tree.rootNode)
+      const parent = vi.spyOn(TreeNode.prototype, 'parent', 'get')
+      const inside = insideOf({ string: true, interpolation: false })
+      leaves.forEach(inside)
+      expect(parent).not.toHaveBeenCalled()
+      parent.mockRestore()
+    }
+    finally {
+      tree.delete()
+    }
+  })
+})

vscode/src/codeReader/languages/shared.ts

@@ -1,3 +1,4 @@
+import type { Tree } from '@vscode/tree-sitter-wasm'
 import type { BaseFacts, BlockShape, ClassBlock, ClassFacts, CodeBlock, CommentBlock, DocumentedBlock, ExceptionBlock, FieldFacts, Node, ParameterFacts, SourceSpan, TextKind } from '../contract'
 
 /** 语言插件共用的建树工具：取子节点、算头部、把一串语句切成块。 */
@@ -280,6 +281,81 @@ export function sameNode(a: Node | null | undefined, b: Node): boolean {
   return !!a && a.startIndex === b.startIndex && a.endIndex === b.endIndex && a.type === b.type
 }
 
+/** 一棵语法树上各段算不算在里面：at 从小到大，inside[i] 管从 at[i] 到下一个 at 之前。 */
+interface InsideSpans {
+  at: number[]
+  inside: boolean[]
+}
+
+/**
+ * 叶子在不在某种节点里面：verdicts 列出要看的节点种类，离叶子最近的那个祖先说了算（f-string 里插值中的名字不在字符串里），
+ * 祖先里一个都没有就不在。web-tree-sitter 找父节点要从根往下走一趟，逐个叶子一层层往上找很慢：
+ * 一棵语法树第一次问时一遍收齐这几种节点的范围，之后按叶子的起点二分查。
+ */
+export function insideOf(verdicts: Readonly<Record<string, boolean>>): (leaf: Node) => boolean {
+  const types = Object.keys(verdicts)
+  const known = new WeakMap<Tree, InsideSpans>()
+  return (leaf) => {
+    // 空的叶子（缺的记号、空的 here document 内容）贴着别的节点的边，按下标分不出在哪边，照父节点一层层找；这种叶子很少
+    if (leaf.startIndex === leaf.endIndex) {
+      for (let node = leaf.parent; node; node = node.parent) {
+        if (Object.hasOwn(verdicts, node.type))
+          return verdicts[node.type]
+      }
+      return false
+    }
+    let spans = known.get(leaf.tree)
+    if (!spans) {
+      spans = insideSpans(leaf.tree.rootNode, types, verdicts)
+      known.set(leaf.tree, spans)
+    }
+    return insideAt(spans, leaf.startIndex)
+  }
+}
+
+/** 这几种节点（有子节点的，没有子节点的当不了谁的祖先）按先序铺开：进一个节点从它的起点起按它算，出来后从它的终点起按外层的算。 */
+function insideSpans(root: Node, types: string[], verdicts: Readonly<Record<string, boolean>>): InsideSpans {
+  const spans: InsideSpans = { at: [], inside: [] }
+  const mark = (index: number, inside: boolean): void => {
+    if (spans.at.at(-1) === index) {
+      spans.inside[spans.inside.length - 1] = inside
+      return
+    }
+    spans.at.push(index)
+    spans.inside.push(inside)
+  }
+  const open: { end: number, inside: boolean }[] = []
+  const closeUntil = (index: number): void => {
+    while (open.length > 0 && open[open.length - 1].end <= index) {
+      const { end } = open.pop()!
+      mark(end, open.at(-1)?.inside ?? false)
+    }
+  }
+  for (const node of root.descendantsOfType(types)) {
+    if (!node || node.childCount === 0)
+      continue
+    closeUntil(node.startIndex)
+    const inside = verdicts[node.type]
+    open.push({ end: node.endIndex, inside })
+    mark(node.startIndex, inside)
+  }
+  closeUntil(Number.POSITIVE_INFINITY)
+  return spans
+}
+
+function insideAt({ at, inside }: InsideSpans, index: number): boolean {
+  let low = 0
+  let high = at.length
+  while (low < high) {
+    const mid = (low + high) >> 1
+    if (at[mid] <= index)
+      low = mid + 1
+    else
+      high = mid
+  }
+  return low > 0 && inside[low - 1]
+}
+
 // ─── 异常类 ───────────────────────────────────────────────────────────────────
 
 /** 语言怎么认异常类、类主体里哪些子块是类属性、哪个是构造函数。 */

vscode/src/codeReader/languages/typescript.ts

@@ -1,7 +1,7 @@
 import type { Access, BaseFacts, BlockShape, CallArgument, CallFacts, ClassFacts, CodeBlock, CollectionFacts, CommentBlock, DeclaratorFacts, ElementFacts, EntryFacts, EnumMemberFacts, FunctionFacts, ImportFacts, ImportItem, LanguageAdapter, Node, ParameterFacts, TextKind, TokenRole } from '../contract'
 import type { CommentRules, ExceptionRules } from './shared'
 import * as path from 'node:path'
-import { attachComments, children, compoundShape, extendHeader, headerSpan, leafBlock, leafShape, listBlocks, named, nodeSpan, sameNode, syntaxChildren, withExceptions } from './shared'
+import { attachComments, children, compoundShape, extendHeader, headerSpan, insideOf, leafBlock, leafShape, listBlocks, named, nodeSpan, sameNode, syntaxChildren, withExceptions } from './shared'
 import { typescriptNameReader } from './typescriptNames'
 
 const GRAMMARS: Record<string, string> = {
@@ -558,7 +558,6 @@ const FLOW_KEYWORDS = new Set([
   'yield',
 ])
 const LITERAL_KEYWORDS = new Set(['true', 'false', 'null', 'undefined', 'this', 'super'])
-const STRING_NODES = new Set(['string', 'template_string', 'regex'])
 const FUNCTION_NAMED = new Set(['function_declaration', 'generator_function_declaration', 'function_expression', 'function', 'generator_function', 'method_definition', 'method_signature', 'abstract_method_signature', 'function_signature'])
 
 /** 调用的结构：`new X(…)` 是新建对象，别的调用不是；带标签的模板字符串（sql`…`）、不带括号的 `new X` 不算。 */
@@ -581,18 +580,11 @@ function typescriptCall(node: Node): CallFacts | undefined {
   }
 }
 
-function insideString(leaf: Node): boolean {
-  for (let node = leaf.parent; node; node = node.parent) {
-    if (node.type === 'template_substitution')
-      return false
-    if (STRING_NODES.has(node.type))
-      return true
-  }
-  return false
-}
+/** 字符串、模板字符串、正则里的词：模板字符串 `${}` 里的不算。 */
+const insideString = insideOf({ string: true, template_string: true, regex: true, template_substitution: false })
 
-function isCallee(node: Node): boolean {
-  const parent = node.parent
+/** node 是不是被调用的那个；手上有它的父节点就传进来（找父节点要从根往下走一趟）。 */
+function isCallee(node: Node, parent: Node | null = node.parent): boolean {
   return !!parent && parent.type === 'call_expression' && sameNode(parent.childForFieldName('function'), node)
 }
 
@@ -602,7 +594,7 @@ function nameRole(leaf: Node): TokenRole {
     return 'variable'
   if (FUNCTION_NAMED.has(parent.type) && sameNode(parent.childForFieldName('name'), leaf))
     return 'function'
-  if (isCallee(leaf))
+  if (isCallee(leaf, parent))
     return 'function'
   if (parent.type === 'member_expression' && sameNode(parent.childForFieldName('property'), leaf) && isCallee(parent))
     return 'function'