superwork/vscode/src/codeReader/languages/rust.ts

import type { Access, ClassFacts, CodeBlock, EnumBlock, FunctionBlock, ImplBlock, ImportBlock, ImportItem, LanguageAdapter, Node, ParameterFacts, TokenRole, TraitBlock } from '../contract'
import type { CommentRules } from './shared'
import * as path from 'node:path'
import { leadingAttributes, visibilityFacts } from '../rules/rustTree'
import { compoundShape, extendHeader, headerSpan, leafBlock, leafShape, listBlocks, named } from './shared'

const STORAGE = new Set(['pub', 'async', 'const', 'static', 'mut', 'let'])

export function createRustAdapter(wasmDir: string): LanguageAdapter {
  return {
    id: 'rust',
    languageIds: ['rust'],
    grammarPath: () => path.join(wasmDir, 'tree-sitter-rust.wasm'),
    toBlocks: root => topBlocks(named(root), -1),
    tokenRole: rustTokenRole,
  }
}

/** extern 块本身不成一块源码，里面的签名跟所在这一层的别的项走同一条路。 */
function unwrapExtern(item: Node): Node[] {
  if (item.type !== 'foreign_mod_item')
    return [item]
  const body = item.childForFieldName('body')
  return body ? named(body).flatMap(unwrapExtern) : []
}

/**
 * `//` 普通注释、讲所在模块的 `//!`、讲下一项的 `///` 各是一种写法，只和同一种连成一块。
 * `///` 和 `/**` 开头的块注释是外部文档注释，讲紧跟着的那一项（中间可以隔着 `#[…]`）。
 */
const COMMENTS: CommentRules = {
  lineStyle: comment => comment.type === 'line_comment' ? `//${docMarker(comment)}` : undefined,
  isDoc: comment => named(comment).some(child => child.type === 'outer_doc_comment_marker'),
  documents: new Set(['function', 'class', 'trait', 'impl', 'module', 'enum', 'const', 'variable']),
}

/** 文档注释记号：外部的 `/`（块注释是 `*`）、内部的 `!`；普通注释没有。 */
function docMarker(comment: Node): string {
  return named(comment).find(child => child.type === 'outer_doc_comment_marker' || child.type === 'inner_doc_comment_marker')?.text ?? ''
}

/** 一层的各项（文件顶层、模块身体、impl 身体）→ 块；afterRow 同 listBlocks。`#[…]` 和多出来的空 `;` 不出块。 */
function topBlocks(items: Node[], afterRow: number): CodeBlock[] {
  return listBlocks(items.flatMap(unwrapExtern), afterRow, (item) => {
    if (item.type === 'attribute_item' || item.type === 'empty_statement')
      return []
    return [withAttributes(itemBlock(item))]
  }, COMMENTS)
}

function itemBlock(item: Node): CodeBlock {
  switch (item.type) {
    case 'use_declaration':
      return useBlock(item)
    case 'struct_item':
      return structBlock(item)
    case 'enum_item':
      return enumBlock(item)
    case 'impl_item':
      return implBlock(item)
    case 'trait_item':
      return traitBlock(item)
    case 'mod_item':
      return moduleBlock(item)
    case 'const_item':
      return valueBlock(item, 'const')
    case 'static_item':
      return valueBlock(item, 'variable')
    case 'function_item':
    case 'function_signature_item':
      return functionBlock(item)
    default:
      return leafBlock('statement', item)
  }
}

/** 没有身体的 `mod 名字;` 不拆。有身体的模块，身体和文件顶层同一套规则。 */
function moduleBlock(node: Node): CodeBlock {
  const body = node.childForFieldName('body')
  if (!body)
    return leafBlock('statement', node)
  const { access, attributes } = visibilityFacts(node)
  return {
    kind: 'module',
    ...compoundShape(node, body.startIndex, named(body), topBlocks),
    ...(attributes && { attributes }),
    facts: classFacts(node, access),
  }
}

/** const 是常量，static 是变量；名称、类型、值交给现成的常量 / 变量属性表。 */
function valueBlock(node: Node, kind: 'const' | 'variable'): CodeBlock {
  const name = node.childForFieldName('name')
  if (!name)
    return leafBlock('statement', node)
  const type = node.childForFieldName('type') ?? undefined
  const value = node.childForFieldName('value') ?? undefined
  const { access, attributes } = visibilityFacts(node)
  return {
    kind,
    ...leafShape(node),
    ...(attributes && { attributes }),
    facts: { access, declarators: [{ name, type, value }] },
  }
}

/** 写在前面的 `#[…]` 算进头部。 */
function withAttributes(block: CodeBlock): CodeBlock {
  const [attribute] = leadingAttributes(block.node)
  return attribute && block.header.startIndex > attribute.startIndex ? extendHeader(block, attribute) : block
}

function useBlock(node: Node): ImportBlock {
  const argument = node.childForFieldName('argument')
  const { source, items } = argument ? useParts(argument) : { items: [] }
  return { kind: 'import', ...leafShape(node), facts: { sources: source ? [source] : [], items } }
}

function useParts(node: Node): { source?: Node, items: ImportItem[] } {
  switch (node.type) {
    case 'scoped_identifier': {
      const source = node.childForFieldName('path') ?? undefined
      const name = node.childForFieldName('name')
      return { source, items: name ? [{ kind: 'named', name }] : [] }
    }
    case 'scoped_use_list': {
      const source = node.childForFieldName('path') ?? undefined
      const list = node.childForFieldName('list')
      return { source, items: list ? useList(list) : [] }
    }
    case 'use_list':
      return { items: useList(node) }
    case 'use_as_clause':
      return { items: asItem(node) }
    case 'use_wildcard':
      return { items: [{ kind: 'namespace' }] }
    case 'identifier':
      return { items: [{ kind: 'named', name: node }] }
    default:
      return { items: [] }
  }
}

function useList(list: Node): ImportItem[] {
  return named(list).flatMap((child) => {
    if (child.type === 'identifier')
      return [{ kind: 'named' as const, name: child }]
    if (child.type === 'use_as_clause')
      return asItem(child)
    if (child.type === 'use_wildcard')
      return [{ kind: 'namespace' as const }]
    if (child.type === 'scoped_identifier') {
      const name = child.childForFieldName('name')
      return name ? [{ kind: 'named' as const, name }] : []
    }
    return []
  })
}

function asItem(node: Node): ImportItem[] {
  const name = node.childForFieldName('path')
  const alias = node.childForFieldName('alias') ?? undefined
  return name ? [{ kind: 'named', name, alias }] : []
}

function structBlock(node: Node): CodeBlock {
  const { access, attributes } = visibilityFacts(node)
  return {
    kind: 'class',
    ...shape(node),
    ...(attributes && { attributes }),
    facts: classFacts(node, access),
  }
}

function enumBlock(node: Node): EnumBlock {
  const { access, attributes } = visibilityFacts(node)
  const body = node.childForFieldName('body')
  const members = named(body).flatMap((variant) => {
    if (variant.type !== 'enum_variant')
      return []
    const name = variant.childForFieldName('name')
    if (!name)
      return []
    const payload = variant.childForFieldName('body') ?? undefined
    return [{ name, ...(payload && { value: payload }), docs: [] }]
  })
  return {
    kind: 'enum',
    ...shape(node),
    ...(attributes && { attributes }),
    facts: { name: node.childForFieldName('name') ?? node, doc: undefined, access, decorators: [], bases: [], members },
  }
}

function traitBlock(node: Node): TraitBlock {
  const { access, attributes } = visibilityFacts(node)
  return { kind: 'trait', ...shape(node), ...(attributes && { attributes }), facts: classFacts(node, access) }
}

function implBlock(node: Node): ImplBlock {
  const { access, attributes } = visibilityFacts(node)
  const body = node.childForFieldName('body')
  return {
    kind: 'impl',
    ...compoundShape(node, body?.startIndex ?? node.endIndex, named(body), topBlocks),
    ...(attributes && { attributes }),
    facts: classFacts(node, access, implType(node)),
  }
}

function functionBlock(node: Node): FunctionBlock {
  const { access, attributes } = visibilityFacts(node)
  const body = node.childForFieldName('body')
  const headerEnd = body?.startIndex ?? node.endIndex
  return {
    kind: 'function',
    node,
    header: headerSpan(node, headerEnd),
    children: [],
    ...(attributes && { attributes }),
    facts: {
      name: node.childForFieldName('name') ?? undefined,
      access,
      ...(isAsync(node) && { async: true }),
      decorators: leadingAttributes(node),
      parameters: parametersOf(node.childForFieldName('parameters')),
      returnType: node.childForFieldName('return_type') ?? undefined,
    },
  }
}

function classFacts(node: Node, access: Access | undefined, name = node.childForFieldName('name') ?? undefined): ClassFacts {
  return { name, access, decorators: leadingAttributes(node), typeParameters: [], bases: [], implements: [] }
}

/** impl 的自己的类型：impl Trait for Type 取 Type，impl Type 取 Type。 */
function implType(node: Node): Node | undefined {
  return node.childForFieldName('type') ?? node.childForFieldName('trait') ?? undefined
}

function shape(node: Node): { node: Node, header: { startIndex: number, endIndex: number, startRow: number, endRow: number }, children: CodeBlock[] } {
  const body = node.childForFieldName('body')
  return body ? { node, header: headerSpan(node, body.startIndex), children: [] } : { ...leafShape(node) }
}

function parametersOf(list: Node | null): ParameterFacts[] {
  if (!list)
    return []
  return named(list).flatMap((child) => {
    if (child.type === 'self_parameter') {
      const name = [...child.children].reverse().find(token => token?.type === 'self')
      return name ? [{ name, optional: false, type: child }] : []
    }
    if (child.type !== 'parameter')
      return []
    const name = child.childForFieldName('pattern')
    const type = child.childForFieldName('type') ?? undefined
    return name ? [{ name, optional: false, type }] : []
  })
}

function isAsync(node: Node): boolean {
  const modifiers = named(node).find(child => child.type === 'function_modifiers')
  return modifiers?.children.some(child => child?.type === 'async') ?? false
}

function rustTokenRole(leaf: Node): TokenRole {
  if (leaf.type === 'line_comment' || leaf.type === 'block_comment' || inside(leaf, 'line_comment') || inside(leaf, 'block_comment'))
    return 'comment'
  if (inside(leaf, 'string_literal') || leaf.type === 'string_content')
    return 'string'
  if (leaf.type === 'integer_literal' || leaf.type === 'float_literal' || leaf.type === 'char_literal')
    return 'number'
  if (leaf.type === 'primitive_type' || leaf.type === 'type_identifier')
    return 'type'
  if (leaf.type === 'mutable_specifier')
    return 'keyword'
  if (!leaf.isNamed)
    return STORAGE.has(leaf.type) ? 'storage' : /^[a-z]+$/.test(leaf.type) ? 'keyword' : 'plain'
  if (leaf.type === 'identifier' || leaf.type === 'field_identifier' || leaf.type === 'self') {
    const parent = leaf.parent?.type
    return parent === 'function_item' || parent === 'function_signature_item' ? 'function' : 'variable'
  }
  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
}