diff --git a/mcp_codebase.py b/mcp_codebase.py index d4d6123..b91a6e3 100644 --- a/mcp_codebase.py +++ b/mcp_codebase.py @@ -2860,11 +2860,11 @@ def build_indexes(): created_nodes = {} # Track language-specific features for cross-language relationships - component_registry = {} # Svelte components - interface_registry = {} # TypeScript/Java interfaces - trait_registry = {} # Rust traits - class_registry = {} # Classes across all languages - package_registry = {} # Packages/modules across languages + component_registry = {} + interface_registry = {} + trait_registry = {} + class_registry = {} + package_registry = {} logger.info("=== Phase 1: Creating nodes ===") @@ -2952,7 +2952,6 @@ def build_indexes(): class_node_id = f"{lang}::class::{fpath}::{class_name}" graph.add_edge(class_node_id, node_id, "contains") else: - # Regular function, async function, or constructor graph.add_node( node_id, **base_attrs, @@ -3089,7 +3088,7 @@ def build_indexes(): type=node_type, variants=m.get("variants", []), members=m.get("members", []), - fields=m.get("fields", []), # For Rust enum variants with fields + fields=m.get("fields", []), visibility=m.get("visibility") ) @@ -3122,7 +3121,7 @@ def build_indexes(): if module_symbol_id not in graph.graph.nodes: graph.add_node( module_symbol_id, - type=name, # Use module name as type + type=name, name=name, file=fpath, lang=lang @@ -3177,7 +3176,6 @@ def build_indexes(): # ===== SVELTE SCRIPTS ===== elif node_type in ["script", "script_module"] and lang in ["typescript", "javascript"]: - # These are handled by their contained declarations pass # ===== VARIABLES & CONSTANTS ===== @@ -3187,7 +3185,7 @@ def build_indexes(): **base_attrs, type=node_type, var_type=m.get("var_type"), - var_kind=m.get("var_kind"), # const, let, var + var_kind=m.get("var_kind"), is_exported=m.get("is_exported", False), is_reactive=m.get("is_reactive", False), is_store=m.get("is_store", False), @@ -3230,8 +3228,12 @@ def build_indexes(): ) graph.add_edge(file_node_id, import_node_id, "contains") - # Create symbolic nodes for imports (rich graph structure) - _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang, created_nodes) + # Create symbolic nodes - DON'T pass created_nodes, collect new nodes separately + new_nodes = _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang) + # Add new nodes after iteration completes + for new_node_id, new_node_data in new_nodes.items(): + if new_node_id not in created_nodes: + created_nodes[new_node_id] = new_node_data node_id = import_node_id @@ -3275,7 +3277,11 @@ def build_indexes(): # ------------------------------------------------- logger.info("=== Phase 2: Building relationships ===") - for node_id, node_data in created_nodes.items(): + # CRITICAL FIX: Create a snapshot of the items to iterate over + # This prevents "dictionary changed size during iteration" errors + nodes_snapshot = list(created_nodes.items()) + + for node_id, node_data in nodes_snapshot: m = node_data.get("metadata", {}) if not m: continue @@ -3297,7 +3303,11 @@ def build_indexes(): if imports: file_node_id = f"file::{fpath}" - _add_import_edges_enhanced(graph, file_node_id, node_id, imports, lang, created_nodes) + # Collect new nodes instead of modifying created_nodes during iteration + new_nodes = _add_import_edges_enhanced(graph, file_node_id, node_id, imports, lang) + for new_node_id, new_node_data in new_nodes.items(): + if new_node_id not in created_nodes: + created_nodes[new_node_id] = new_node_data # ===== FUNCTION CALLS ===== calls_raw = m.get("calls", []) @@ -3311,7 +3321,6 @@ def build_indexes(): # ===== CLASS/STRUCT INHERITANCE ===== if node_type in ["class", "struct"]: - # Extends relationships extends = m.get("extends") if extends: if isinstance(extends, list): @@ -3324,18 +3333,15 @@ def build_indexes(): if class_name == parent_name: graph.add_edge(node_id, class_info["node_id"], "extends") - # Implements relationships implements = m.get("implements", []) if isinstance(implements, str): implements = [implements] for interface_name in implements: - # Check interface registry for iface_name, iface_info in interface_registry.items(): if iface_name == interface_name: graph.add_edge(node_id, iface_info["node_id"], "implements") - # Check trait registry (Rust) for trait_name, trait_info in trait_registry.items(): if trait_name == interface_name: graph.add_edge(node_id, trait_info["node_id"], "implements") @@ -3358,12 +3364,10 @@ def build_indexes(): traits = [traits] for trait_name in traits: - # Link to trait for t_name, t_info in trait_registry.items(): if t_name == trait_name: graph.add_edge(node_id, t_info["node_id"], "implements") - # Link to struct/type being implemented impl_name = m.get("name", "") for class_name, class_info in class_registry.items(): if class_name == impl_name or class_name in impl_name: @@ -3375,7 +3379,6 @@ def build_indexes(): receiver_base = m.get("receiver_base_type") if receiver_base: - # Find the struct this method belongs to for struct_name, struct_info in class_registry.items(): if struct_name == receiver_base and struct_info["lang"] == "go": graph.add_edge(struct_info["node_id"], node_id, "has_method") @@ -3390,7 +3393,6 @@ def build_indexes(): if type_name == clean_return: graph.add_edge(node_id, type_info["node_id"], "returns") - # Parameter types parameters = m.get("parameters", []) if isinstance(parameters, list): for param in parameters: @@ -3402,17 +3404,14 @@ def build_indexes(): # ===== SVELTE COMPONENT RELATIONSHIPS ===== if node_type == "component" or (node_type == "markup" and lang == "svelte"): - # Link to used components used_components = m.get("used_components", []) for comp_name in used_components: if comp_name in component_registry: comp_info = component_registry[comp_name] graph.add_edge(node_id, comp_info["node_id"], "uses_component") - # Link props to type definitions props = m.get("component_props", []) for prop_name in props: - # Try to find corresponding interface/type for iface_name, iface_info in interface_registry.items(): if f"{m.get('component_name')}Props" == iface_name: graph.add_edge(node_id, iface_info["node_id"], "uses_props_type") @@ -3434,21 +3433,17 @@ def build_indexes(): # ------------------------------------------------- logger.info("=== Phase 3: Advanced relationships ===") - # Build type relationships type_connections = build_type_relationships(graph, created_nodes) logger.info(f"Added {type_connections} type relationships") - # Build framework-specific relationships bevy_connections = build_bevy_relationships(graph, created_nodes) logger.info(f"Added {bevy_connections} Bevy-specific relationships") - # Build cross-language relationships cross_lang_connections = build_cross_language_relationships( graph, created_nodes, component_registry, interface_registry, class_registry ) logger.info(f"Added {cross_lang_connections} cross-language relationships") - # Save graph graph.save() graph.to_json() graph.to_toon() @@ -3511,12 +3506,15 @@ def build_indexes(): elapsed = index_build_time - start_time logger.info(f"✅ Index build complete. Total time: {elapsed:.1f}s") + # ------------------------------------------------- # Helper Functions # ------------------------------------------------- -def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, lang, created_nodes): - """Enhanced import edge creation that creates rich symbolic nodes for all languages.""" +def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, lang): + """Enhanced import edge creation. Returns dict of new nodes to add.""" + new_nodes = {} + for imp_path in imports: if not imp_path or not isinstance(imp_path, str): continue @@ -3525,7 +3523,6 @@ def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, if lang == "python": normalized_path = imp_path.replace('.', '::') elif lang in ["javascript", "typescript", "svelte"]: - # Skip very short relative imports if imp_path.startswith('.') and len(imp_path) < 3: continue normalized_path = imp_path.replace('/', '::').replace('@', '') @@ -3549,20 +3546,19 @@ def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, if package_node_id not in graph.graph.nodes: graph.add_node( package_node_id, - type=package_name, # Rich type name + type=package_name, name=package_name, file=imp_path, lang=lang, is_external=True ) - created_nodes[package_node_id] = { + new_nodes[package_node_id] = { "type": package_name, "name": package_name, "file": imp_path, "lang": lang } - # Connect file to package if not graph.graph.has_edge(file_node_id, package_node_id): graph.add_edge(file_node_id, package_node_id, "imports") @@ -3573,12 +3569,9 @@ def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, for i, part in enumerate(parts[1:], 1): current_path = f"{current_path}::{part}" - # Determine node type based on position and context if i == len(parts) - 1: - # Last part - could be a class, function, or module node_type_name = part else: - # Intermediate part - likely a module/namespace node_type_name = part module_node_id = f"symbol::{lang}::{current_path}" @@ -3586,30 +3579,29 @@ def _add_import_edges_enhanced(graph, file_node_id, importing_node_id, imports, if module_node_id not in graph.graph.nodes: graph.add_node( module_node_id, - type=node_type_name, # Rich type name + type=node_type_name, name=current_path, file=imp_path, lang=lang, is_symbol=True ) - created_nodes[module_node_id] = { + new_nodes[module_node_id] = { "type": node_type_name, "name": current_path, "file": imp_path, "lang": lang } - # Connect to parent if not graph.graph.has_edge(parent_node, module_node_id): graph.add_edge(parent_node, module_node_id, "contains") parent_node = module_node_id - # Connect the importing node to the final symbol if parent_node != package_node_id: if not graph.graph.has_edge(importing_node_id, parent_node): graph.add_edge(importing_node_id, parent_node, "imports") + return new_nodes def _clean_type_string(type_str: str) -> str: """Clean up type strings for matching across all languages.""" @@ -3652,39 +3644,27 @@ def _clean_type_string(type_str: str) -> str: return type_str - def _extract_param_type(param: str) -> str: """Extract type from parameter string for all languages.""" if not param or not isinstance(param, str): return "" - # Handle different parameter formats - # Python: "name: Type" or "name" - # Rust: "name: Type" or "Type" - # Go: "name: Type" or "name Type" - # Java/TypeScript: "Type name" or "name: Type" - if ':' in param: - # Python, Rust, TypeScript style parts = param.split(':') if len(parts) < 2: return "" param_type = parts[-1].strip() elif ' ' in param: - # Go, Java style parts = param.strip().split() - # Check if first part looks like a type (starts with uppercase or is a keyword) - if parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']: + if parts[0] and parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']: param_type = parts[0] else: - # Assume last part is the type - param_type = parts[-1] + param_type = parts[-1] if parts else "" else: return "" - # Clean up type annotations and default values - param_type = param_type.split('=')[0].strip() # Remove default values - param_type = param_type.split('`')[0].strip() # Remove Go struct tags + param_type = param_type.split('=')[0].strip() + param_type = param_type.split('`')[0].strip() return _clean_type_string(param_type) @@ -3694,37 +3674,25 @@ def _extract_field_type(field: str) -> str: if not field or not isinstance(field, str): return "" - # Handle different field formats - # Python: "name: Type" - # Rust: "name: Type" or "embedded:Type" - # Go: "name Type" or "name Type `tag`" - # Java: "Type name" - # TypeScript: "name: Type" - - # Handle embedded fields (Rust/Go) if field.startswith('embedded:'): return _clean_type_string(field[9:]) if ':' in field: - # Python, Rust, TypeScript style parts = field.split(':') if len(parts) < 2: return "" field_type = parts[-1].strip() elif ' ' in field: - # Go, Java style parts = field.strip().split() if len(parts) < 2: return "" - # Check if first part looks like a type - if parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']: + if parts[0] and parts[0][0].isupper() or parts[0] in ['int', 'string', 'bool', 'float', 'double', 'long']: field_type = parts[0] else: field_type = parts[1] if len(parts) > 1 else parts[0] else: return "" - # Remove struct tags, JSON tags, etc. field_type = re.sub(r'`.*?`', '', field_type).strip() field_type = re.sub(r'".*?"', '', field_type).strip() @@ -3736,7 +3704,6 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, """Build relationships across different programming languages.""" connections = 0 - # TypeScript/Java interfaces <-> Rust structs with matching names for ts_name, ts_info in interface_registry.items(): if ts_info["lang"] in ["typescript", "javascript"]: for rust_name, rust_info in class_registry.items(): @@ -3745,7 +3712,6 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, graph.add_edge(ts_info["node_id"], rust_info["node_id"], "mirrors") connections += 1 - # Java classes <-> Go structs with matching names for java_name, java_info in class_registry.items(): if java_info["lang"] == "java": for go_name, go_info in class_registry.items(): @@ -3754,9 +3720,7 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, graph.add_edge(java_info["node_id"], go_info["node_id"], "mirrors") connections += 1 - # Svelte components to TypeScript type definitions for comp_name, comp_info in component_registry.items(): - # Look for matching Props interface props_name = f"{comp_name}Props" for iface_name, iface_info in interface_registry.items(): if iface_name == props_name or iface_name == comp_name: @@ -3764,14 +3728,12 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, graph.add_edge(comp_info["node_id"], iface_info["node_id"], "uses_type") connections += 1 - # Look for matching class for class_name, class_info in class_registry.items(): if class_name == comp_name or class_name == props_name: if not graph.graph.has_edge(comp_info["node_id"], class_info["node_id"]): graph.add_edge(comp_info["node_id"], class_info["node_id"], "uses_type") connections += 1 - # Python classes <-> TypeScript interfaces (common in web backends) for py_name, py_info in class_registry.items(): if py_info["lang"] == "python": for ts_name, ts_info in interface_registry.items(): @@ -3780,15 +3742,6 @@ def build_cross_language_relationships(graph, created_nodes, component_registry, graph.add_edge(py_info["node_id"], ts_info["node_id"], "api_type") connections += 1 - # Rust traits <-> Go/Java interfaces (behavioral contracts) - for trait_name, trait_info in interface_registry.items(): - if trait_info.get("lang") == "rust": - for iface_name, iface_info in interface_registry.items(): - if iface_info["lang"] in ["go", "java"] and trait_name == iface_name: - if not graph.graph.has_edge(trait_info["node_id"], iface_info["node_id"]): - graph.add_edge(trait_info["node_id"], iface_info["node_id"], "similar_contract") - connections += 1 - return connections @@ -3796,7 +3749,6 @@ def build_type_relationships(graph, created_nodes): """Build enhanced type relationships within language boundaries.""" connections = 0 - # Track all type definitions by name for quick lookup type_map = {} for node_id, node_data in created_nodes.items(): node_type = node_data.get("type") @@ -3809,7 +3761,6 @@ def build_type_relationships(graph, created_nodes): type_map[key] = [] type_map[key].append(node_id) - # Build relationships based on type usage for node_id, node_data in created_nodes.items(): m = node_data.get("metadata", {}) if not m: @@ -3860,7 +3811,6 @@ def build_type_relationships(graph, created_nodes): generics = m.get("generics", []) if isinstance(generics, list): for generic in generics: - # Generics might have constraints (e.g., T: Trait) generic_parts = generic.split(':') if len(generic_parts) > 1: constraint = generic_parts[1].strip() @@ -3874,6 +3824,89 @@ def build_type_relationships(graph, created_nodes): return connections +def build_bevy_relationships(graph, created_nodes): + """Build Bevy-specific relationships (ECS patterns).""" + connections = 0 + + systems = [] + components = [] + resources = [] + bundles = [] + + for node_id, node_data in created_nodes.items(): + m = node_data.get("metadata", {}) + if not m: + continue + + lang = node_data.get("lang") + if lang != "rust": + continue + + node_type = node_data.get("type") + name = node_data.get("name") + + if node_type in ["function", "method"]: + params = m.get("parameters", []) + has_query = any("Query" in str(p) for p in params) + has_commands = any("Commands" in str(p) for p in params) + has_res = any("Res" in str(p) or "ResMut" in str(p) for p in params) + + if has_query or has_commands or has_res: + systems.append(node_id) + + if node_type == "struct": + derives = m.get("derives", []) + if "Component" in str(derives): + components.append(node_id) + + if "Resource" in str(derives): + resources.append(node_id) + + if "Bundle" in str(derives): + bundles.append(node_id) + + for system_id in systems: + system_data = created_nodes.get(system_id, {}) + m = system_data.get("metadata", {}) + params = m.get("parameters", []) + + for param in params: + param_str = str(param) + + query_match = re.findall(r'Query<[^>]*?([A-Z]\w+)', param_str) + for comp_name in query_match: + for comp_id in components: + comp_data = created_nodes.get(comp_id, {}) + if comp_data.get("name") == comp_name: + if not graph.graph.has_edge(system_id, comp_id): + graph.add_edge(system_id, comp_id, "queries_component") + connections += 1 + + res_match = re.findall(r'Res(?:Mut)?<([A-Z]\w+)', param_str) + for res_name in res_match: + for res_id in resources: + res_data = created_nodes.get(res_id, {}) + if res_data.get("name") == res_name: + if not graph.graph.has_edge(system_id, res_id): + graph.add_edge(system_id, res_id, "uses_resource") + connections += 1 + + for bundle_id in bundles: + bundle_data = created_nodes.get(bundle_id, {}) + m = bundle_data.get("metadata", {}) + fields = m.get("fields", []) + + for field in fields: + field_type = _extract_field_type(field) + for comp_id in components: + comp_data = created_nodes.get(comp_id, {}) + if comp_data.get("name") == field_type: + if not graph.graph.has_edge(bundle_id, comp_id): + graph.add_edge(bundle_id, comp_id, "bundles_component") + connections += 1 + + return connections + def build_bevy_relationships(graph, created_nodes): """Build Bevy-specific relationships (ECS patterns).""" connections = 0 @@ -3970,18 +4003,19 @@ def build_bevy_relationships(graph, created_nodes): return connections -def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang, created_nodes): - """Create rich symbolic nodes for imports to build dense graph structure.""" +def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_path, lang): + """Create rich symbolic nodes for imports. Returns dict of new nodes to add.""" + new_nodes = {} + if not import_path: - return + return new_nodes # Normalize path separators based on language if lang == "python": parts = import_path.replace('.', '::').split('::') elif lang in ["javascript", "typescript"]: - # Handle relative imports if import_path.startswith('.'): - return # Skip relative imports for symbolic nodes + return new_nodes parts = import_path.replace('/', '::').replace('@', '').split('::') elif lang == "rust": parts = import_path.split('::') @@ -3994,29 +4028,28 @@ def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_pa parts = [p.strip() for p in parts if p.strip()] if not parts: - return + return new_nodes - # Create package/crate node (first part) + # Create package/crate node package_name = parts[0] package_node_id = f"package::{lang}::{package_name}" if package_node_id not in graph.graph.nodes: graph.add_node( package_node_id, - type=package_name, # Use package name as type for rich node types + type=package_name, name=package_name, file=import_path, lang=lang, is_external=True ) - created_nodes[package_node_id] = { + new_nodes[package_node_id] = { "type": package_name, "name": package_name, "file": import_path, "lang": lang } - # Connect file to package if not graph.graph.has_edge(file_node_id, package_node_id): graph.add_edge(file_node_id, package_node_id, "imports") @@ -4031,27 +4064,27 @@ def _create_import_symbolic_nodes(graph, import_node_id, file_node_id, import_pa if module_node_id not in graph.graph.nodes: graph.add_node( module_node_id, - type=part, # Use part name as type for rich node types + type=part, name=current_path, file=import_path, lang=lang, is_symbol=True ) - created_nodes[module_node_id] = { + new_nodes[module_node_id] = { "type": part, "name": current_path, "file": import_path, "lang": lang } - # Connect parent to child if not graph.graph.has_edge(parent_node, module_node_id): graph.add_edge(parent_node, module_node_id, "contains") parent_node = module_node_id - # Connect the import node to the final symbol - if parent_node != package_node_id and not graph.graph.has_edge(import_node_id, parent_node): - graph.add_edge(import_node_id, parent_node, "imports") + + if parent_node != package_node_id and not graph.graph.has_edge(import_node_id, parent_node): + graph.add_edge(import_node_id, parent_node, "imports") + return new_nodes def _extract_parameter_type(param): """Extract type from parameter string, handling multiple formats"""