superwork/vscode/src/codeReader/rules/axum.ts

import type { ChainGroup, ChainStepFacts, FunctionBlock, Inline, KindBadge, LibraryRule, Node, PropertyItem, PropertyRow, RuleCard, RuleContext, ThemeColor } from '../contract'
import { returnNote, throwsRows, writtenParamRows } from './docRows'
import { chainRows, methodRuns, returnValue, valueCard } from './rustRows'
import { attributeName, docSummary, named, stringContent } from './rustTree'

/** 和 FastAPI 接口同一个蓝。 */
const COLOR: ThemeColor = { dark: '#61afef', light: '#1a6fb5' }
const BADGE: KindBadge = { label: '接口', color: COLOR, icon: 'globe' }
/** 和 FastAPI 路由器同一个蓝、同一个图标。 */
const ROUTER_BADGE: KindBadge = { label: '路由器', color: { dark: '#7aa2f7', light: '#34548a' }, icon: 'network' }
const METHODS = new Set(['get', 'post', 'put', 'delete', 'patch', 'head', 'options', 'trace'])

/** Router 链里这几个方法的叫法；连着调用几次的合成一步。连着的 `.route` 另合成路由表（见 [routerGroups]）。 */
const STEP_NAMES: Record<string, string> = {
  with_state: '共享状态',
  layer: '中间件',
  route_layer: '路由中间件',
  nest: '嵌套路由',
  nest_service: '嵌套服务',
  merge: '合并路由',
  fallback: '兜底处理',
}

/**
 * 处理函数上的 get / post 等路由属性；语句、变量的初始值、返回值是 Router 的链时排成「路由器」的链式调用卡片，每条路由都列出来
 * （见 [routerChainCard]）；只写了一步 `.route` 的照旧是那一个接口的卡片。请求表和 FastAPI 接口一样：方法、路径、路径参数，再加返回。
 * 文件里导入了 axum 才启用。
 */
export const axumRule: LibraryRule = {
  id: 'axum',
  label: 'axum 接口',
  languageIds: ['rust'],
  libraries: ['axum'],
  claim: (block, context) => {
    if (block.kind === 'function') {
      const route = attributeRoute(block)
      return route ? card(block, route, context) : undefined
    }
    return valueCard(block, context, (value, own) => {
      const chained = routerChainCard(value, context)
      if (chained)
        return chained
      const [route] = routesOn(value, own).filter(item => !handled(value, item.handler))
      return route ? routerCard(value, route, context) : undefined
    })
  },
  attributeNames: () => ({ attributes: {}, values: { kind: { 'axum.route': '接口', 'axum.router': '路由器' } } }),
}

/**
 * Router 的链（起点是 `Router::new()` / `Router::default()`，或链里有 `.route`）排成链式调用卡片：连着的 `.route` 合成路由表，
 * 别的几步按 [STEP_NAMES] 起叫法；不是这种链时没有。
 */
function routerChainCard(value: Node, context: RuleContext): RuleCard | undefined {
  const facts = context.document.language.chain?.(value)
  const start = facts?.start.text.replace(/\s+/g, '')
  if (!facts || !(/^(?:axum::)?Router::(?:new|default)\(\)$/.test(start ?? '') || facts.steps.some(step => step.method.text === 'route')))
    return undefined
  const view = context.chain(value, routerGroups(facts.steps, context))
  return view && { kind: 'axum.router', badge: ROUTER_BADGE, rows: chainRows(view), absorbs: [] }
}

/** 连着的 `.route` 合成「路由 N 个」表（路径 | 方法 | 处理函数，一个路径挂几个方法的各占一行），别的按 [STEP_NAMES] 起叫法。 */
function routerGroups(steps: readonly ChainStepFacts[], context: RuleContext): ChainGroup[] {
  return methodRuns(steps).flatMap(({ from, to, method }): ChainGroup[] => {
    if (method === 'route') {
      const items = steps.slice(from, to + 1).flatMap(step => routeItems(step, context))
      return [{ from, to, label: `路由 ${items.length} 个`, table: { columns: ['路径', '方法', '处理函数'], items } }]
    }
    const label = STEP_NAMES[method]
    return label ? [{ from, to, label }] : []
  })
}

/** 一步 `.route(路径, get(a).post(b))` 的各行：路径去掉引号、带源码下标，方法大写，处理函数是名字时按函数着色、带源码下标。 */
function routeItems(step: ChainStepFacts, context: RuleContext): PropertyItem[] {
  const [path, router] = step.arguments
  const content = path?.node.children.find(child => child?.type === 'string_content')
  const pathCell: Inline[] = content ? [{ text: content.text, role: 'string', offset: content.startIndex }] : path ? context.verbatim(path.node) : [hint('—')]
  return methodRoutes(router?.node).map(({ method, handler }) => ({
    cells: [
      pathCell,
      [keyword(method.toUpperCase())],
      handler?.type === 'identifier' ? [{ text: handler.text, role: 'function', offset: handler.startIndex }] : handler ? context.verbatim(handler) : [hint('—')],
    ],
  }))
}

/** 方法路由 `get(a)`、`get(a).post(b)` 里的各个方法和它的处理函数，照源码顺序；认不出时是空的。 */
function methodRoutes(node: Node | undefined): { method: string, handler?: Node }[] {
  if (node?.type !== 'call_expression')
    return []
  const fn = node.childForFieldName('function')
  const handler = named(node.childForFieldName('arguments')).find(arg => !arg.isExtra)
  if (fn?.type === 'field_expression') {
    const method = fn.childForFieldName('field')?.text
    return method && METHODS.has(method) ? [...methodRoutes(fn.childForFieldName('value') ?? undefined), { method, handler }] : []
  }
  const method = fn?.text.split('::').pop()
  return method && METHODS.has(method) ? [{ method, handler }] : []
}

interface Route {
  method: string
  path: string
  handler?: string
}

function attributeRoute(block: FunctionBlock): Route | undefined {
  for (const attr of block.facts.decorators) {
    const method = attributeName(attr)?.split('::').pop()
    if (!method || !METHODS.has(method))
      continue
    const path = attributeString(attr)
    if (path)
      return { method: method.toUpperCase(), path }
  }
  return undefined
}

function card(block: FunctionBlock, route: Route, context: RuleContext): RuleCard {
  const doc = block.facts.doc && context.doc(block.facts.doc)
  const summary = docSummary(doc)
  const params = pathParams(route.path, block.node, context)
  const rows: PropertyRow[] = [
    nameRow(block.facts.name),
    ...(summary ? [{ label: '说明', value: { markdown: summary } }] : []),
    ...(block.facts.async ? [{ label: '异步', value: [keyword('是')] }] : []),
    requestRow(route, params),
    ...writtenParamRows(doc?.params ?? [], '参数'),
    { label: '返回', value: returnValue(block.facts.returnType, context), ...returnNote(doc) },
    ...throwsRows(doc),
  ]
  return {
    kind: 'axum.route',
    badge: BADGE,
    rows,
    ...(block.facts.name && { defines: range(block.facts.name) }),
    absorbs: [],
  }
}

/** 同一文件里的处理函数已经标了路由属性时，Router 上这句不再单独成卡。 */
function handled(node: Node, handler: string | undefined): boolean {
  if (!handler)
    return false
  const root = node.tree.rootNode
  for (const item of named(root)) {
    if (item.type !== 'function_item' || item.childForFieldName('name')?.text !== handler)
      continue
    return itemAttributesHasRoute(item)
  }
  return false
}

function routerCard(node: Node, route: Route & { handlerNode?: Node }, context: RuleContext): RuleCard {
  const fn = route.handler ? functionNamed(node, route.handler) : undefined
  const params = fn ? pathParams(route.path, fn, context) : []
  const comment = fn && leadingDoc(fn)
  const doc = docSummary(comment && context.doc(comment))
  return {
    kind: 'axum.route',
    badge: BADGE,
    rows: [
      { label: '名称', value: route.handler ? [plain(route.handler)] : [hint('路由')] },
      ...(doc ? [{ label: '说明', value: { markdown: doc } }] : []),
      requestRow(route, params),
      { label: '返回', value: fn ? returnValue(fn.childForFieldName('return_type'), context) : [plain('—')] },
    ],
    absorbs: [],
  }
}

function requestRow(route: Route, params: Inline[][]): PropertyRow {
  const items: PropertyItem[] = [
    { cells: [[plain('方法')], [keyword(route.method)]] },
    { cells: [[plain('路径')], [{ text: route.path, role: 'string' }]] },
    ...(params.length > 0 ? [{ cells: [[plain('路径参数')], params.flatMap((part, index) => index === 0 ? part : [plain(', '), ...part])] }] : []),
  ]
  return {
    label: '请求',
    value: [keyword(route.method), plain(` ${route.path}`)],
    columns: ['项', '值'],
    items,
  }
}

/** 式子里写的各个 `.route(路径, 方法(函数))`，外层在前；拆成子块的那几段（`own` 为假）不看。 */
function routesOn(node: Node, own: (node: Node) => boolean, found: (Route & { handlerNode?: Node })[] = []): (Route & { handlerNode?: Node })[] {
  if (!own(node))
    return found
  if (node.type === 'call_expression') {
    const callee = node.childForFieldName('function')
    const field = callee?.type === 'field_expression' ? callee.childForFieldName('field')?.text : undefined
    if (field === 'route') {
      const args = named(node.childForFieldName('arguments'))
      const pathNode = args.find(arg => arg.type === 'string_literal')
      const call = args.find(arg => arg.type === 'call_expression')
      const method = call?.childForFieldName('function')?.text.split('::').pop()
      const handlerNode = named(call?.childForFieldName('arguments'))[0]
      const path = stringContent(pathNode)
      if (path && method && METHODS.has(method)) {
        found.push({
          method: method.toUpperCase(),
          path,
          ...(handlerNode?.type === 'identifier' && { handler: handlerNode.text, handlerNode }),
        })
      }
    }
  }
  for (const child of named(node))
    routesOn(child, own, found)
  return found
}

/** 路径里的 `{名字}` 对上函数的参数：`Path(id): Path<Uuid>` 写成 `id: Uuid`，类型照抄成一行、名字各自带源码下标。 */
function pathParams(path: string, fn: Node, context: RuleContext): Inline[][] {
  const names = [...path.matchAll(/\{(\w+)\}/g)].map(match => match[1] ?? '')
  return names.map((name) => {
    const type = innerType(parameterNamed(fn, name)?.childForFieldName('type'))
    return type ? [plain(`${name}: `), ...context.verbatim(type)] : [plain(name)]
  })
}

function parameterNamed(fn: Node, name: string): Node | undefined {
  const list = fn.childForFieldName('parameters')
  return named(list).find((param) => {
    if (param.type !== 'parameter' && param.type !== 'self_parameter')
      return false
    const pattern = param.childForFieldName('pattern') ?? param
    return namedDeep(pattern).some(node => (node.type === 'identifier' || node.type === 'self') && node.text === name)
  })
}

/** 路径参数的类型：`Path<Uuid>` 取尖括号里的第一个，别的写法就是它自己。 */
function innerType(type: Node | null | undefined): Node | undefined {
  if (type?.type === 'generic_type' && type.childForFieldName('type')?.text === 'Path')
    return named(type.childForFieldName('type_arguments'))[0] ?? type
  return type ?? undefined
}

function itemAttributesHasRoute(fn: Node): boolean {
  let prev = fn.previousNamedSibling
  while (prev?.type === 'attribute_item') {
    const method = attributeName(prev)?.split('::').pop()
    if (method && METHODS.has(method))
      return true
    prev = prev.previousNamedSibling
  }
  return false
}

function functionNamed(node: Node, name: string): Node | undefined {
  return named(node.tree.rootNode).find(item => item.type === 'function_item' && item.childForFieldName('name')?.text === name)
}

/** 函数上面的文档注释（中间可以隔着属性）：连着几行的 `///` 取第一条，读说明时从它往下连。 */

function leadingDoc(fn: Node): Node | undefined {
  let prev = fn.previousNamedSibling
  while (prev?.type === 'attribute_item')
    prev = prev.previousNamedSibling
  if (!prev || (prev.type !== 'line_comment' && prev.type !== 'block_comment'))
    return undefined
  let first = prev
  for (let before = first.previousNamedSibling; first.type === 'line_comment' && before?.type === 'line_comment' && before.startPosition.row === first.startPosition.row - 1; before = before.previousNamedSibling)
    first = before
  return first
}

function attributeString(attributeItem: Node): string | undefined {
  const attribute = named(attributeItem).find(child => child.type === 'attribute')
  const tree = attribute?.children.find(child => child?.type === 'token_tree')
  return stringContent(tree?.children.find(child => child?.type === 'string_literal') ?? undefined)
}

function namedDeep(node: Node): Node[] {
  return [node, ...named(node).flatMap(namedDeep)]
}

function nameRow(name: Node | undefined): PropertyRow {
  return {
    label: '名称',
    value: name ? [{ text: name.text, role: 'function', offset: name.startIndex }] : [hint('匿名')],
    ...(name && { defines: range(name) }),
  }
}

function keyword(text: string): Inline {
  return { text, role: 'keyword' }
}

function plain(text: string): Inline {
  return { text, role: 'plain' }
}

function hint(text: string): Inline {
  return { text, role: 'hint' }
}

function range(node: Node): { startIndex: number, endIndex: number } {
  return { startIndex: node.startIndex, endIndex: node.endIndex }
}