Skip to content

Commit 7c78413

Browse files
sahil-manglasahil-mangla
committed
fix(registry): restrict same-module suffix resolution to self/namespace receivers
1 parent 23fd0e8 commit 7c78413

4 files changed

Lines changed: 230 additions & 58 deletions

File tree

‎internal/cbm/extract_defs.c‎

Lines changed: 54 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3502,8 +3502,9 @@ static void extract_class_def(CBMExtractCtx *ctx, TSNode node, const CBMLangSpec
35023502
}
35033503

35043504
TSNode name_node = ts_node_child_by_field_name(node, TS_FIELD("name"));
3505-
// ObjC: class name is first identifier child
3506-
if (ts_node_is_null(name_node) && ctx->language == CBM_LANG_OBJC) {
3505+
// ObjC/Java: class name is first identifier child (Java fallback for enum_declaration)
3506+
if (ts_node_is_null(name_node) &&
3507+
(ctx->language == CBM_LANG_OBJC || ctx->language == CBM_LANG_JAVA)) {
35073508
name_node = cbm_find_child_by_kind(node, "identifier");
35083509
}
35093510
// Swift and newer tree-sitter-kotlin: class/object name is a type_identifier
@@ -3913,77 +3914,93 @@ static void extract_class_def(CBMExtractCtx *ctx, TSNode node, const CBMLangSpec
39133914

39143915
// Find the body/members node inside a class node
39153916
static TSNode find_class_body(TSNode class_node, CBMLanguage lang) {
3917+
TSNode result = {0};
39163918
// Try field names first
39173919
static const char *body_fields[] = {"body", "members", "class_body", "declaration_list", NULL};
39183920
for (const char **f = body_fields; *f; f++) {
39193921
TSNode body = ts_node_child_by_field_name(class_node, *f, (uint32_t)strlen(*f));
39203922
if (!ts_node_is_null(body)) {
3921-
return body;
3923+
result = body;
3924+
break;
39223925
}
39233926
}
39243927
// Go: type_spec -> type field (interface_type or struct_type)
3925-
if (lang == CBM_LANG_GO) {
3928+
if (ts_node_is_null(result) && lang == CBM_LANG_GO) {
39263929
TSNode type_inner = ts_node_child_by_field_name(class_node, TS_FIELD("type"));
39273930
if (!ts_node_is_null(type_inner)) {
3928-
return type_inner;
3931+
result = type_inner;
39293932
}
39303933
}
39313934
// ObjC: class_implementation/class_interface has no single body node
39323935
// Methods are inside implementation_definition children directly
3933-
if (lang == CBM_LANG_OBJC) {
3936+
if (ts_node_is_null(result) && lang == CBM_LANG_OBJC) {
39343937
return class_node; // iterate children of the class node itself
39353938
}
39363939
// Squirrel: class_declaration has no body field — member_declaration nodes
39373940
// (each wrapping a function_declaration) are direct children of the class.
3938-
if (lang == CBM_LANG_SQUIRREL) {
3941+
if (ts_node_is_null(result) && lang == CBM_LANG_SQUIRREL) {
39393942
return class_node;
39403943
}
39413944
// Smali: field_definition nodes are direct children of class_definition (no
39423945
// dedicated body node) — iterate the class node itself.
3943-
if (lang == CBM_LANG_SMALI) {
3946+
if (ts_node_is_null(result) && lang == CBM_LANG_SMALI) {
39443947
return class_node;
39453948
}
39463949
// GraphQL: object/interface fields live in a fields_definition child.
3947-
if (lang == CBM_LANG_GRAPHQL) {
3950+
if (ts_node_is_null(result) && lang == CBM_LANG_GRAPHQL) {
39483951
TSNode b = cbm_find_child_by_kind(class_node, "fields_definition");
39493952
if (!ts_node_is_null(b)) {
3950-
return b;
3953+
result = b;
39513954
}
39523955
}
39533956
// Prisma: model columns live in a statement_block child. Gated to Prisma so
3954-
// the common "statement_block" kind can never hijack another language's
3955-
// class body via the generic fallback below.
3956-
if (lang == CBM_LANG_PRISMA) {
3957+
// general fallback behavior is unchanged.
3958+
if (ts_node_is_null(result) && lang == CBM_LANG_PRISMA) {
39573959
TSNode b = cbm_find_child_by_kind(class_node, "statement_block");
39583960
if (!ts_node_is_null(b)) {
3959-
return b;
3961+
result = b;
39603962
}
39613963
}
39623964
// Fallback: search children for known body node types
3963-
static const char *body_types[] = {"class_body",
3964-
"interface_body",
3965-
"enum_body",
3966-
"template_body",
3967-
"interface_type",
3968-
"struct_type",
3969-
"field_declaration_list",
3970-
"compound_statement",
3971-
"block",
3972-
"closure",
3973-
"implementation_definition",
3974-
NULL};
3975-
uint32_t count = ts_node_child_count(class_node);
3976-
for (uint32_t i = 0; i < count; i++) {
3977-
TSNode child = ts_node_child(class_node, i);
3978-
const char *ck = ts_node_type(child);
3979-
for (const char **t = body_types; *t; t++) {
3980-
if (strcmp(ck, *t) == 0) {
3981-
return child;
3965+
if (ts_node_is_null(result)) {
3966+
static const char *body_types[] = {"class_body",
3967+
"interface_body",
3968+
"enum_body",
3969+
"template_body",
3970+
"interface_type",
3971+
"struct_type",
3972+
"field_declaration_list",
3973+
"compound_statement",
3974+
"block",
3975+
"closure",
3976+
"implementation_definition",
3977+
NULL};
3978+
uint32_t count = ts_node_child_count(class_node);
3979+
for (uint32_t i = 0; i < count; i++) {
3980+
TSNode child = ts_node_child(class_node, i);
3981+
const char *ck = ts_node_type(child);
3982+
for (const char **t = body_types; *t; t++) {
3983+
if (strcmp(ck, *t) == 0) {
3984+
result = child;
3985+
break;
3986+
}
3987+
}
3988+
if (!ts_node_is_null(result)) {
3989+
break;
39823990
}
39833991
}
39843992
}
3985-
TSNode null_node = {0};
3986-
return null_node;
3993+
if (!ts_node_is_null(result) && strcmp(ts_node_type(result), "enum_body") == 0) {
3994+
TSNode decls = cbm_find_child_by_kind(result, "enum_body_declarations");
3995+
if (!ts_node_is_null(decls)) {
3996+
return decls;
3997+
}
3998+
}
3999+
if (ts_node_is_null(result)) {
4000+
TSNode null_node = {0};
4001+
return null_node;
4002+
}
4003+
return result;
39874004
}
39884005

39894006
// Dart: resolve method name from method_signature/function_signature.
@@ -5836,7 +5853,8 @@ static void push_class_body_children(TSNode node, const CBMLangSpec *spec, wd_st
58365853
const char *ck = ts_node_type(child);
58375854
if (strcmp(ck, "field_declaration_list") == 0 || strcmp(ck, "class_body") == 0 ||
58385855
strcmp(ck, "declaration_list") == 0 || strcmp(ck, "body") == 0 ||
5839-
strcmp(ck, "block") == 0 || strcmp(ck, "suite") == 0 ||
5856+
strcmp(ck, "block") == 0 || strcmp(ck, "suite") == 0 || strcmp(ck, "enum_body") == 0 ||
5857+
strcmp(ck, "enum_body_declarations") == 0 ||
58405858
// Groovy class bodies are a `closure` node; routing through the
58415859
// nested-class path keeps methods from being re-walked (and thus
58425860
// double-extracted) as top-level functions. Gated to Groovy so other

‎src/pipeline/registry.c‎

Lines changed: 127 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -663,20 +663,48 @@ static cbm_resolution_t resolve_import_map(const cbm_registry_t *r, const char *
663663
return empty_result();
664664
}
665665

666+
static bool is_same_module_receiver(const char *prefix, const char *module_qn) {
667+
if (!prefix || !prefix[0]) {
668+
return true; /* Bare names are always candidates for same-module resolution */
669+
}
670+
/* 1. Check known self-receivers */
671+
static const char *const self_receivers[] = {"self", "this", "cls", "@self", NULL};
672+
for (int i = 0; self_receivers[i]; i++) {
673+
if (strcmp(prefix, self_receivers[i]) == 0) {
674+
return true;
675+
}
676+
}
677+
/* 2. Check if prefix matches the module name or its last segment (namespace qualified) */
678+
size_t plen = strlen(prefix);
679+
size_t mlen = strlen(module_qn);
680+
if (mlen == plen && strcmp(module_qn, prefix) == 0) {
681+
return true;
682+
}
683+
if (mlen > plen && module_qn[mlen - plen - 1] == '.' &&
684+
strcmp(module_qn + (mlen - plen), prefix) == 0) {
685+
return true;
686+
}
687+
return false;
688+
}
689+
666690
/* Strategy 2: Same-module match */
667691
static cbm_resolution_t resolve_same_module(const cbm_registry_t *r, const char *callee_name,
668-
const char *suffix, const char *module_qn) {
692+
const char *prefix, const char *suffix,
693+
const char *module_qn) {
669694
char candidate[CBM_SZ_512];
670695
snprintf(candidate, sizeof(candidate), "%s.%s", module_qn, callee_name);
671696
const char *stored_key = cbm_ht_get_key(r->exact, candidate);
672697
if (stored_key) {
673698
return (cbm_resolution_t){stored_key, "same_module", CONF_SAME_MODULE, REG_RESOLVED};
674699
}
675700
if (suffix && suffix[0]) {
676-
snprintf(candidate, sizeof(candidate), "%s.%s", module_qn, suffix);
677-
stored_key = cbm_ht_get_key(r->exact, candidate);
678-
if (stored_key) {
679-
return (cbm_resolution_t){stored_key, "same_module", CONF_SAME_MODULE, REG_RESOLVED};
701+
if (is_same_module_receiver(prefix, module_qn)) {
702+
snprintf(candidate, sizeof(candidate), "%s.%s", module_qn, suffix);
703+
stored_key = cbm_ht_get_key(r->exact, candidate);
704+
if (stored_key) {
705+
return (cbm_resolution_t){stored_key, "same_module", CONF_SAME_MODULE,
706+
REG_RESOLVED};
707+
}
680708
}
681709
}
682710
return empty_result();
@@ -768,6 +796,66 @@ static const char *qualified_suffix_match(const qn_array_t *arr, const char *cal
768796
}
769797
return match;
770798
}
799+
static bool qn_ends_with_qualified(const char *qn, const char *callee_name) {
800+
char dotted[CBM_SZ_512];
801+
size_t w = 0;
802+
for (const char *s = callee_name; *s && w + 1 < sizeof(dotted);) {
803+
if (s[0] == ':' && s[1] == ':') {
804+
dotted[w++] = '.';
805+
s += 2;
806+
} else {
807+
dotted[w++] = *s++;
808+
}
809+
}
810+
dotted[w] = '\0';
811+
812+
size_t qlen = strlen(qn);
813+
if (qlen < w) {
814+
return false;
815+
}
816+
const char *tail = qn + (qlen - w);
817+
if (strcmp(tail, dotted) != 0) {
818+
return false;
819+
}
820+
if (tail != qn && tail[-1] != '.') {
821+
return false;
822+
}
823+
return true;
824+
}
825+
static bool is_type_like_label(const char *label) {
826+
if (!label) {
827+
return false;
828+
}
829+
return strcmp(label, "Class") == 0 || strcmp(label, "Struct") == 0 ||
830+
strcmp(label, "Interface") == 0 || strcmp(label, "Enum") == 0 ||
831+
strcmp(label, "Type") == 0 || strcmp(label, "Trait") == 0;
832+
}
833+
834+
static bool is_candidate_method(const cbm_registry_t *r, const char *qn) {
835+
const char *label = cbm_registry_label_of(r, qn);
836+
if (label && strcmp(label, "Method") == 0) {
837+
return true;
838+
}
839+
840+
char parent_qn[CBM_SZ_512];
841+
size_t len = strlen(qn);
842+
if (len >= sizeof(parent_qn)) {
843+
return false;
844+
}
845+
strcpy(parent_qn, qn);
846+
char *last_dot = strrchr(parent_qn, '.');
847+
if (!last_dot) {
848+
return false;
849+
}
850+
*last_dot = '\0';
851+
852+
const char *parent_label = cbm_registry_label_of(r, parent_qn);
853+
return is_type_like_label(parent_label);
854+
}
855+
856+
static bool is_qualified_callee(const char *callee_name) {
857+
return strchr(callee_name, '.') != NULL || strstr(callee_name, "::") != NULL;
858+
}
771859

772860
/* Strategy 3+4: Name lookup + suffix match */
773861
static cbm_resolution_t resolve_name_lookup(const cbm_registry_t *r, const char *callee_name,
@@ -782,36 +870,53 @@ static cbm_resolution_t resolve_name_lookup(const cbm_registry_t *r, const char
782870
return empty_result(); /* unresolvably ambiguous — see REG_MAX_CANDIDATES */
783871
}
784872

873+
cbm_resolution_t res = empty_result();
874+
785875
/* Strategy 3.5: a qualified callee disambiguates among multiple same-name
786876
* candidates by full qualified tail, before bare-name scoring collapses
787877
* them onto a single winner. */
788878
if (arr->count > 1) {
789879
const char *q = qualified_suffix_match(arr, callee_name);
790880
if (q) {
791-
return (cbm_resolution_t){q, "qualified_suffix", CONF_QUALIFIED_SUFFIX, REG_RESOLVED};
881+
res = (cbm_resolution_t){q, "qualified_suffix", CONF_QUALIFIED_SUFFIX, REG_RESOLVED};
792882
}
793883
}
794884

795-
/* Strategy 3: unique name */
796-
if (arr->count == SKIP_ONE) {
797-
double conf = CONF_UNIQUE_NAME;
798-
if (import_vals && import_count > 0 &&
799-
!is_import_reachable(arr->items[0], import_vals, import_count)) {
800-
conf *= DEFAULT_CONFIDENCE;
885+
if (!(res.qualified_name && res.qualified_name[0])) {
886+
/* Strategy 3: unique name */
887+
if (arr->count == 1) {
888+
double conf = CONF_UNIQUE_NAME;
889+
if (import_vals && import_count > 0 &&
890+
!is_import_reachable(arr->items[0], import_vals, import_count)) {
891+
conf *= DEFAULT_CONFIDENCE;
892+
}
893+
res = (cbm_resolution_t){arr->items[0], "unique_name", conf, REG_RESOLVED};
801894
}
802-
return (cbm_resolution_t){arr->items[0], "unique_name", conf, REG_RESOLVED};
803895
}
804896

805-
/* Strategy 4: multiple candidates */
806-
if (import_vals && import_count > 0) {
807-
return resolve_multi_with_imports(arr, module_qn, import_vals, import_count);
897+
if (!(res.qualified_name && res.qualified_name[0])) {
898+
/* Strategy 4: multiple candidates */
899+
if (import_vals && import_count > 0) {
900+
res = resolve_multi_with_imports(arr, module_qn, import_vals, import_count);
901+
} else {
902+
const char *best =
903+
best_by_import_distance((const char **)arr->items, arr->count, module_qn);
904+
if (best) {
905+
double conf = candidate_count_penalty(CONF_SUFFIX_MATCH, arr->count);
906+
res = (cbm_resolution_t){best, "suffix_match", conf, arr->count};
907+
}
908+
}
808909
}
809-
const char *best = best_by_import_distance((const char **)arr->items, arr->count, module_qn);
810-
if (best) {
811-
double conf = candidate_count_penalty(CONF_SUFFIX_MATCH, arr->count);
812-
return (cbm_resolution_t){best, "suffix_match", conf, arr->count};
910+
911+
if (res.qualified_name && is_qualified_callee(callee_name)) {
912+
if (!is_candidate_method(r, res.qualified_name)) {
913+
if (!qn_ends_with_qualified(res.qualified_name, callee_name)) {
914+
return empty_result();
915+
}
916+
}
813917
}
814-
return empty_result();
918+
919+
return res;
815920
}
816921

817922
cbm_resolution_t cbm_registry_resolve(const cbm_registry_t *r, const char *callee_name,
@@ -864,7 +969,7 @@ cbm_resolution_t cbm_registry_resolve(const cbm_registry_t *r, const char *calle
864969
resolve_import_map(r, prefix, suffix, import_map_keys, import_map_vals, import_map_count);
865970
if (!(res.qualified_name && res.qualified_name[0])) {
866971
/* Strategy 2: same module */
867-
res = resolve_same_module(r, callee_name, suffix, module_qn);
972+
res = resolve_same_module(r, callee_name, prefix, suffix, module_qn);
868973
}
869974
if (!(res.qualified_name && res.qualified_name[0])) {
870975
/* Strategy 3+4: name lookup */

‎tests/test_extraction.c‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,23 @@ TEST(java_class_extends_and_implements) {
408408
PASS();
409409
}
410410

411+
TEST(java_enum_method) {
412+
CBMFileResult *r = extract(
413+
"enum Day {\n"
414+
" MON, TUE, WED, THU, FRI, SAT, SUN;\n\n"
415+
" public boolean isWeekend() { return this == SAT || this == SUN; }\n"
416+
" public String label() { return name().toLowerCase(); }\n"
417+
"}\n",
418+
CBM_LANG_JAVA, "t", "Day.java");
419+
ASSERT_NOT_NULL(r);
420+
ASSERT_FALSE(r->has_error);
421+
ASSERT(has_def(r, "Enum", "Day"));
422+
ASSERT(has_def(r, "Method", "isWeekend"));
423+
ASSERT(has_def(r, "Method", "label"));
424+
cbm_free_result(r);
425+
PASS();
426+
}
427+
411428
/* REPRODUCTION (RED until fixed) — Python `class Animal(Base):` must extract the
412429
* BARE base name "Base", but extract_base_classes captures the whole
413430
* `superclasses` argument_list text "(Base)" instead: collect_bases_from_field
@@ -3524,6 +3541,7 @@ SUITE(extraction) {
35243541
RUN_TEST(java_method);
35253542
RUN_TEST(java_interface);
35263543
RUN_TEST(java_class_extends_and_implements);
3544+
RUN_TEST(java_enum_method);
35273545
RUN_TEST(python_class_base_extracted_bare);
35283546
RUN_TEST(php_class);
35293547
RUN_TEST(php_function);

0 commit comments

Comments
 (0)