From 7231d35fc9ee7013ac0175883353071c86f75476 Mon Sep 17 00:00:00 2001
From: dance858 <danielcederberg1@gmail.com>
Date: Sun, 30 Aug 2026 20:12:03 -0700
Subject: [PATCH] affine optimization

---
 src/expr.c                                 | 11 +++---
 tests/all_tests.c                          |  1 +
 tests/jacobian_tests/test_values_version.h | 45 ++++++++++++++++++++++
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/src/expr.c b/src/expr.c
index 8e0ef2c..0990727 100644
--- a/src/expr.c
+++ b/src/expr.c
@@ -130,16 +130,15 @@ void wsum_hess_init(expr *node)
 
 /* Runs the atom's eval and bumps the Jacobian's values_version — except
    for an affine node already evaluated this parameter epoch: its values
-   are constant, so consumers correctly see "unchanged". (The impl still
-   runs; skipping it entirely is a planned follow-up.) */
+   are constant, so the impl (and with it the whole subtree walk) is
+   skipped and consumers correctly see "unchanged". expr_set_needs_refresh
+   re-arms the eval. */
 void eval_jacobian(expr *node)
 {
     if (node == NULL) return;
+    if (node->work->is_affine_cached && node->work->jacobian_evaluated) return;
     node->eval_jacobian_impl(node);
-    if (!(node->work->is_affine_cached && node->work->jacobian_evaluated))
-    {
-        matrix_values_changed(node->jacobian);
-    }
+    matrix_values_changed(node->jacobian);
     node->work->jacobian_evaluated = true;
 }
 
diff --git a/tests/all_tests.c b/tests/all_tests.c
index c6f137d..e1bb557 100644
--- a/tests/all_tests.c
+++ b/tests/all_tests.c
@@ -212,6 +212,7 @@ int main(void)
     mu_run_test(test_quad_form, tests_run);
     mu_run_test(test_values_version_non_affine, tests_run);
     mu_run_test(test_values_version_affine, tests_run);
+    mu_run_test(test_impl_skip_affine, tests_run);
     mu_run_test(test_values_version_csc_mirror_dedup, tests_run);
     mu_run_test(test_values_version_stacked_pd_to_csr, tests_run);
     mu_run_test(test_values_version_param_under_hstack, tests_run);
diff --git a/tests/jacobian_tests/test_values_version.h b/tests/jacobian_tests/test_values_version.h
index 74ffbba..3b352b2 100644
--- a/tests/jacobian_tests/test_values_version.h
+++ b/tests/jacobian_tests/test_values_version.h
@@ -60,6 +60,51 @@ const char *test_values_version_affine(void)
     return 0;
 }
 
+/* Affine node: once evaluated this parameter epoch, eval_jacobian skips the
+ * impl entirely — proven by poisoning the jacobian values and observing that
+ * a re-eval does not rewrite them. expr_set_needs_refresh re-arms the eval,
+ * which then restores the true values. */
+const char *test_impl_skip_affine(void)
+{
+    double u[3] = {0.1, 0.2, 0.3};
+    expr *x = new_variable(3, 1, 0, 3);
+    expr *m = new_neg(x);
+
+    jacobian_init(m);
+    m->forward(m, u);
+    eval_jacobian(m);
+
+    uint64_t v1 = m->jacobian->values_version;
+    int nnz = m->jacobian->nnz;
+    mu_assert("neg jacobian must have entries", nnz == 3);
+    for (int ii = 0; ii < nnz; ii++)
+    {
+        mu_assert("neg jacobian value must be -1", m->jacobian->x[ii] == -1.0);
+        m->jacobian->x[ii] = 42.0; /* poison */
+    }
+
+    eval_jacobian(m);
+    mu_assert("re-eval of affine node must not bump",
+              m->jacobian->values_version == v1);
+    for (int ii = 0; ii < nnz; ii++)
+    {
+        mu_assert("impl must not have run (poison must survive)",
+                  m->jacobian->x[ii] == 42.0);
+    }
+
+    expr_set_needs_refresh(m);
+    eval_jacobian(m);
+    mu_assert("eval after refresh must bump", m->jacobian->values_version == v1 + 1);
+    for (int ii = 0; ii < nnz; ii++)
+    {
+        mu_assert("eval after refresh must restore true values",
+                  m->jacobian->x[ii] == -1.0);
+    }
+
+    free_expr(m);
+    return 0;
+}
+
 /* sparse_matrix CSC mirror: refresh_csc_values fills once per write and
  * dedupes repeated calls with no write in between. */
 const char *test_values_version_csc_mirror_dedup(void)