codegraph/src/extraction/languages/rust.ts

import type { Node as SyntaxNode } from 'web-tree-sitter';
import { getNodeText, getChildByField } from '../tree-sitter-helpers';
import type { LanguageExtractor } from '../tree-sitter-types';

/**
 * A Rust function's declared return type, normalized to the bare type a chained
 * `Foo::new().bar()` could be called on (the #645/#608 mechanism). Reads the
 * `return_type` field: `-> Self` yields the marker `self` (resolved to the impl's
 * own type at resolution time, like PHP's `self`/`static`); a concrete `-> Foo` /
 * `-> FooBuilder` its name; a reference (`&Foo`) is unwrapped; generics are reduced
 * to the base type (`Vec<Foo>` → `Vec`); primitives / unit / tuple yield undefined.
 * Stdlib types that aren't in the graph simply fail the later existence check.
 */
function extractRustReturnType(node: SyntaxNode, source: string): string | undefined {
  let rt = getChildByField(node, 'return_type');
  if (!rt) return undefined;
  if (rt.type === 'reference_type') {
    rt =
      rt.namedChildren.find(
        (c: SyntaxNode) =>
          c.type === 'type_identifier' ||
          c.type === 'scoped_type_identifier' ||
          c.type === 'generic_type',
      ) ?? rt;
  }
  if (!rt || rt.type === 'primitive_type' || rt.type === 'unit_type' || rt.type === 'tuple_type') {
    return undefined;
  }
  const text = getNodeText(rt, source).trim().replace(/<[^>]*>/g, '');
  const last = text.split('::').pop()?.trim();
  if (!last || !/^[A-Za-z_]\w*$/.test(last)) return undefined;
  return last === 'Self' ? 'self' : last;
}

export const rustExtractor: LanguageExtractor = {
  // `function_signature_item` is a trait method DECLARATION (`fn render(&self);`,
  // no body). Extracting it makes a trait's method set first-class, which
  // impl-navigation and trait-dispatch synthesis need (a struct's method set is
  // matched against the trait's).
  functionTypes: ['function_item', 'function_signature_item'],
  classTypes: [], // Rust has impl blocks
  methodTypes: ['function_item', 'function_signature_item'],
  interfaceTypes: ['trait_item'],
  structTypes: ['struct_item'],
  enumTypes: ['enum_item'],
  enumMemberTypes: ['enum_variant'],
  typeAliasTypes: ['type_item'], // Rust type aliases
  importTypes: ['use_declaration'],
  callTypes: ['call_expression'],
  variableTypes: ['let_declaration', 'const_item', 'static_item'],
  interfaceKind: 'trait',
  nameField: 'name',
  bodyField: 'body',
  paramsField: 'parameters',
  returnField: 'return_type',
  getReturnType: extractRustReturnType,
  getSignature: (node, source) => {
    const params = getChildByField(node, 'parameters');
    const returnType = getChildByField(node, 'return_type');
    if (!params) return undefined;
    let sig = getNodeText(params, source);
    if (returnType) {
      sig += ' -> ' + getNodeText(returnType, source);
    }
    return sig;
  },
  isAsync: (node) => {
    for (let i = 0; i < node.childCount; i++) {
      const child = node.child(i);
      if (child?.type === 'async') return true;
    }
    return false;
  },
  getVisibility: (node) => {
    for (let i = 0; i < node.childCount; i++) {
      const child = node.child(i);
      if (child?.type === 'visibility_modifier') {
        return child.text.includes('pub') ? 'public' : 'private';
      }
    }
    return 'private'; // Rust defaults to private
  },
  getReceiverType: (node, source) => {
    // Walk up the tree-sitter AST to find a parent impl_item
    let parent = node.parent;
    while (parent) {
      if (parent.type === 'impl_item') {
        // For `impl Type { ... }` — the type is a direct type_identifier child
        // For `impl Trait for Type { ... }` — the type is the LAST type_identifier
        // (the first is part of the trait path)
        const children = parent.namedChildren;
        // Find all direct type_identifier children (not nested in scoped paths)
        const typeIdents = children.filter(
          (c: SyntaxNode) => c.type === 'type_identifier'
        );
        if (typeIdents.length > 0) {
          // Last type_identifier is always the implementing type
          const typeNode = typeIdents[typeIdents.length - 1]!;
          return source.substring(typeNode.startIndex, typeNode.endIndex);
        }
        // Handle generic types: impl<T> MyStruct<T> { ... }
        const genericType = children.find(
          (c: SyntaxNode) => c.type === 'generic_type'
        );
        if (genericType) {
          const innerType = genericType.namedChildren.find(
            (c: SyntaxNode) => c.type === 'type_identifier'
          );
          if (innerType) {
            return source.substring(innerType.startIndex, innerType.endIndex);
          }
        }
        return undefined;
      }
      parent = parent.parent;
    }
    return undefined;
  },

  extractImport: (node, source) => {
    const importText = source.substring(node.startIndex, node.endIndex).trim();

    // Helper to get the root crate/module from a scoped path
    const getRootModule = (scopedNode: SyntaxNode): string => {
      const firstChild = scopedNode.namedChild(0);
      if (!firstChild) return source.substring(scopedNode.startIndex, scopedNode.endIndex);
      if (firstChild.type === 'identifier' ||
          firstChild.type === 'crate' ||
          firstChild.type === 'super' ||
          firstChild.type === 'self') {
        return source.substring(firstChild.startIndex, firstChild.endIndex);
      } else if (firstChild.type === 'scoped_identifier') {
        return getRootModule(firstChild);
      }
      return source.substring(firstChild.startIndex, firstChild.endIndex);
    };

    // Find the use argument (scoped_use_list or scoped_identifier)
    const useArg = node.namedChildren.find((c: SyntaxNode) =>
      c.type === 'scoped_use_list' ||
      c.type === 'scoped_identifier' ||
      c.type === 'use_list' ||
      c.type === 'identifier'
    );

    if (useArg) {
      return { moduleName: getRootModule(useArg), signature: importText };
    }
    return null;
  },
};