Perf #7: Investigate Dict{Int, Any} type stability in KKT construction

RETROSPECTIVES.md

@@ -2977,3 +2977,68 @@ Replaced `Vector{Function}` with `Vector{MNFunctionWrapper}` for M_fns!/N_fns! i
 ### Action Items for Next PR
 
 - When wrapping external function outputs, always test with the most complex available test case, not just the simplest one.
+
+---
+
+## PR: perf/07-typed-dicts
+
+**Date:** 2026-02-16
+**Commits:** 3
+**Tests:** 1405 passing (24 new)
+
+### Summary
+
+Investigation PR for Perf #7: Dict{Int, Any} type instability in KKT construction.
+Traced the full call chain from `solve()` through `run_nonlinear_solver` through
+`compute_K_evals` to determine which dicts are hot-path vs cold-path.
+
+**Key finding:** All `Dict{Int, Any}` instances in `qp_kkt.jl` (lines 92-95) and
+`nonlinear_kkt.jl` (line 256) are **cold-path only** — used during one-time symbolic
+construction, never accessed in the Newton iteration loop. The hot-path containers
+were already well-typed from previous PRs (Perf #4, #5, #6).
+
+### TDD Compliance
+
+**Score: N/A (investigation PR)**
+
+This was primarily an investigation PR. Tests were written to document the existing
+type guarantees as regression tests, not to drive new implementation. The one code
+change (explicit `Dict{Int, Int}` annotation on `π_sizes_trimmed`) was already
+correctly inferred by Julia — the annotation is for documentation only.
+
+### Clean Code
+
+- Tests clearly document the hot-path vs cold-path distinction
+- Each testset has a descriptive name explaining WHY the type is what it is
+- Benchmark script included for reproducibility
+
+### Commits
+
+- 3 small, focused commits: tests, implementation+tier-registration, benchmarks
+- Each commit has descriptive message explaining the finding
+
+### CLAUDE.md Compliance
+
+- [x] Reviewed CLAUDE.md at PR start
+- [x] Expert review protocol followed (hot-path identification checkpoint)
+- [x] Investigation documented before implementation
+- [x] Pre-merge retrospective recorded
+
+### What Went Well
+
+- Thorough call chain analysis before writing any code prevented wasted effort
+- Correctly identified that no hot-path changes were needed
+- Regression tests add value even when no code changes are needed
+
+### What Could Be Improved
+
+- The initial task description assumed Dict{Int, Any} was on the hot path. Earlier
+  investigation (before writing the task) would have saved the entire PR effort.
+- The test_test_tiers.jl expected file list was initially missed, causing a stale
+  test failure on the first full test run.
+
+### Action Items for Next PR
+
+- Always update test_test_tiers.jl when adding new test files (easy to forget)
+- For performance investigation PRs, run a quick @code_warntype check BEFORE
+  creating the task to validate that the suspected instability exists

scripts/benchmark_typed_kkt_dicts.jl

@@ -0,0 +1,113 @@
+#=
+Benchmark script for Perf #7: Dict{Int, Any} type investigation in KKT construction.
+
+Finding: All Dict{Int, Any} instances in qp_kkt.jl and nonlinear_kkt.jl are COLD PATH
+(one-time symbolic construction). The hot-path containers were already typed:
+- Vector{Union{Matrix{Float64}, Nothing}} for K/M/N evals
+- Dict{Int, Matrix{Float64}} for M/N buffers
+- Vector{MNFunctionWrapper} for compiled M/N functions
+- Vector{Int} for π_sizes
+
+This benchmark verifies that the hot-path dict access pattern (Dict{Int, Matrix{Float64}})
+is not a bottleneck compared to the dominant cost of matrix operations.
+=#
+
+using MixedHierarchyGames
+using Graphs: SimpleDiGraph, add_edge!
+using BenchmarkTools
+
+# Setup: 2-player Stackelberg hierarchy
+G = SimpleDiGraph(2)
+add_edge!(G, 1, 2)
+
+state_dim = 4
+control_dim = 2
+T = 5
+primal_dim = (state_dim + control_dim) * (T + 1)
+primal_dims = fill(primal_dim, 2)
+
+backend = default_backend()
+θs = setup_problem_parameter_variables(fill(state_dim, 2); backend)
+Δt = 0.1
+
+Js = Dict(
+    1 => (z1, z2; θ=nothing) -> begin
+        (; xs, us) = MixedHierarchyGames.unflatten_trajectory(z1, state_dim, control_dim)
+        sum(sum(x.^2) for x in xs) + 0.1 * sum(sum(u.^2) for u in us)
+    end,
+    2 => (z1, z2; θ=nothing) -> begin
+        (; xs, us) = MixedHierarchyGames.unflatten_trajectory(z2, state_dim, control_dim)
+        sum(sum(x.^2) for x in xs) + 0.1 * sum(sum(u.^2) for u in us)
+    end
+)
+
+gs = [
+    z -> begin
+        (; xs, us) = MixedHierarchyGames.unflatten_trajectory(z, state_dim, control_dim)
+        dyn = mapreduce(vcat, 1:T) do t; xs[t+1] - xs[t] - Δt * us[t] end
+        ic = xs[1] - θs[1]
+        vcat(dyn, ic)
+    end,
+    z -> begin
+        (; xs, us) = MixedHierarchyGames.unflatten_trajectory(z, state_dim, control_dim)
+        dyn = mapreduce(vcat, 1:T) do t; xs[t+1] - xs[t] - Δt * us[t] end
+        ic = xs[1] - θs[2]
+        vcat(dyn, ic)
+    end
+]
+
+println("Building NonlinearSolver (one-time cost)...")
+solver = NonlinearSolver(G, Js, gs, primal_dims, θs, state_dim, control_dim)
+
+initial_states = Dict(1 => randn(state_dim), 2 => randn(state_dim))
+
+println("\n=== Benchmark: Full solve (includes hot-path typed containers) ===")
+# Warmup
+solve_raw(solver, initial_states; max_iters=5)
+
+b = @benchmark solve_raw($solver, $initial_states; max_iters=20) samples=20
+display(b)
+
+println("\n=== Benchmark: compute_K_evals (hot-path dict access) ===")
+problem_vars = solver.precomputed.problem_vars
+setup_info = solver.precomputed.setup_info
+z_current = zeros(length(solver.precomputed.all_variables))
+
+# Pre-allocate buffers (as done in run_nonlinear_solver)
+N_players = 2
+M_buffers = Dict{Int, Matrix{Float64}}()
+N_buffers = Dict{Int, Matrix{Float64}}()
+k_eval_buffers = (;
+    M_evals = Vector{Union{Matrix{Float64}, Nothing}}(nothing, N_players),
+    N_evals = Vector{Union{Matrix{Float64}, Nothing}}(nothing, N_players),
+    K_evals = Vector{Union{Matrix{Float64}, Nothing}}(nothing, N_players),
+    follower_cache = Vector{Union{Vector{Int}, Nothing}}(nothing, N_players),
+    buffer_cache = Vector{Union{Vector{Float64}, Nothing}}(nothing, N_players),
+    all_K_vec = Vector{Float64}(undef, solver.precomputed.mcp_obj.parameter_dimension - sum(length(initial_states[k]) for k in keys(initial_states))),
+)
+
+# Warmup
+compute_K_evals(z_current, problem_vars, setup_info; M_buffers, N_buffers, buffers=k_eval_buffers)
+
+b2 = @benchmark compute_K_evals($z_current, $problem_vars, $setup_info;
+    M_buffers=$M_buffers, N_buffers=$N_buffers, buffers=$k_eval_buffers) samples=100
+display(b2)
+
+println("\n=== Dict{Int, Matrix{Float64}} access micro-benchmark ===")
+# Show that Dict access with concrete value types is fast
+typed_dict = Dict{Int, Matrix{Float64}}(2 => randn(10, 10))
+untyped_dict = Dict{Int, Any}(2 => randn(10, 10))
+
+println("Typed Dict{Int, Matrix{Float64}} access:")
+b3 = @benchmark $typed_dict[2]
+display(b3)
+
+println("\nUntyped Dict{Int, Any} access:")
+b4 = @benchmark $untyped_dict[2]
+display(b4)
+
+println("\n=== Summary ===")
+println("All Dict{Int, Any} in KKT construction are cold-path (one-time symbolic setup).")
+println("Hot-path containers use concrete types: Vector{Union{Matrix{Float64}, Nothing}},")
+println("Dict{Int, Matrix{Float64}}, Vector{MNFunctionWrapper}, Vector{Int}.")
+println("No performance improvement from this PR — the hot path was already correctly typed.")

src/nonlinear_kkt.jl

@@ -473,7 +473,7 @@ function preoptimize_nonlinear_solver(
 
         # Strip policy constraints for MCP construction
         πs_solve = strip_policy_constraints(πs, hierarchy_graph, zs, gs)
-        π_sizes_trimmed = Dict(ii => length(πs_solve[ii]) for ii in keys(πs_solve))
+        π_sizes_trimmed = Dict{Int, Int}(ii => length(πs_solve[ii]) for ii in keys(πs_solve))
 
         # Build MCP function vector
         π_order = ordered_player_indices(πs_solve)

test/test_test_tiers.jl

@@ -39,6 +39,7 @@
             "test_allocation_optimization.jl",
             "test_jacobian_buffer_safety.jl",
             "test_regularization.jl",
+            "test_typed_kkt_dicts.jl",
         ]
         @test Set(SLOW_TEST_FILES) == Set(expected_slow)
     end
@@ -78,6 +79,7 @@
             "test_regularization.jl",
             "test_unified_interface.jl",
             "test_vcat_ordered.jl",
+            "test_typed_kkt_dicts.jl",
         ])
         @test all_classified == expected_all
     end

test/test_tiers.jl

@@ -37,6 +37,7 @@ const SLOW_TEST_FILES = [
     "test_allocation_optimization.jl",
     "test_jacobian_buffer_safety.jl",
     "test_regularization.jl",
+    "test_typed_kkt_dicts.jl",
 ]
 
 """