chore: import upstream snapshot with attribution

graphify/extractors/rust.py

@@ -0,0 +1,410 @@
+"""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}