# codeflash-python Roadmap

What the original codeflash Python pipeline does, mapped to
codeflash-python's implementation.  Every stage is ported and tested.

Reference implementation:
`/codeflash/languages/python/` and `/codeflash/optimization/` in the
main codeflash repo.

---

## 1. Discovery — **done**

Find optimizable functions in Python source.

- `discover_functions(source, file_path)` — libcst visitor
  (`PositionProvider`, `ParentNodeProvider`)
- Filters: must have `return`, skip properties, fixtures, nested
  functions
- Output: `list[FunctionToOptimize]`

Module: `analysis/_discovery.py`

---

## 2. Source extraction — **done**

Read a function's source text from disk.

- `extract_function_source(fn)` — line slicing from
  `starting_line`/`ending_line`
- Output: `str`

Module: `analysis/_extraction.py`

---

## 3. Code replacement — **done**

Splice optimized code back into the original file.

- `replace_function_source(source, fn, new_source)` — libcst
  `.with_changes(body, decorators)`
- Handles top-level functions and class methods
- New helper functions / classes introduced by optimized code
- `__init__` replacement
- Preexisting-object dedup
- Import addition after replacement
- Global assignment insertion (caches, constants)
- Zero-diff detection

Module: `codegen/_replacement.py`, `verification/_ranking.py`,
`context/imports.py`

---

## 4. Context extraction — **done**

Build the full code context sent to the AI alongside the target
function.  Depends on Jedi for static analysis and libcst for CST
pruning.

### 4a. Function-to-FunctionSource via Jedi — **done**

- `get_function_source()` — Jedi-based resolution to `FunctionSource`

Module: `context/resolve.py`

### 4b. Helper discovery via Jedi — **done**

- `discover_helpers()` — two-level Jedi resolution (helpers + helpers-of-helpers)
- Output: `dict[Path, set[FunctionSource]]`

Module: `context/helpers.py`

### 4c. Four context types — **done**

- `collect_top_level_defs_with_dependencies()` — CST dependency collection
- `mark_defs_for_functions()` / `remove_unused_definitions_by_function_names()` — reachability marking + pruning
- `parse_code_and_prune_cst()` — per-context-type CST filtering
- `add_needed_imports_from_module()` — import handling
- `extract_all_contexts()` → `CodeOptimizationContext`
- Single CST parse per file, reused for all four contexts

Modules: `context/dependencies.py`, `context/pruning.py`,
`context/imports.py`, `context/orchestration.py`

### 4d. Testgen enrichment — **done**

Ref: `code_context_extractor.py:1484-1583`

- For imported classes in testgen context, resolve via Jedi to their
  definitions
- Extract class source or stub (for third-party)
- Recursive base-class extraction
- `extract_parameter_type_constructors()` — find types used in
  function signature, extract `__init__` stubs (handles dataclasses,
  attrs, NamedTuple)
- `extract_init_stub_from_class()` — synthetic `__init__` for
  declarative classes

### 4e. Token-limit fallback — **done**

Ref: `code_context_extractor.py:166-189`

Progressive degradation when context exceeds token budget:
1. Full context
2. Strip docstrings from read-only
3. Remove read-only entirely

---

## 5. Test discovery — **done**

Find existing tests that exercise a target function.

Ref: `discovery/discover_unit_tests.py`

### 5a. File-level filtering — **done**

- Glob for `test_*.py` / `*_test.py` in test root
- `filter_test_files_by_imports()` — parse imports, keep only
  files that import target function names (cheap pre-filter before
  Jedi)

### 5b. Jedi-based test linking — **done**

Ref: `discover_unit_tests.py:859-956`

- For each test function, use Jedi references to find calls to
  target functions
- Record call positions (line/col) for instrumentation
- Output: `dict[qualified_name, set[FunctionCalledInTest]]`

### 5c. Replay test discovery — **done**

- Find tests generated by the tracer (recorded input/output pairs)
- Parse replay test metadata

---

## 6. Test instrumentation — **done**

Modify existing tests to capture behavior and performance baselines.

Ref: `code_utils/instrument_existing_tests.py`

### 6a. `InjectPerfOnly` AST transformer — **done**

- Walk test AST, find calls to target function at known positions
- **Behavior mode**: wrap calls with `codeflash_capture.record_call()`
  to capture inputs/outputs
- **Performance mode**: wrap with timing measurement
- Handles nested calls, multiple calls per test, try/except blocks

### 6b. Output — **done**

Modified test source files (written to temp directory), ready to run
with pytest.

---

## 7. Test generation (AI) — **done**

Generate new test cases via the AI service.

Ref: `verification/generate_tests.py`

- Send testgen context to AI
- Receive pytest test functions
- Post-process: syntax validation, dedup, formatting
- Merge with existing tests

Module: `testing/_testgen.py`

---

## 8. Baseline establishment — **done**

Run all tests (existing + generated) on original code.

Ref: `function_optimizer.py:setup_and_establish_baseline`

### 8a. JIT detection — **done**

- `contains_jit_decorator()` — AST-based detection of JIT decorators
  from numba, torch, tensorflow, jax (handles aliases and arguments)
- `jit_disabled_env()` — environment variables to disable JIT during
  coverage measurement

Module: `verification/_baseline.py`

### 8b. Test execution and results parsing — **done**

- Run pytest with codeflash plugin to capture:
  - Behavioral baseline (return values via pickle/sqlite)
  - Performance baseline (timing)
  - Coverage baseline (optional, via `coverage` package)

Module: `testing/_test_runner.py`, `testing/_parse_results.py`

### 8c. Baseline orchestration — **done**

- `establish_original_code_baseline()` — orchestrates behavioral
  testing, coverage validation, line profiling, and performance
  benchmarking

Module: `verification/_baseline.py`

---

## 9. Candidate processing & verification — **done**

Evaluate each AI-generated optimization.

Ref: `function_optimizer.py:process_single_candidate`,
`determine_best_candidate`

### 9a. Per-candidate loop — **done**

1. Syntax validation (`ast.parse`)
2. Code replacement in module
3. Run behavioral tests — compare outputs via deep comparator
4. Run performance tests — measure speedup
5. Coverage check (optional)

### 9b. Deep comparator — **done**

Ref: `verification/comparator.py`

- Recursive equality for Python objects
- Special handling for: NumPy arrays, PyTorch tensors, Pandas
  DataFrames, NaN equality, temp path normalization, attrs/slots
  objects

Module: `verification/_comparator.py`

### 9c. Candidate ranking — **done**

- `CandidateEvaluationContext` tracks speedup ratios per candidate
- AST-level dedup (normalized code comparison)
- Best candidate = highest verified speedup that passes all tests

Module: `verification/_ranking.py`

---

## 10. Refinement & repair — **done**

Iteratively improve candidates.

### 10a. Refinement — **done**

Ref: `api/aiservice.py:optimize_code_refinement`

- Send original + optimized code + line profiler results to
  `/ai/refinement`
- Get improved candidate

### 10b. Code repair — **done**

Ref: `api/aiservice.py:code_repair`

- When tests fail, send original + broken code + test diffs to
  `/ai/code_repair`
- Get fixed candidate

### 10c. Adaptive optimization — **done**

- Send multiple candidates + their speedups to `/ai/adaptive_optimize`
- Get new candidate informed by what worked

Module: `ai/_refinement.py`

---

## 11. Reference graph (caching) — **done**

Persistent index of function call relationships for faster
dependency analysis on subsequent runs.

Ref: `languages/python/reference_graph.py`

- `ReferenceGraph` backed by SQLite
  (`~/.codeflash/codeflash_cache.db`)
- `build_index()` — parallel Jedi analysis of all project files
- `get_call_graph()` — return `CallGraph` for given functions
- File-hash–based cache invalidation
- Augment graph with trace profiling data

Module: `analysis/_call_graph.py`, `analysis/_reference_graph.py`

---

## 12. Function ranking — **done**

Module: `analysis/_function_ranking.py`

Prioritize functions by optimization impact.

Ref: `benchmarking/function_ranker.py`

- `FunctionRanker.rank_functions()` — rank by addressable time
  (own_time + callee_time / call_count)
- Importance threshold (~0.1% of total runtime)
- Fallback: rank by dependency count from call graph
- Fallback: original discovery order

---

## 13. Tracing / profiling **done**

Collect runtime data to power ranking and replay tests.

Ref: `tracer.py`

- Python tracer using `sys.setprofile`
- Captures: call counts, timing, call graph edges
- Stores in SQLite (`pstats` table)
- Generates replay tests from recorded inputs/outputs

---

## 14. Line profiling — **done**

Per-line performance data to guide optimization.

Ref: `static_analysis/line_profile_utils.py`

- Add `@profile` decorator to target function
- Run with `line_profiler`
- Parse results → send to AI for targeted optimization
- Separate API endpoint: `/ai/optimize-line-profiler`

---

## 15. Pytest plugin — **done**

Timing and stability instrumentation for test execution.

Ref: `verification/pytest_plugin.py`

- Pytest hooks for injecting timing markers into test output
- Stability thresholds for performance verification
- Numpy-based performance statistics (mean, std, outlier removal)
- System time measurement and validation
- Pytest pattern filtering (excludes pytest/plugin infrastructure overhead)

---

## 16. Static analysis utilities — **done**

Import validation, code extraction helpers, and generated-test
editing that support the orchestrator.

### 16a. Import validation

Ref: `languages/python/static_analysis/static_analysis.py`

- `ImportedInternalModuleAnalysis` — validates imported callees
  and their dependencies before optimization

### 16b. Code extraction

Ref: `languages/python/static_analysis/code_extractor.py`

- `GlobalFunctionCollector` — CST visitor for module-level functions
- Source import gathering and pre-existing object detection
- Insertion point finding (after imports)
- Global assignment collection for context
- Token-based context limiting (`OPTIMIZATION_CONTEXT_TOKEN_LIMIT`,
  `TESTGEN_CONTEXT_TOKEN_LIMIT`)

### 16c. Generated test editing

Ref: `languages/python/static_analysis/edit_generated_tests.py`

- Extract dependent functions from generated test code
- Identify single dependent functions for focused optimization

---

## 17. Unused helper detection — **done**

Post-optimization cleanup: detect and revert helper functions
introduced by the optimizer that are never called.

Ref: `languages/python/context/unused_definition_remover.py`

- `detect_unused_helper_functions()` — find unreferenced helpers
  in optimized code
- `revert_unused_helper_functions()` — remove unused helpers,
  restore original code for those symbols
- Attribute assignment tracking for usage analysis

---

## 18. Coverage integration — **done**

Optional coverage checking during baseline and verification.

Ref: `languages/python/static_analysis/coverage_utils.py`,
`verification/coverage_utils.py`

- Integration with `coverage.py` for line-level coverage data
- Coverage-based filtering of test relevance
- Coverage diff between original and optimized code

---

## 19. Subprocess runners — **done**

Process isolation for test discovery and benchmark execution.
Prevents import pollution and resource leaks from affecting the
main optimizer process.

### 19a. Test discovery subprocess

Ref: `discovery/pytest_new_process_discovery.py`

- Spawn isolated subprocess for pytest test collection
- Collect test metadata and file paths without importing
  user code in the main process

### 19b. Benchmark subprocess

Ref: `benchmarking/pytest_new_process_trace_benchmarks.py`

- Spawn isolated subprocess for benchmark execution
- Capture trace data in separate process

---

## 20. Benchmark orchestration **done**

Full benchmark pipeline: trace collection, replay, comparison,
and result formatting.

Ref: `benchmarking/`

- `trace_benchmarks.py` — orchestrates benchmark tracing in
  subprocess, error pattern detection
- `instrument_codeflash_trace.py` — injects tracing decorators
  into benchmark functions
- `replay_test.py` — re-executes tests with performance measurement
- `compare.py` — before/after benchmark comparison and reporting
- `utils.py` — `validate_and_format_benchmark_table()`,
  multithreading/multiprocessing edge cases
- `plugin/plugin.py` — pytest plugin for benchmark collection

---

## 21. Platform API — **done**

CodeFlash platform client for non-AI service interactions.

Ref: `api/cfapi.py`

Package: **codeflash-core** (shared client alongside `_client.py`)

- Get blocklisted functions (functions to skip optimization)
- Check if function is being re-optimized
- Send completion emails
- Git repository owner/name extraction
- Authentication and telemetry integration

---

## 22. Concolic testing — **done**

CrossHair symbolic execution for test generation.

Ref: `languages/python/static_analysis/concolic_utils.py`

- Validate concolic-generated tests
- Filter known CrossHair limitations (`<locals>`, iterator objects)
- Syntax error detection and subprocess test validation

---

## 23. Pipeline orchestration — **done**

End-to-end optimizer that sequences all stages into a complete
optimization run. This is the top-level entry point.

### 23a. Module preparation

Ref: `languages/python/optimizer.py`

- `prepare_python_module()` — parse module, normalize code,
  validate imported callees
- Resolve function AST and coordinate optimization workflow

### 23b. Per-function optimization loop

Ref: `languages/python/function_optimizer.py`

- Extends base `FunctionOptimizer`
- Sequences: testgen → baseline → candidates → verify → rank
- Numerical code detection and optimization review metrics
- Unused helper detection post-verification

### 23c. High-level orchestrator

Ref: `optimization/optimizer.py`

- Worktree management for safe code replacement
- Progress display and call graph visualization
- Telemetry integration (Sentry, PostHog)
- Git-based diff generation
- Email notification on completion

### 23d. Language support

Ref: `languages/python/support.py`

- Python `Language` subclass
- Integrates libcst caching, test discovery dispatch
- LRU cache infrastructure for performance

---

## 24. Public API — **done**

Programmatic API for LLM agents, MCP tools, and external
automation.  Replaces direct submodule imports with a stable,
serializable interface.

### 24a. Subpackage re-exports — **done**

- Populated all empty `__init__.py` files with `__all__` and
  re-exports (analysis, codegen, context, testing, verification,
  benchmarking, pipeline, ai, runtime)
- Follows `test_discovery/__init__.py` pattern
- Moved `find_preexisting_objects` from `codegen` to `analysis`
  to break circular import

### 24b. Model serialization — **done**

- `to_dict()`/`from_dict()` on `FunctionParent`,
  `FunctionToOptimize`, `FunctionSource`, `FunctionResult`,
  `CodeOptimizationContext`, `CodeStringsMarkdown`
- JSON-serializable inputs/outputs for future MCP exposure

### 24c. OptimizationConfig — **done**

- `api/_config.py` — frozen attrs class
- Fields: `project_root`, `module_root`, `tests_root`,
  `test_framework`, `pytest_cmd`, `ignore_paths`, `api_key`,
  `n_candidates`, `ai_timeout`
- `to_dict()`/`from_dict()` roundtrip

### 24d. OptimizationSession — **done**

- `api/_session.py` — mutable session with lazy AI client
- Step-by-step methods: `discover_functions`, `extract_context`,
  `generate_candidates`, `apply`
- Context manager support (`__enter__`/`__exit__`)
- `optimize_function()` one-shot facade
- Experiment loop stubs: `profile`, `build_targets`, `measure`,
  `evaluate` (raise `NotImplementedError`)

### 24e. Top-level re-exports — **done**

- `FunctionParent`, `FunctionSource`, `FunctionToOptimize`,
  `OptimizationConfig`, `OptimizationSession`, `optimize_function`
  importable from `codeflash_python`

---

## 25. Experiment loop — **planned**

Flesh out the agent experiment loop on `OptimizationSession`:
`profile()`, `build_targets()`, `measure()`, `evaluate()`.

This is the workflow where an agent iteratively profiles code,
selects optimization targets, applies fixes, measures results,
and decides whether to keep or discard each change.

---

## 26. MCP tool layer — **planned**

Expose `OptimizationSession` methods as MCP tools so LLM agents
can drive optimization via tool calls.  The serialization
groundwork (stage 24b) is in place — each method already accepts
and returns JSON-serializable types.

---

## 27. Agent migration — **planned**

Port the `codeflash-agent` Claude Code plugin from shell-command
invocations to the `OptimizationSession` programmatic API.  This
removes the dependency on CLI parsing and makes the agent's
actions testable and composable.

---

## Dependencies by stage

```
                    1. Discovery (libcst) ✅
                         │
                    2. Extraction ✅
                         │
              ┌──────────┼──────────┐
              │          │          │
         3. Replacement  │     4. Context ✅
            (libcst) ✅  │     (Jedi + libcst)
                         │          │
                    5. Test discovery ✅
                       (Jedi)
                         │
                    6. Instrumentation ✅
                       (ast)
                         │
                    7. Test generation ✅
                       (AI service)
                         │
                    8. Baseline ✅
                       (pytest)
                         │
                    9. Verification ✅
                       (pytest + comparator)
                         │
              ┌──────────┼──────────┐
              │          │          │
         10. Refine ✅  11. Ref graph ✅  12. Ranking ✅
                        (Jedi+SQLite)     (trace data)
                         │
                    13. Tracing ✅
                    (sys.settrace)
                         │
                    14. Line profiling ✅
                    (line_profiler)

         15. Pytest plugin ✅ ←── 8 (baseline)
              │
    ┌─────────┼──────────┐
    │         │          │
  18. Coverage ✅  20. Bench ✅  19. Subprocess ✅
  (optional)       orchestr.      runners

  16. Static analysis ✅ ←── 4 (context)
  17. Unused helpers ✅ ←── 9 (verification)
  21. Platform API ✅ (core) ←── independent
  22. Concolic ✅ ←── 7 (testgen)

              │
         23. Pipeline orchestration ✅
             (depends on ALL above)
              │
         24. Public API ✅
         (session, config, serialization)
              │
    ┌─────────┼──────────┐
    │         │          │
  25. Experiment  26. MCP    27. Agent
      loop        tools      migration
```
