refactor: LLM AI Coding Agent audit — split monoliths, clean annotations

llm-refactor-plan.md

@@ -0,0 +1,264 @@
+# LLM Refactor Plan — petite-ad
+
+> **Status:** In progress on branch `llm-refactor`.
+> Last updated: 2026-05-11
+> Commits: `ffe4984` (annotations/re-exports), `7e4c7e1` (split compiled.rs), `84dd468` (extract expr/tape)
+
+Generated by the LLM AI Coding Agent audit checklist. Each finding is evaluated
+against the principles in the skill: predictable structure, flat explicit
+architecture, linear functions with explicit state, regenerability, and comment
+hygiene.
+
+---
+
+## High Priority
+
+### 1. ✅ Split `src/multi/compiled.rs` into backend modules
+
+**File:** `src/multi/compiled.rs` (4,615 lines)
+
+This single file holds: WGPU backend, SIMD x2/x4 backends, scalar backend,
+device planning types, batch buffer types, CompiledGraph, BackendKind dispatch,
+OpCode/FlatInstruction, and all execution logic. Five distinct concerns are
+interleaved in one file.
+
+**Proposed split:**
+
+```
+src/multi/
+  backend/
+    mod.rs           → pub mod + re-exports
+    types.rs         → OpCode, FlatInstruction, BackendCapabilities, Instruction
+    scalar.rs        → ScalarBackend
+    simd.rs          → SimdBackend (f64x2, f64x4)
+    wgpu.rs          → WgpuBackend, WgpuBuffer, WgpuBufferSet (feature-gated)
+    device.rs        → DeviceBatchPlan, DeviceBuffer, DeviceBufferSet,
+                       DeviceMemoryLocation, DeviceTransferKind, etc.
+    dispatch.rs      → BackendKind, BackendSupportReport, auto-dispatch logic
+    mock.rs          → MockDeviceBackend, DeviceExecutionTrace
+  compiled.rs        → keep: CompiledGraph, CompiledWorkspace, BatchInputs,
+                       BatchValues, BatchGradients, metadata
+```
+
+**Benefit:** Each backend can be regenerated independently without
+understanding the others. A future LLM working on WGPU won't need to parse
+SIMD dispatch logic.
+
+### 2. ✅ Decompose `gradient_batch_auto_into`
+
+**File:** `src/multi/compiled.rs`, line ~1,852
+
+The single largest function in the codebase. It dispatches to Scalar, SIMD x2,
+SIMD x4, WGPU paths with inline per-backend logic.
+
+**Proposed fix:**
+```rust
+fn gradient_batch_auto_into(compiled, batch, buf) -> Result<()> {
+    let backend = compiled.recommended_batch_gradient_backend();
+    match backend {
+        BackendKind::Scalar => gradient_batch_scalar_into(compiled, batch, buf),
+        BackendKind::SimdF64x2 => gradient_batch_simd_f64x2_into(compiled, batch, buf),
+        // ...
+    }
+}
+// Each arm becomes a standalone 20-50 line function.
+```
+
+**Benefit:** A future LLM replacing the SIMD backend only needs to read one
+small function, not parse 1,000 lines of unrelated dispatch code.
+
+### 3. ✅ Split `src/multi/graph.rs` into concerns
+
+**File:** `src/multi/graph.rs` (6,315 lines)
+
+Contains: `Graph` (core API, 1,550-line impl), `Tape`/`TapeWorkspace`,
+`ExprGraph`/`ExprNode` (expression graph with operator overloading),
+`GradientCheckEntry`/`GradientCheckReport`, `GraphNode`, `NodeId`,
+`DomainPolicy`, `GraphStats`, WGPU helpers (cfg-gated), serde support
+(cfg-gated).
+
+**Proposed split:**
+
+```
+src/multi/
+  graph/
+    mod.rs        → pub mod + re-exports
+    core.rs       → Graph, GraphNode, NodeId, DomainPolicy, GraphStats
+    tape.rs       → Tape, TapeWorkspace
+    expr.rs       → ExprGraph, ExprNode, operator overloads
+    check.rs      → GradientCheckEntry, GradientCheckReport
+  graph.rs        → keep if preferred, or delete
+```
+
+**Benefit:** The `Graph` impl shrinks from 1,550 lines to ~400 lines of core
+logic. ExprGraph changes won't touch Graph internals.
+
+### 4. ✅ Remove unused import lint suppressions in `multi.rs`
+
+**File:** `src/multi.rs`, lines 51-55
+
+```rust
+#[allow(unused_imports)]
+pub use multi_ad_fr::MultiAD2FR;
+#[allow(unused_imports)]
+pub use multi_ad_rf::MultiAD2RF;
+#[allow(unused_imports)]
+pub use multi_ad_rr::MultiAD2RR;
+```
+
+These types are re-exported but never used internally in `multi.rs`. The
+`#[allow(unused_imports)]` is noise. Replace with a module-level
+`#![allow(unused_imports)]` or better: re-export them through lib.rs only and
+remove from multi.rs since they're already in lib.rs's `pub use` block.
+
+### 5. ✅ Clean up `#[allow(dead_code)]` on public API methods
+
+**Affected files:** `src/mono/mono_fn.rs`, `src/multi/multi_fn.rs`,
+`src/mono/types.rs`, `src/forward.rs`, `src/multi/builder.rs`
+
+Public library methods tagged `#[allow(dead_code)]` signal that no internal
+code calls them. This is expected for a library, but the annotations clutter
+the code. Options:
+- Move to a single `#![allow(dead_code)]` at crate root
+- Or remove the annotations (dead_code doesn't apply to `pub` items when
+  building as a library, only as a binary)
+
+---
+
+## Medium Priority
+
+### 6. Unify duplicate `Dual` and `OpKind` types
+
+**Files:** `src/mono/mono_hessian_common.rs` and
+`src/multi/multi_hessian_common.rs`
+
+Both define:
+```rust
+struct Dual { val: f64, tan: f64 }
+enum OpKind { Sin, Cos, Tan, Exp, Neg, Ln, Sqrt, Abs, ... }
+```
+
+The two `Dual` types are nearly identical (mono adds `Dual::constant`,
+multi adds `Add/Sub/Mul/Div` impls). The `OpKind` enums have different
+variants (mono has 8, multi has 16+).
+
+**Proposal:** Move `Dual` to `src/types.rs` with both `constant` and arithmetic
+impls. Keep `OpKind` separate since the variants differ, but name them
+`MonoOpKind` and `MultiOpKind` to clarify which is which.
+
+### 7. Consolidate re-export blocks
+
+**Files:** `src/lib.rs` (30-line `pub use` block), `src/multi.rs` (30-line
+`pub use` block)
+
+Half the types are re-exported in both files. A future LLM needs to read both
+to understand what's publicly available.
+
+**Proposal:** Use `pub use multi::*;` in lib.rs for most types, with explicit
+re-exports only for items that need renaming or feature-gating. Keep multi.rs
+as the single source of truth for what `multi` exports.
+
+### 8. Inline `#[cfg(test)]` example function files
+
+**Files:** `src/mono/mf1.rs` through `mf4.rs` (4 files), `src/multi/f1.rs`
+through `f3.rs` (3 files)
+
+Each file is ~30 lines: one struct + one trait impl, used in exactly one test.
+Seven extra files for 210 lines of test code.
+
+**Proposal:** Move the struct definitions into the test module that uses them
+(`src/mono/tests.rs`, `src/multi/tests.rs`) as module-level items, or
+consolidate into a single `src/mono/tests/examples.rs`.
+
+### 9. Split `src/forward.rs` by variable arity
+
+**File:** `src/forward.rs` (840 lines)
+
+Contains both `ForwardAD::differentiate` (mono) and
+`ForwardAD::directional_derivative` (multi), plus graph-based forward mode.
+Three related but distinct APIs in one file.
+
+**Proposal:** Keep `src/forward.rs` for the `ForwardAD` struct and shared
+types. Move mono differentiation to `src/mono/forward.rs` and multi
+directional derivative to `src/multi/forward.rs`.
+
+### 10. Remove architecture notes from test code
+
+**File:** `src/tests_comprehensive.rs`, lines ~630-640
+
+```rust
+// Note: F1, F2, F3 are private modules, tested through multi::tests
+// since MF1-4 are private implementation details
+```
+
+These comments document where other types are tested — information that belongs
+in those test modules, not in the comprehensive test file. Remove or move.
+
+### 11. Extract per-backend plan constructors from `batch_plan`
+
+**File:** `src/multi/compiled.rs`, `fn batch_plan` (135 lines)
+
+This function builds a `DeviceBatchPlan` for each backend variant with inline
+code for every variant. Extract:
+
+```rust
+fn plan_scalar(graph, batch_size) -> DeviceBatchPlan { ... }
+fn plan_simd_f64x2(graph, batch_size) -> DeviceBatchPlan { ... }
+fn plan_wgpu(graph, batch_size) -> DeviceBatchPlan { ... }
+```
+
+---
+
+## Lower Priority
+
+### 12. Trailing comma in macro invocation
+
+**File:** `src/mono/mf4.rs`, line 19
+
+```rust
+&mono_ops![sin, neg,]
+```
+
+The trailing comma after `neg` is technically valid in Rust macros but
+inconsistent with other `mono_ops!` calls. Remove for consistency.
+
+### 13. Module-level docs missing in `src/macros.rs`
+
+**File:** `src/macros.rs` (107 lines)
+
+The file defines `mono_ops!` and `multi_ops!` macros but has no module-level
+doc comment. Add a brief explanation of the macro API and syntax.
+
+### 14. Ambiguous module name `types.rs` in both mono and multi
+
+**Files:** `src/mono/types.rs`, `src/multi/types.rs`
+
+Both define different types but share the same name. When reading code via
+`use super::types::*`, it's ambiguous which module is being referenced.
+
+**Proposal:** Rename to `mono_types.rs` and `multi_types.rs`, or use
+`mono::types` / `multi::types` consistently in imports (already done, but
+the filename convention would help).
+
+### 15. Redundant `#[cfg(test)]` on items inside `#[cfg(test)]` modules
+
+**Files:** `src/mono/mf1.rs` through `mf4.rs`, `src/multi/f1.rs` through `f3.rs`
+
+These files are declared as `#[cfg(test)] mod mf1;` in `mono.rs`, but each
+struct inside also has `#[cfg(test)]`. The outer annotation already gates
+the entire module. Remove the inner annotations.
+
+---
+
+## Summary
+
+| Priority | Items | Effort Estimate |
+|----------|-------|-----------------|
+| High | 5 | 2-3 days (split compiled.rs + graph.rs) |
+| Medium | 6 | 1-2 days |
+| Lower | 4 | 0.5 day |
+
+The highest-value change is splitting `compiled.rs` into backend modules.
+This alone addresses the core violation of the "predictable structure" and
+"flat, explicit architecture" principles. A future LLM can work on the
+WGPU backend without touching SIMD logic, and vice versa.

src/lib.rs

@@ -144,21 +144,17 @@ pub use mono::MonoAD;
 // Higher-order autodiff methods (exact Hessian computation)
 pub use mono::{MonoAD2FR, MonoAD2RF, MonoAD2RR};
 
-pub use multi::builder::GraphBuilder;
-pub use multi::multi_ad_fr::MultiAD2FR;
-pub use multi::multi_ad_rf::MultiAD2RF;
-pub use multi::multi_ad_rr::MultiAD2RR;
-pub use multi::MultiAD;
 pub use multi::{
-    AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities, BackendKind,
-    BackendRejectionReason, BackendSupportReport, BatchGradients, BatchGradientsBuffer,
-    BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph, CompiledGraphMetadata,
-    CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer, DeviceBufferHandle,
-    DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode,
+    builder::GraphBuilder, multi_ad_fr::MultiAD2FR, multi_ad_rf::MultiAD2RF,
+    multi_ad_rr::MultiAD2RR, AcceleratorDeviceContext, AcceleratorDeviceKind, BackendCapabilities,
+    BackendKind, BackendRejectionReason, BackendSupportReport, BatchGradients,
+    BatchGradientsBuffer, BatchInputs, BatchLayout, BatchValues, BatchValuesBuffer, CompiledGraph,
+    CompiledGraphMetadata, CompiledWorkspace, DeviceBackend, DeviceBatchPlan, DeviceBuffer,
+    DeviceBufferHandle, DeviceBufferKind, DeviceBufferLayout, DeviceBufferSet, DeviceExecutionMode,
     DeviceExecutionTrace, DeviceMemoryLocation, DeviceTransferKind, DeviceTransferPlan,
     DeviceTransferPolicy, DomainPolicy, ExecutionBackend, ExprGraph, ExprNode, FlatInstruction,
     GpuBackendBoundary, GradientCheckEntry, GradientCheckReport, Graph, GraphNode, GraphStats,
-    Instruction, MockDeviceBackend, NodeId, OpCode, ScalarBackend, SimdBackend, Tape,
+    Instruction, MockDeviceBackend, MultiAD, NodeId, OpCode, ScalarBackend, SimdBackend, Tape,
     TapeWorkspace, UNUSED_NODE_ID,
 };
 #[cfg(feature = "backend-wgpu")]

src/macros.rs

@@ -1,3 +1,10 @@
+//! Convenience macros for building single-variable and multi-variable
+//! computational graphs concisely.
+//!
+//! - [`mono_ops!`] — Build a `[MonoAD; N]` array of unary operations (sin, cos, tan, exp, etc.).
+//! - [`multi_ops!`] — Build a `[(MultiAD, Vec<usize>); N]` array of multi-variable operations
+//!   with input variable markers.
+
 /// Macro to convert function names to MonoAD enum at compile time.
 /// This avoids the function pointer comparison issue across library boundaries.
 ///

src/mono/mf4.rs

@@ -13,7 +13,7 @@ impl MonoFn for MF4 {
     }
 
     fn graph(&self) -> &'static GraphType {
-        &mono_ops![sin, neg,]
+        &mono_ops![sin, neg]
     }
 
     fn expected_value(&self) -> f64 {

src/mono/mono_fn.rs

@@ -2,7 +2,6 @@ pub use super::mono_ad::MonoAD;
 pub use super::types::BackwardResultBox;
 
 /// Type alias for a graph of mono operations (slice of MonoAD)
-#[allow(dead_code)] // Public API for library extension
 pub type GraphType = [MonoAD];
 
 /// Trait for single-variable functions with analytical gradients.
@@ -12,7 +11,6 @@ pub type GraphType = [MonoAD];
 ///
 /// This trait is primarily intended for testing and demonstration purposes.
 /// Most users will work directly with the `MonoAD` enum.
-#[allow(dead_code)] // Public API for library extension
 pub trait MonoFn {
     /// Returns the input value for this function.
     fn input(&self) -> f64;