wehub-resource-sync
411 lines
19 KiB
Python
411 lines
19 KiB
Python
"""Rust extractor. Moved verbatim from graphify/extract.py."""
|
|
from __future__ import annotations
|
|
|
|
|
|
from pathlib import Path
|
|
from graphify.extractors.base import _LANGUAGE_BUILTIN_GLOBALS, _file_stem, _make_id, _read_text
|
|
|
|
|
|
def _rust_collect_type_refs(node, source: bytes, generic: bool, out: list[tuple[str, str]]) -> None:
|
|
"""Walk a Rust type expression; append (name, role) tuples."""
|
|
if node is None:
|
|
return
|
|
t = node.type
|
|
if t == "primitive_type":
|
|
return
|
|
if t == "type_identifier":
|
|
text = _read_text(node, source)
|
|
if text:
|
|
out.append((text, "generic_arg" if generic else "type"))
|
|
return
|
|
if t == "scoped_type_identifier":
|
|
text = _read_text(node, source).rsplit("::", 1)[-1]
|
|
if text:
|
|
out.append((text, "generic_arg" if generic else "type"))
|
|
return
|
|
if t == "generic_type":
|
|
name_node = node.child_by_field_name("type")
|
|
if name_node is None:
|
|
for c in node.children:
|
|
if c.type in ("type_identifier", "scoped_type_identifier"):
|
|
name_node = c
|
|
break
|
|
if name_node is not None:
|
|
text = _read_text(name_node, source).rsplit("::", 1)[-1]
|
|
if text:
|
|
out.append((text, "generic_arg" if generic else "type"))
|
|
for c in node.children:
|
|
if c.type == "type_arguments":
|
|
for arg in c.children:
|
|
if arg.is_named:
|
|
_rust_collect_type_refs(arg, source, True, out)
|
|
return
|
|
if t in ("reference_type", "pointer_type", "array_type", "tuple_type", "slice_type"):
|
|
for c in node.children:
|
|
if c.is_named:
|
|
_rust_collect_type_refs(c, source, generic, out)
|
|
return
|
|
if node.is_named:
|
|
for c in node.children:
|
|
if c.is_named:
|
|
_rust_collect_type_refs(c, source, generic, out)
|
|
|
|
_RUST_TRAIT_METHOD_BLOCKLIST: frozenset[str] = frozenset({
|
|
"new", "default", "parse", "from_str", "now", "clone", "into", "from",
|
|
"to_string", "to_owned", "len", "is_empty", "iter", "next", "build",
|
|
"start", "run", "init", "app", "get", "set", "push", "pop", "insert",
|
|
"remove", "contains", "collect", "map", "filter", "unwrap", "expect",
|
|
"ok", "err", "some", "none", "send", "recv", "lock", "read", "write",
|
|
})
|
|
|
|
def extract_rust(path: Path) -> dict:
|
|
"""Extract functions, structs, enums, traits, impl methods, and use declarations from a .rs file."""
|
|
try:
|
|
import tree_sitter_rust as tsrust
|
|
from tree_sitter import Language, Parser
|
|
except ImportError:
|
|
return {"nodes": [], "edges": [], "error": "tree-sitter-rust not installed"}
|
|
|
|
try:
|
|
language = Language(tsrust.language())
|
|
parser = Parser(language)
|
|
source = path.read_bytes()
|
|
tree = parser.parse(source)
|
|
root = tree.root_node
|
|
except Exception as e:
|
|
return {"nodes": [], "edges": [], "error": str(e)}
|
|
|
|
stem = _file_stem(path)
|
|
str_path = str(path)
|
|
nodes: list[dict] = []
|
|
edges: list[dict] = []
|
|
seen_ids: set[str] = set()
|
|
function_bodies: list[tuple[str, object]] = []
|
|
|
|
def add_node(nid: str, label: str, line: int) -> None:
|
|
if nid not in seen_ids:
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": label,
|
|
"file_type": "code",
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
})
|
|
|
|
def add_edge(src: str, tgt: str, relation: str, line: int,
|
|
confidence: str = "EXTRACTED", weight: float = 1.0,
|
|
context: str | None = None) -> None:
|
|
edge = {
|
|
"source": src,
|
|
"target": tgt,
|
|
"relation": relation,
|
|
"confidence": confidence,
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
"weight": weight,
|
|
}
|
|
if context:
|
|
edge["context"] = context
|
|
edges.append(edge)
|
|
|
|
file_nid = _make_id(str(path))
|
|
add_node(file_nid, path.name, 1)
|
|
|
|
def ensure_named_node(name: str, line: int) -> str:
|
|
nid = _make_id(stem, name)
|
|
if nid in seen_ids:
|
|
return nid
|
|
nid = _make_id(name)
|
|
if nid not in seen_ids:
|
|
# The name isn't defined in this file, so this is a cross-file reference
|
|
# (e.g. a `Thing` type annotation imported from another module). Emit a
|
|
# SOURCELESS stub — like the inheritance-base path below — so the
|
|
# corpus-level rewire can collapse it onto the real definition. A sourced
|
|
# stub here makes _disambiguate_colliding_node_ids bake the referencing
|
|
# file's path (with extension) into the id and blocks the rewire, which is
|
|
# the phantom-duplicate-node bug (#1402).
|
|
seen_ids.add(nid)
|
|
nodes.append({
|
|
"id": nid,
|
|
"label": name,
|
|
"file_type": "code",
|
|
"source_file": "",
|
|
"source_location": "",
|
|
"origin_file": str_path,
|
|
})
|
|
return nid
|
|
|
|
def emit_param_return_refs(func_node, func_nid: str, line: int) -> None:
|
|
params = func_node.child_by_field_name("parameters")
|
|
if params is not None:
|
|
for p in params.children:
|
|
if p.type != "parameter":
|
|
continue
|
|
type_node = p.child_by_field_name("type")
|
|
refs: list[tuple[str, str]] = []
|
|
_rust_collect_type_refs(type_node, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "parameter_type"
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt != func_nid:
|
|
add_edge(func_nid, tgt, "references", line, context=ctx)
|
|
return_type = func_node.child_by_field_name("return_type")
|
|
if return_type is not None:
|
|
refs = []
|
|
_rust_collect_type_refs(return_type, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "return_type"
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt != func_nid:
|
|
add_edge(func_nid, tgt, "references", line, context=ctx)
|
|
|
|
def walk(node, parent_impl_nid: str | None = None) -> None:
|
|
t = node.type
|
|
|
|
if t == "function_item":
|
|
name_node = node.child_by_field_name("name")
|
|
if name_node:
|
|
func_name = _read_text(name_node, source)
|
|
line = node.start_point[0] + 1
|
|
if parent_impl_nid:
|
|
func_nid = _make_id(parent_impl_nid, func_name)
|
|
add_node(func_nid, f".{func_name}()", line)
|
|
add_edge(parent_impl_nid, func_nid, "method", line)
|
|
else:
|
|
func_nid = _make_id(stem, func_name)
|
|
add_node(func_nid, f"{func_name}()", line)
|
|
add_edge(file_nid, func_nid, "contains", line)
|
|
emit_param_return_refs(node, func_nid, line)
|
|
body = node.child_by_field_name("body")
|
|
if body:
|
|
function_bodies.append((func_nid, body))
|
|
return
|
|
|
|
if t in ("struct_item", "enum_item", "trait_item"):
|
|
name_node = node.child_by_field_name("name")
|
|
if name_node:
|
|
item_name = _read_text(name_node, source)
|
|
line = node.start_point[0] + 1
|
|
item_nid = _make_id(stem, item_name)
|
|
add_node(item_nid, item_name, line)
|
|
add_edge(file_nid, item_nid, "contains", line)
|
|
if t == "trait_item":
|
|
for c in node.children:
|
|
if c.type != "trait_bounds":
|
|
continue
|
|
for sub in c.children:
|
|
if not sub.is_named:
|
|
continue
|
|
refs: list[tuple[str, str]] = []
|
|
_rust_collect_type_refs(sub, source, False, refs)
|
|
for idx, (ref_name, _role) in enumerate(refs):
|
|
tgt = ensure_named_node(ref_name, line)
|
|
if tgt == item_nid:
|
|
continue
|
|
rel = "inherits" if idx == 0 else "references"
|
|
if rel == "inherits":
|
|
add_edge(item_nid, tgt, "inherits", line)
|
|
else:
|
|
add_edge(item_nid, tgt, "references", line,
|
|
context="generic_arg")
|
|
if t == "struct_item":
|
|
for c in node.children:
|
|
if c.type != "field_declaration_list":
|
|
continue
|
|
for field in c.children:
|
|
if field.type != "field_declaration":
|
|
continue
|
|
type_node = field.child_by_field_name("type")
|
|
if type_node is None:
|
|
for fc in field.children:
|
|
if fc.type in ("type_identifier", "generic_type",
|
|
"scoped_type_identifier",
|
|
"reference_type", "primitive_type"):
|
|
type_node = fc
|
|
break
|
|
refs = []
|
|
_rust_collect_type_refs(type_node, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "field"
|
|
tgt = ensure_named_node(ref_name, field.start_point[0] + 1)
|
|
if tgt != item_nid:
|
|
add_edge(item_nid, tgt, "references",
|
|
field.start_point[0] + 1, context=ctx)
|
|
# Tuple structs (`struct Wrapper(pub Logger, Config);`) nest their
|
|
# positional field types directly under ordered_field_declaration_list
|
|
# with no field_declaration wrapper -- the same shape handled for tuple
|
|
# enum variants below. Without this branch these field type references
|
|
# are silently dropped.
|
|
for c in node.children:
|
|
if c.type != "ordered_field_declaration_list":
|
|
continue
|
|
fline = c.start_point[0] + 1
|
|
for tc in c.children:
|
|
if tc.type not in ("type_identifier", "generic_type",
|
|
"scoped_type_identifier", "reference_type",
|
|
"primitive_type", "tuple_type", "array_type"):
|
|
continue
|
|
refs = []
|
|
_rust_collect_type_refs(tc, source, False, refs)
|
|
for ref_name, role in refs:
|
|
ctx = "generic_arg" if role == "generic_arg" else "field"
|
|
tgt = ensure_named_node(ref_name, fline)
|
|
if tgt != item_nid:
|
|
add_edge(item_nid, tgt, "references", fline, context=ctx)
|
|
if t == "enum_item":
|
|
# Variant payload types nest under enum_variant_list ->
|
|
# enum_variant -> ordered_field_declaration_list (tuple variant,
|
|
# `Click(Logger)`) | field_declaration_list (struct variant,
|
|
# `Resize { size: Dim }`). Neither was traversed, so every
|
|
# enum-variant type reference was silently dropped.
|
|
_TYPE_NODES = ("type_identifier", "generic_type",
|
|
"scoped_type_identifier", "reference_type",
|
|
"primitive_type", "tuple_type", "array_type")
|
|
|
|
def _emit_enum_type(type_node, at_line):
|
|
if type_node is None:
|
|
return
|
|
refs2: list[tuple[str, str]] = []
|
|
_rust_collect_type_refs(type_node, source, False, refs2)
|
|
for ref_name, role in refs2:
|
|
ctx = "generic_arg" if role == "generic_arg" else "field"
|
|
tgt = ensure_named_node(ref_name, at_line)
|
|
if tgt != item_nid:
|
|
add_edge(item_nid, tgt, "references", at_line, context=ctx)
|
|
|
|
for c in node.children:
|
|
if c.type != "enum_variant_list":
|
|
continue
|
|
for variant in c.children:
|
|
if variant.type != "enum_variant":
|
|
continue
|
|
vline = variant.start_point[0] + 1
|
|
for vc in variant.children:
|
|
if vc.type == "ordered_field_declaration_list":
|
|
for tc in vc.children:
|
|
if tc.type in _TYPE_NODES:
|
|
_emit_enum_type(tc, vline)
|
|
elif vc.type == "field_declaration_list":
|
|
for field in vc.children:
|
|
if field.type != "field_declaration":
|
|
continue
|
|
type_node = field.child_by_field_name("type")
|
|
_emit_enum_type(type_node, field.start_point[0] + 1)
|
|
return
|
|
|
|
if t == "impl_item":
|
|
type_node = node.child_by_field_name("type")
|
|
trait_node = node.child_by_field_name("trait")
|
|
impl_nid: str | None = None
|
|
if type_node:
|
|
type_name = _read_text(type_node, source).strip()
|
|
impl_nid = _make_id(stem, type_name)
|
|
add_node(impl_nid, type_name, node.start_point[0] + 1)
|
|
if trait_node is not None and impl_nid is not None:
|
|
refs: list[tuple[str, str]] = []
|
|
_rust_collect_type_refs(trait_node, source, False, refs)
|
|
for idx, (ref_name, _role) in enumerate(refs):
|
|
tgt = ensure_named_node(ref_name, node.start_point[0] + 1)
|
|
if tgt == impl_nid:
|
|
continue
|
|
if idx == 0:
|
|
add_edge(impl_nid, tgt, "implements", node.start_point[0] + 1)
|
|
else:
|
|
add_edge(impl_nid, tgt, "references", node.start_point[0] + 1,
|
|
context="generic_arg")
|
|
body = node.child_by_field_name("body")
|
|
if body:
|
|
for child in body.children:
|
|
walk(child, parent_impl_nid=impl_nid)
|
|
return
|
|
|
|
if t == "use_declaration":
|
|
arg = node.child_by_field_name("argument")
|
|
if arg:
|
|
raw = _read_text(arg, source)
|
|
clean = raw.split("{")[0].rstrip(":").rstrip("*").rstrip(":")
|
|
module_name = clean.split("::")[-1].strip()
|
|
if module_name:
|
|
tgt_nid = _make_id(module_name)
|
|
add_edge(file_nid, tgt_nid, "imports_from", node.start_point[0] + 1, context="import")
|
|
return
|
|
|
|
for child in node.children:
|
|
walk(child, parent_impl_nid=None)
|
|
|
|
walk(root)
|
|
|
|
label_to_nid: dict[str, str] = {}
|
|
for n in nodes:
|
|
raw = n["label"]
|
|
normalised = raw.strip("()").lstrip(".")
|
|
label_to_nid[normalised] = n["id"]
|
|
|
|
seen_call_pairs: set[tuple[str, str]] = set()
|
|
raw_calls: list[dict] = []
|
|
|
|
def walk_calls(node, caller_nid: str) -> None:
|
|
if node.type == "function_item":
|
|
return
|
|
if node.type == "call_expression":
|
|
func_node = node.child_by_field_name("function")
|
|
callee_name: str | None = None
|
|
is_member_call: bool = False
|
|
is_scoped_call: bool = False
|
|
if func_node:
|
|
if func_node.type == "identifier":
|
|
callee_name = _read_text(func_node, source)
|
|
elif func_node.type == "field_expression":
|
|
is_member_call = True
|
|
field = func_node.child_by_field_name("field")
|
|
if field:
|
|
callee_name = _read_text(field, source)
|
|
elif func_node.type == "scoped_identifier":
|
|
# Type::method() — still allow in-file EXTRACTED match, but
|
|
# skip cross-file resolution: bare last-segment lookup ignores
|
|
# crate boundaries and produces spurious INFERRED edges (#908).
|
|
is_scoped_call = True
|
|
name = func_node.child_by_field_name("name")
|
|
if name:
|
|
callee_name = _read_text(name, source)
|
|
if callee_name and callee_name not in _LANGUAGE_BUILTIN_GLOBALS:
|
|
tgt_nid = label_to_nid.get(callee_name)
|
|
if tgt_nid and tgt_nid != caller_nid:
|
|
pair = (caller_nid, tgt_nid)
|
|
if pair not in seen_call_pairs:
|
|
seen_call_pairs.add(pair)
|
|
line = node.start_point[0] + 1
|
|
edges.append({
|
|
"source": caller_nid,
|
|
"target": tgt_nid,
|
|
"relation": "calls",
|
|
"context": "call",
|
|
"confidence": "EXTRACTED",
|
|
"source_file": str_path,
|
|
"source_location": f"L{line}",
|
|
"weight": 1.0,
|
|
})
|
|
elif not is_scoped_call and callee_name.lower() not in _RUST_TRAIT_METHOD_BLOCKLIST:
|
|
raw_calls.append({
|
|
"caller_nid": caller_nid,
|
|
"callee": callee_name,
|
|
"is_member_call": is_member_call,
|
|
"source_file": str_path,
|
|
"source_location": f"L{node.start_point[0] + 1}",
|
|
})
|
|
for child in node.children:
|
|
walk_calls(child, caller_nid)
|
|
|
|
for caller_nid, body_node in function_bodies:
|
|
walk_calls(body_node, caller_nid)
|
|
|
|
valid_ids = seen_ids
|
|
clean_edges = []
|
|
for edge in edges:
|
|
src, tgt = edge["source"], edge["target"]
|
|
if src in valid_ids and (tgt in valid_ids or edge["relation"] in ("imports", "imports_from")):
|
|
clean_edges.append(edge)
|
|
|
|
return {"nodes": nodes, "edges": clean_edges, "raw_calls": raw_calls}
|