refactor: simplify mocha builders, deduplicate treesitter, remove dead code

codeflash/discovery/functions_to_optimize.py

@@ -183,25 +183,6 @@ def _is_js_ts_function_exists_but_not_exported(file_path: Path, function_name: s
     return False
 
 
-def _find_all_functions_via_language_support(file_path: Path) -> dict[Path, list[FunctionToOptimize]]:
-    """Find all optimizable functions using the language support abstraction.
-
-    This function uses the registered language support for the file's language
-    to discover functions, then converts them to FunctionToOptimize instances.
-    """
-    from codeflash.languages.base import FunctionFilterCriteria
-
-    functions: dict[Path, list[FunctionToOptimize]] = {}
-
-    try:
-        lang_support = get_language_support(file_path)
-        criteria = FunctionFilterCriteria(require_return=True)
-        functions[file_path] = lang_support.discover_functions(file_path, criteria)
-    except Exception as e:
-        logger.debug(f"Failed to discover functions in {file_path}: {e}")
-
-    return functions
-
 
 def get_functions_to_optimize(
     optimize_all: str | None,

codeflash/languages/javascript/mocha_runner.py

@@ -207,15 +207,15 @@ def _extract_mocha_json(stdout: str) -> str | None:
     return None
 
 
-def _build_mocha_behavioral_command(
-    test_files: list[Path], timeout: int | None = None, project_root: Path | None = None
+def _build_mocha_command(
+    test_files: list[Path], timeout: int | None = None, default_timeout_ms: int = 60000
 ) -> list[str]:
-    """Build Mocha command for behavioral tests.
+    """Build a Mocha command.
 
     Args:
         test_files: List of test files to run.
         timeout: Optional timeout in seconds (converted to ms for Mocha).
-        project_root: Project root directory.
+        default_timeout_ms: Default timeout in milliseconds when timeout is not provided.
 
     Returns:
         Command list for subprocess execution.
@@ -226,59 +226,7 @@ def _build_mocha_behavioral_command(
     if timeout:
         cmd.extend(["--timeout", str(timeout * 1000)])
     else:
-        cmd.extend(["--timeout", "60000"])
-
-    cmd.extend(str(f.resolve()) for f in test_files)
-
-    return cmd
-
-
-def _build_mocha_benchmarking_command(
-    test_files: list[Path], timeout: int | None = None, project_root: Path | None = None
-) -> list[str]:
-    """Build Mocha command for benchmarking tests.
-
-    Args:
-        test_files: List of test files to run.
-        timeout: Optional timeout in seconds (converted to ms for Mocha).
-        project_root: Project root directory.
-
-    Returns:
-        Command list for subprocess execution.
-
-    """
-    cmd = ["npx", "mocha", "--reporter", "json", "--jobs", "1", "--exit"]
-
-    if timeout:
-        cmd.extend(["--timeout", str(timeout * 1000)])
-    else:
-        cmd.extend(["--timeout", "120000"])
-
-    cmd.extend(str(f.resolve()) for f in test_files)
-
-    return cmd
-
-
-def _build_mocha_line_profile_command(
-    test_files: list[Path], timeout: int | None = None, project_root: Path | None = None
-) -> list[str]:
-    """Build Mocha command for line profiling tests.
-
-    Args:
-        test_files: List of test files to run.
-        timeout: Optional timeout in seconds (converted to ms for Mocha).
-        project_root: Project root directory.
-
-    Returns:
-        Command list for subprocess execution.
-
-    """
-    cmd = ["npx", "mocha", "--reporter", "json", "--jobs", "1", "--exit"]
-
-    if timeout:
-        cmd.extend(["--timeout", str(timeout * 1000)])
-    else:
-        cmd.extend(["--timeout", "60000"])
+        cmd.extend(["--timeout", str(default_timeout_ms)])
 
     cmd.extend(str(f.resolve()) for f in test_files)
 
@@ -391,7 +339,7 @@ def run_mocha_behavioral_tests(
 
     _ensure_runtime_files(effective_cwd)
 
-    mocha_cmd = _build_mocha_behavioral_command(test_files=test_files, timeout=timeout, project_root=effective_cwd)
+    mocha_cmd = _build_mocha_command(test_files=test_files, timeout=timeout)
 
     mocha_env = test_env.copy()
     codeflash_sqlite_file = get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite"))
@@ -475,7 +423,7 @@ def run_mocha_benchmarking_tests(
 
     _ensure_runtime_files(effective_cwd)
 
-    mocha_cmd = _build_mocha_benchmarking_command(test_files=test_files, timeout=timeout, project_root=effective_cwd)
+    mocha_cmd = _build_mocha_command(test_files=test_files, timeout=timeout, default_timeout_ms=120000)
 
     mocha_env = test_env.copy()
     codeflash_sqlite_file = get_run_tmp_file(Path("test_return_values_0.sqlite"))
@@ -563,7 +511,7 @@ def run_mocha_line_profile_tests(
 
     _ensure_runtime_files(effective_cwd)
 
-    mocha_cmd = _build_mocha_line_profile_command(test_files=test_files, timeout=timeout, project_root=effective_cwd)
+    mocha_cmd = _build_mocha_command(test_files=test_files, timeout=timeout)
 
     mocha_env = test_env.copy()
     codeflash_sqlite_file = get_run_tmp_file(Path("test_return_values_line_profile.sqlite"))

codeflash/languages/javascript/treesitter.py

@@ -579,39 +579,7 @@ class TreeSitterAnalyzer:
             module.exports = utils;
         → checks if prop_name is a key of the object literal assigned to var_name.
         """
-        root = node
-        while root.parent:
-            root = root.parent
-
-        for child in root.children:
-            # Look for: const/let/var varName = { ... }
-            if child.type in ("lexical_declaration", "variable_declaration"):
-                for decl in child.children:
-                    if decl.type == "variable_declarator":
-                        name_node = decl.child_by_field_name("name")
-                        value_node = decl.child_by_field_name("value")
-                        if (
-                            name_node
-                            and self.get_node_text(name_node, source_bytes) == var_name
-                            and value_node
-                            and value_node.type == "object"
-                        ):
-                            for obj_child in value_node.children:
-                                if obj_child.type == "method_definition":
-                                    method_name_node = obj_child.child_by_field_name("name")
-                                    if (
-                                        method_name_node
-                                        and self.get_node_text(method_name_node, source_bytes) == prop_name
-                                    ):
-                                        return True
-                                elif obj_child.type == "shorthand_property_identifier":
-                                    if self.get_node_text(obj_child, source_bytes) == prop_name:
-                                        return True
-                                elif obj_child.type == "pair":
-                                    key_node = obj_child.child_by_field_name("key")
-                                    if key_node and self.get_node_text(key_node, source_bytes) == prop_name:
-                                        return True
-        return False
+        return any(name == prop_name for name, _ in self._resolve_variable_object_properties(node, var_name, source_bytes))
 
     def _find_preceding_jsdoc(self, node: Node, source_bytes: bytes) -> int | None:
         """Find JSDoc comment immediately preceding a function node.

tests/test_languages/test_mocha_runner.py

@@ -226,17 +226,17 @@ class TestFindMochaProjectRoot:
             assert result == root
 
 
-class TestMochaBehavioralCommand:
-    """Tests for building Mocha behavioral commands."""
+class TestBuildMochaCommand:
+    """Tests for building Mocha commands."""
 
     def test_basic_command(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_behavioral_command
+        from codeflash.languages.javascript.mocha_runner import _build_mocha_command
 
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = Path(tmpdir) / "test.js"
             test_file.write_text("// test")
 
-            cmd = _build_mocha_behavioral_command(test_files=[test_file])
+            cmd = _build_mocha_command(test_files=[test_file])
             assert "npx" in cmd
             assert "mocha" in cmd
             assert "--reporter" in cmd
@@ -246,78 +246,46 @@ class TestMochaBehavioralCommand:
             assert "--exit" in cmd
 
     def test_timeout_flag(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_behavioral_command
+        from codeflash.languages.javascript.mocha_runner import _build_mocha_command
 
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = Path(tmpdir) / "test.js"
             test_file.write_text("// test")
 
-            cmd = _build_mocha_behavioral_command(test_files=[test_file], timeout=30)
+            cmd = _build_mocha_command(test_files=[test_file], timeout=30)
             timeout_idx = cmd.index("--timeout")
             assert cmd[timeout_idx + 1] == "30000"
 
     def test_default_timeout(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_behavioral_command
+        from codeflash.languages.javascript.mocha_runner import _build_mocha_command
 
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = Path(tmpdir) / "test.js"
             test_file.write_text("// test")
 
-            cmd = _build_mocha_behavioral_command(test_files=[test_file])
+            cmd = _build_mocha_command(test_files=[test_file])
             timeout_idx = cmd.index("--timeout")
             assert cmd[timeout_idx + 1] == "60000"
 
-
-class TestMochaBenchmarkingCommand:
-    """Tests for building Mocha benchmarking commands."""
-
-    def test_basic_command(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_benchmarking_command
+    def test_benchmarking_default_timeout_is_longer(self):
+        from codeflash.languages.javascript.mocha_runner import _build_mocha_command
 
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = Path(tmpdir) / "test.js"
             test_file.write_text("// test")
 
-            cmd = _build_mocha_benchmarking_command(test_files=[test_file])
-            assert "npx" in cmd
-            assert "mocha" in cmd
-            assert "--exit" in cmd
-
-    def test_default_timeout_is_longer(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_benchmarking_command
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            test_file = Path(tmpdir) / "test.js"
-            test_file.write_text("// test")
-
-            cmd = _build_mocha_benchmarking_command(test_files=[test_file])
+            cmd = _build_mocha_command(test_files=[test_file], default_timeout_ms=120000)
             timeout_idx = cmd.index("--timeout")
             assert cmd[timeout_idx + 1] == "120000"
 
-
-class TestMochaLineProfileCommand:
-    """Tests for building Mocha line profile commands."""
-
-    def test_basic_command(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_line_profile_command
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            test_file = Path(tmpdir) / "test.js"
-            test_file.write_text("// test")
-
-            cmd = _build_mocha_line_profile_command(test_files=[test_file])
-            assert "npx" in cmd
-            assert "mocha" in cmd
-            assert "--exit" in cmd
-
     def test_timeout_conversion(self):
-        from codeflash.languages.javascript.mocha_runner import _build_mocha_line_profile_command
+        from codeflash.languages.javascript.mocha_runner import _build_mocha_command
 
         with tempfile.TemporaryDirectory() as tmpdir:
             test_file = Path(tmpdir) / "test.js"
             test_file.write_text("// test")
 
-            cmd = _build_mocha_line_profile_command(test_files=[test_file], timeout=45)
+            cmd = _build_mocha_command(test_files=[test_file], timeout=45)
             timeout_idx = cmd.index("--timeout")
             assert cmd[timeout_idx + 1] == "45000"