feat(java): resolve calls on chained, new and cast receivers through declared return types

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01B8gwEpjajiWHVW3c9smuKf

graphify/cache.py

@@ -37,7 +37,7 @@ except Exception:
     _EXTRACTOR_VERSION = "unknown"
 
 # Bump when AST cache-key semantics change independently of the package version.
-_AST_CACHE_SCHEMA = 10  # C# partial types carry `type_path`, the nested/generic partial merge key; prior (9): Python property-read raw calls and `_py_property` nodes; (8): Rust `use`-group per-binding edges; (7): Python opaque-base/super(args) markers + nested-class enclosing ids.
+_AST_CACHE_SCHEMA = 11  # Java `java_receiver_chain` raw-call marker and `_java_ret` return types; prior (10): C# partial types carry `type_path`, the nested/generic partial merge key; (9): Python property-read raw calls and `_py_property` nodes; (8): Rust `use`-group per-binding edges; (7): Python opaque-base/super(args) markers + nested-class enclosing ids.
 
 
 # Per-process memo: distributions() scans site-packages metadata, which is

graphify/cli.py

@@ -4106,7 +4106,7 @@ def dispatch_command(cmd: str) -> None:
                         for _marker in (
                             "_callable", "_callable_class", "_elixir_module",
                             "_rust_impl_key", "_rust_declaration_count",
-                            "_ts_public_twin", "_py_property",
+                            "_ts_public_twin", "_py_property", "_java_ret",
                         ):
                             if _node.get(_marker):
                                 _ctx_node[_marker] = _node[_marker]

graphify/extract.py

@@ -5449,6 +5449,10 @@ def _resolve_java_member_calls(
     caller's class plus method parameters and explicit locals are inferred from
     the extractor's method-scoped type table. A missing or ambiguous receiver
     type is skipped rather than falling back to a bare method-name match.
+
+    A receiver that is itself a call (`builder().size(3).build()`), a `new T()`
+    or a cast is typed step by step from the declared return types the
+    extractor put on method nodes (`_java_ret`).
     """
     def key(label: str) -> str:
         return str(label).strip().removeprefix(".").removesuffix("()")
@@ -5531,6 +5535,32 @@ def _resolve_java_member_calls(
             frontier.extend(bases)
         return next(iter(hits)) if len(hits) == 1 else None
 
+    def _type_named(name: str, near: str | None) -> str | None:
+        # A name declared more than once (a nested `Builder` in many classes)
+        # is the one in the file that names it, or none.
+        defs = type_def_nids.get(key(name), [])
+        if len(defs) > 1:
+            here = node_by_id.get(near, {}).get("source_file")
+            defs = [nid for nid in defs if node_by_id.get(nid, {}).get("source_file") == here]
+        return defs[0] if len(defs) == 1 else None
+
+    def _chain_type(chain: list, caller: str, depth: int = 0) -> str | None:
+        kind = chain[0] if chain and depth <= 16 else None
+        if kind == "type":
+            return _type_named(chain[1], caller)
+        if kind == "this":
+            return enclosing_type.get(caller)
+        if kind == "super":
+            bases = _jvm_bases(enclosing_type.get(caller, ""))
+            return bases[0] if bases and len(bases) == 1 else None
+        if kind == "call":
+            owner = (enclosing_type.get(caller) if chain[2] is None
+                     else _chain_type(chain[2], caller, depth + 1))
+            method_nid = _method_on_type_or_bases(owner, key(chain[1])) if owner else None
+            returned = node_by_id.get(method_nid or "", {}).get("_java_ret")
+            return _type_named(returned, method_nid) if returned else None
+        return None
+
     for result in per_file:
         for raw_call in result.get("raw_calls", []):
             if raw_call.get("lang") != "java" or not raw_call.get("is_member_call"):
@@ -5538,6 +5568,24 @@ def _resolve_java_member_calls(
             receiver = raw_call.get("receiver")
             callee = raw_call.get("callee")
             caller = raw_call.get("caller_nid")
+            chain = raw_call.get("java_receiver_chain")
+            if chain and not receiver and callee and caller:
+                type_nid = _chain_type(chain, caller)
+                method_nid = _method_on_type_or_bases(type_nid, key(callee)) if type_nid else None
+                if method_nid and method_nid != caller and (caller, method_nid) not in existing_pairs:
+                    existing_pairs.add((caller, method_nid))
+                    all_edges.append({
+                        "source": caller,
+                        "target": method_nid,
+                        "relation": "calls",
+                        "context": "call",
+                        "confidence": "INFERRED",
+                        "confidence_score": 0.85,
+                        "source_file": raw_call.get("source_file", ""),
+                        "source_location": raw_call.get("source_location"),
+                        "weight": 1.0,
+                    })
+                continue
             if not receiver or not callee or not caller:
                 continue
 

graphify/extractors/engine.py

@@ -463,6 +463,60 @@ def _java_receiver_type_name(type_node, source: bytes) -> str | None:
     return name
 
 
+def _java_receiver_chain(node, source: bytes, depth: int = 0) -> list | None:
+    """Describe a Java receiver that is itself an expression, for the resolver
+    to type from declared return types: `new T()` / `(T) x` -> ["type", "T"],
+    `f()` -> ["call", "f", None], `x.f()` -> ["call", "f", <x>], and the leaves
+    ["this"], ["super"], ["name", "x"]. None when any part is something else."""
+    if node is None or depth > 16:
+        return None
+    t = node.type
+    if t == "parenthesized_expression":
+        inner = [c for c in node.named_children if c.type != "comment"]
+        return _java_receiver_chain(inner[0], source, depth + 1) if len(inner) == 1 else None
+    if t in ("object_creation_expression", "cast_expression"):
+        name = _java_receiver_type_name(node.child_by_field_name("type"), source)
+        return ["type", name] if name else None
+    if t in ("this", "super"):
+        return [t]
+    if t == "identifier":
+        return ["name", _read_text(node, source)]
+    if t == "field_access":
+        owner, field = node.child_by_field_name("object"), node.child_by_field_name("field")
+        if owner is not None and owner.type == "this" and field is not None:
+            return ["name", f"this.{_read_text(field, source)}"]
+        return None
+    if t == "method_invocation":
+        name_node = node.child_by_field_name("name")
+        if name_node is None:
+            return None
+        receiver = node.child_by_field_name("object")
+        inner = None
+        if receiver is not None:
+            inner = _java_receiver_chain(receiver, source, depth + 1)
+            if inner is None:
+                return None
+        return ["call", _read_text(name_node, source), inner]
+    return None
+
+
+def _java_type_chain_leaf(chain: list, receiver_types: dict) -> list | None:
+    """Replace the ["name", x] leaf of a receiver chain with x's declared type
+    (["type", T]); a capitalized name nothing declares is a type itself
+    (`Foo.create().run()`), as for a plain receiver. None when x is untyped."""
+    if chain[0] == "call":
+        if chain[2] is None:
+            return chain
+        inner = _java_type_chain_leaf(chain[2], receiver_types)
+        return ["call", chain[1], inner] if inner else None
+    if chain[0] == "name":
+        declared = receiver_types.get(chain[1])
+        if declared:
+            return ["type", declared]
+        return ["type", chain[1]] if chain[1][:1].isupper() else None
+    return chain
+
+
 def _java_declarator_names(declaration_node, source: bytes) -> list[str]:
     names: list[str] = []
     for child in declaration_node.children:
@@ -4943,6 +4997,7 @@ def _extract_generic(
     # Java receiver typing is method-scoped: current-class fields are shared,
     # while parameters and locals belong only to their declaring method.
     java_field_types: dict[str, dict[str, str]] = {}
+    java_return_types: dict[str, str] = {}
     java_method_scopes: dict[int, tuple[object, str]] = {}
     # C# receiver typing is method-scoped too (#2299): class fields/properties
     # are shared, parameters and locals belong only to their declaring method —
@@ -6985,6 +7040,14 @@ def _extract_generic(
                             add_edge(func_nid, target_nid, "references",
                                      line, context=ctx)
 
+            if (config.ts_module == "tree_sitter_java" and parent_class_nid
+                    and node.type == "method_declaration"):
+                # Overloads share one node: keep a return type they all declare.
+                ret = _java_receiver_type_name(node.child_by_field_name("type"), source) or ""
+                java_return_types[func_nid] = (
+                    ret if java_return_types.get(func_nid, ret) == ret else ""
+                )
+
             body = _find_body(node, config)
             # JS/TS: capture callable members assigned directly in a function
             # body. Besides constructor-style `this.X = fn`, factories commonly
@@ -7729,6 +7792,7 @@ def _extract_generic(
             kotlin_object_receiver: str | None = None
             csharp_qualified_prefix: str | None = None
             csharp_receiver_chain: list[str] | None = None
+            java_receiver_chain: list | None = None
             lua_self_qualified: bool = False  # Lua self:m() rewritten to Table:m
 
             # Special handling per language
@@ -8124,6 +8188,8 @@ def _extract_generic(
                             if owner is not None and owner.type == "this" and field is not None:
                                 member_receiver = f"this.{_read_text(field, source)}"
                                 is_this_field_call = True
+                        else:
+                            java_receiver_chain = _java_receiver_chain(receiver, source)
             elif config.ts_module == "tree_sitter_ruby":
                 # Ruby's `call` node carries `receiver` and `method` as direct
                 # fields (no intermediate accessor node), so the generic accessor
@@ -8482,6 +8548,13 @@ def _extract_generic(
                             receiver_type = (receiver_types or {}).get(member_receiver or "")
                             if receiver_type:
                                 rc_entry["receiver_type"] = receiver_type
+                            if java_receiver_chain:
+                                chain = _java_type_chain_leaf(
+                                    java_receiver_chain,
+                                    receiver_types if isinstance(receiver_types, dict) else {},
+                                )
+                                if chain:
+                                    rc_entry["java_receiver_chain"] = chain
                         # Kotlin fully-qualified call (#2550): the dotted prefix +
                         # lang tag let _resolve_kotlin_qualified_calls claim it.
                         if kotlin_qualified_prefix:
@@ -8947,6 +9020,9 @@ def _extract_generic(
     ]
     if _field_table_export:
         result["member_field_tables"] = _field_table_export
+    for _n in nodes:
+        if java_return_types.get(_n["id"]):
+            _n["_java_ret"] = java_return_types[_n["id"]]
     # #2551: the parser recovered from syntax errors, so extraction may be
     # partial (in the worst case, nothing but the file node). Record the first
     # error's line so extract() can warn instead of reporting silent success.

graphify/watch.py

@@ -2034,7 +2034,7 @@ def _rebuild_code(
                     for marker in (
                         "_callable", "_callable_class", "_elixir_module",
                         "_rust_impl_key", "_rust_declaration_count",
-                        "_ts_public_twin", "_py_property",
+                        "_ts_public_twin", "_py_property", "_java_ret",
                     ):
                         if node.get(marker):
                             ctx_node[marker] = node[marker]

tests/test_java_member_calls.py

@@ -7,6 +7,7 @@ ambiguous receivers must stay unlinked rather than creating a false call edge.
 """
 from __future__ import annotations
 
+import json
 from pathlib import Path
 
 from graphify.extract import extract
@@ -241,11 +242,10 @@ def test_ambiguous_receiver_type_emits_no_edge(tmp_path: Path):
     assert send_targets == set()
 
 
-def test_inherited_field_resolves_and_chained_receiver_stays_deferred(tmp_path: Path):
+def test_inherited_field_resolves_and_chained_receiver_follows_return_type(tmp_path: Path):
     """#3151: a field declared on the superclass now types `this.<field>` in
     the subclass, so `this.gateway.charge()` resolves. A chained receiver
-    (`factory.create().charge()`) still carries no declared type and stays
-    deferred rather than guessed."""
+    (`factory.create().charge()`) is typed by `create()`'s declared return."""
     calls, result = _calls(tmp_path, {
         "Services.java": (
             "class Gateway { void charge() {} Gateway create() { return this; } }\n"
@@ -262,8 +262,8 @@ def test_inherited_field_resolves_and_chained_receiver_stays_deferred(tmp_path:
     chained = _find(result, ".chained()", "checkout")
     assert any(source == inherited and "charge" in target
                for source, target in calls), "superclass field must type the receiver"
-    assert not any(source == chained and "charge" in target
-                   for source, target in calls)
+    assert any(source == chained and "charge" in target
+               for source, target in calls), "create() declares Gateway"
 
 
 def test_unqualified_call_still_resolves(tmp_path: Path):
@@ -385,3 +385,163 @@ def test_inherited_call_is_not_bound_to_another_language(tmp_path: Path):
 
     run = _find(result, ".run()", "worker_run")
     assert not any(source == run for source, _ in calls)
+
+
+# Builder steps named like methods of an unrelated class, so only the declared
+# return types can pick the right ones.
+_BUILDER = {
+    "Product.java": "class Product { void ship() {} }\n",
+    "Builder.java": (
+        "class Builder {\n"
+        "    Builder size(int n) { return this; }\n"
+        "    Product build() { return new Product(); }\n"
+        "}\n"
+    ),
+    "Factory.java": "class Factory { static Builder builder() { return new Builder(); } }\n",
+    "Decoy.java": "class Decoy { Decoy size(int n) { return this; } void build() {} void ship() {} }\n",
+}
+
+
+def _owner_of(result: dict, method_nid: str) -> str:
+    nodes = {n["id"]: n for n in result["nodes"]}
+    return next(nodes[e["source"]]["label"] for e in result["edges"]
+                if e["relation"] == "method" and e["target"] == method_nid)
+
+
+def _call_owners(calls, result: dict, caller: str) -> set[tuple[str, str]]:
+    nodes = {n["id"]: n for n in result["nodes"]}
+    return {(_owner_of(result, t), nodes[t]["label"]) for s, t in calls if s == caller}
+
+
+def test_call_chain_is_typed_from_declared_return_types(tmp_path: Path):
+    calls, result = _calls(tmp_path, {
+        **_BUILDER,
+        "Client.java": "class Client { void run() { Factory.builder().size(3).build().ship(); } }\n",
+    })
+
+    owners = _call_owners(calls, result, _find(result, ".run()", "client_run"))
+    assert {("Builder", ".size()"), ("Builder", ".build()"), ("Product", ".ship()")} <= owners
+    assert not any(owner == "Decoy" for owner, _ in owners)
+
+
+def test_new_cast_and_bare_call_receivers_are_typed(tmp_path: Path):
+    calls, result = _calls(tmp_path, {
+        **_BUILDER,
+        "Client.java": (
+            "class Client {\n"
+            "    Builder mine() { return new Builder(); }\n"
+            "    void viaNew() { new Builder().size(1); }\n"
+            "    void viaCast(Object o) { ((Builder) o).build(); }\n"
+            "    void viaBare() { mine().build(); }\n"
+            "    void viaField() { this.factory().size(2); }\n"
+            "    Builder factory() { return Factory.builder(); }\n"
+            "}\n"
+        ),
+    })
+
+    assert ("Builder", ".size()") in _call_owners(calls, result, _find(result, ".viaNew()", "vianew"))
+    assert ("Builder", ".build()") in _call_owners(calls, result, _find(result, ".viaCast()", "viacast"))
+    assert ("Builder", ".build()") in _call_owners(calls, result, _find(result, ".viaBare()", "viabare"))
+    assert ("Builder", ".size()") in _call_owners(calls, result, _find(result, ".viaField()", "viafield"))
+
+
+def test_return_type_named_in_several_files_is_the_returning_files_own(tmp_path: Path):
+    calls, result = _calls(tmp_path, {
+        "Alpha.java": (
+            "class Alpha {\n"
+            "    static Builder builder() { return new Builder(); }\n"
+            "    static class Builder { Alpha get() { return null; } }\n"
+            "}\n"
+        ),
+        "Beta.java": (
+            "class Beta {\n"
+            "    static class Builder { Beta get() { return null; } }\n"
+            "}\n"
+        ),
+        "Client.java": "class Client { void run() { Alpha.builder().get(); } }\n",
+    })
+
+    run = _find(result, ".run()", "client_run")
+    alpha_get = _find(result, ".get()", "alpha")
+    beta_get = _find(result, ".get()", "beta")
+    assert (run, alpha_get) in calls
+    assert (run, beta_get) not in calls
+
+
+def test_untyped_chain_link_stays_unlinked(tmp_path: Path):
+    # A type-parameter return, a library return and an untyped `var` receiver
+    # carry no type to follow.
+    calls, result = _calls(tmp_path, {
+        **_BUILDER,
+        "Client.java": (
+            "class Client<T> {\n"
+            "    T self() { return null; }\n"
+            "    String text() { return \"\"; }\n"
+            "    void viaTypeParam() { self().build(); }\n"
+            "    void viaLibrary() { text().length(); }\n"
+            "    void viaVar() { var b = Factory.builder(); b.size(1).build(); }\n"
+            "}\n"
+        ),
+    })
+
+    for caller in ("viaTypeParam", "viaLibrary", "viaVar"):
+        owners = _call_owners(calls, result, _find(result, f".{caller}()", caller.lower()))
+        assert not any(label in (".build()", ".length()") for _, label in owners), (caller, owners)
+
+
+def test_chain_through_an_unchanged_file_survives_a_changed_files_rebuild(tmp_path: Path):
+    from graphify.watch import _rebuild_code
+
+    files = {**_BUILDER, "Client.java": "class Client { void run() { Factory.builder().size(3).build(); } }\n"}
+    roots = []
+    for name in ("inc", "full"):
+        root = tmp_path / name
+        root.mkdir()
+        for fname, body in files.items():
+            (root / fname).write_text(body, encoding="utf-8")
+        roots.append(root)
+    inc, full = roots
+    assert _rebuild_code(inc, no_cluster=True, acquire_lock=False, force=True)
+    for root in roots:
+        p = root / "Client.java"
+        p.write_text(p.read_text(encoding="utf-8") + "// touched\n", encoding="utf-8")
+    assert _rebuild_code(inc, changed_paths=[inc / "Client.java"], no_cluster=True, acquire_lock=False)
+    assert _rebuild_code(full, no_cluster=True, acquire_lock=False, force=True)
+
+    def calls_of(root):
+        g = json.loads((root / "graphify-out" / "graph.json").read_text(encoding="utf-8"))
+        return {(e["source"], e["target"]) for e in g.get("links", g.get("edges", [])) if e["relation"] == "calls"}
+
+    assert calls_of(inc) == calls_of(full)
+    assert any(s.endswith("client_run") and t.endswith("builder_build") for s, t in calls_of(full))
+
+
+def test_schema_bump_retires_java_entries_cached_without_receiver_chains(tmp_path: Path, monkeypatch):
+    """An entry from schema 7 has no `java_receiver_chain` / `_java_ret`; replaying
+    it would leave `Factory.builder().size(3)` unlinked."""
+    import graphify.cache as cache_mod
+    from graphify.cache import save_cached
+    from graphify.extract import extract_java
+
+    files = {**_BUILDER, "Client.java": "class Client { void run() { Factory.builder().size(3); } }\n"}
+    paths = []
+    for name, body in files.items():
+        (tmp_path / name).write_text(body, encoding="utf-8")
+        paths.append(tmp_path / name)
+    current_schema = cache_mod._AST_CACHE_SCHEMA
+    monkeypatch.setattr(cache_mod, "_EXTRACTOR_VERSION", "same-version")
+    monkeypatch.setattr(cache_mod, "_AST_CACHE_SCHEMA", 7)  # last schema without them
+    monkeypatch.setattr(cache_mod, "_cleaned_ast_dirs", set())
+    for p in paths:
+        stale = extract_java(p)
+        for raw_call in stale.get("raw_calls", []):
+            raw_call.pop("java_receiver_chain", None)
+        for node in stale["nodes"]:
+            node.pop("_java_ret", None)
+        save_cached(p, stale, root=tmp_path, cache_root=tmp_path, kind="ast")
+
+    monkeypatch.setattr(cache_mod, "_AST_CACHE_SCHEMA", current_schema)
+    monkeypatch.setattr(cache_mod, "_cleaned_ast_dirs", set())
+    r = extract(paths, root=tmp_path, cache_root=tmp_path, parallel=False)
+    calls = {(e["source"], e["target"]) for e in r["edges"] if e["relation"] == "calls"}
+    assert ("Builder", ".size()") in _call_owners(calls, r, _find(r, ".run()", "client_run"))