Size the gathered Jacobian by the true selected-row nnz sum (#105)

sparse_index_alloc capped the output capacity at the source's total nnz,
a bound that only holds for distinct (selection/permutation) indices.
A gather with enough DUPLICATED indices selects more entries than the
source has (e.g. x[[0,0,1,2,2,3]] on a size-4 variable: 6 > 4), and the
fill loop wrote the excess past the end of the buffers, corrupting the
heap (cvxpy#3442 — DNLP segfault during derivative-structure init).

Sum the selected rows' nnz for an exact allocation. The sum replaces the
sat_mul_int dense bound; a checked add errors out before wrapping if the
true count is not CSR-representable.

New tests (verified to hit heap-buffer-overflow under ASan without the
fix): the gather Jacobian itself, and the jacobian + weighted-sum
Hessian of multiply(a[idx], b[idx]) mirroring the issue.

Co-authored-by: Transurgeon <peter.zijie@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

src/utils/sparse_matrix.c

@@ -24,6 +24,7 @@
 #include "utils/mini_numpy.h"
 #include "utils/tracked_alloc.h"
 #include "utils/utils.h"
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -124,8 +125,24 @@ static void sparse_transpose_fill_values(const matrix *self, matrix *out)
 static matrix *sparse_index_alloc(matrix *self, const int *indices, int n_idxs)
 {
     CSR_matrix *Jx = ((sparse_matrix *) self)->csr;
-    int alloc_ub = sat_mul_int(n_idxs, self->n);
-    CSR_matrix *J = new_CSR_matrix(n_idxs, self->n, MIN(Jx->nnz, alloc_ub));
+
+    /* Exact output nnz: sum the selected rows' nnz. Jx->nnz is NOT an upper
+       bound — duplicated indices select the same source row more than once
+       (cvxpy#3442). Duplicated gathers of dense rows can push the true count
+       past INT_MAX, which a CSR cannot represent, so fail before wrapping. */
+    int nnz = 0;
+    for (int i = 0; i < n_idxs; i++)
+    {
+        int len = Jx->p[indices[i] + 1] - Jx->p[indices[i]];
+        if (len > INT_MAX - nnz)
+        {
+            fprintf(stderr, "Error in sparse_index_alloc: gathered nnz "
+                            "exceeds INT_MAX.\n");
+            exit(1);
+        }
+        nnz += len;
+    }
+    CSR_matrix *J = new_CSR_matrix(n_idxs, self->n, nnz);
 
     J->p[0] = 0;
     for (int i = 0; i < n_idxs; i++)

tests/all_tests.c

@@ -194,6 +194,7 @@ int main(void)
     mu_run_test(test_jacobian_elementwise_mult_2, tests_run);
     mu_run_test(test_jacobian_elementwise_mult_3, tests_run);
     mu_run_test(test_jacobian_elementwise_mult_4, tests_run);
+    mu_run_test(test_jacobian_elementwise_mult_duplicate_gathers, tests_run);
     mu_run_test(test_quad_over_lin1, tests_run);
     mu_run_test(test_quad_over_lin2, tests_run);
     mu_run_test(test_quad_over_lin3, tests_run);
@@ -223,6 +224,7 @@ int main(void)
     mu_run_test(test_index_jacobian_of_variable, tests_run);
     mu_run_test(test_index_jacobian_of_log, tests_run);
     mu_run_test(test_index_jacobian_repeated, tests_run);
+    mu_run_test(test_index_jacobian_duplicates_exceed_source_nnz, tests_run);
     mu_run_test(test_sum_of_index, tests_run);
     mu_run_test(test_promote_scalar_jacobian, tests_run);
     mu_run_test(test_promote_scalar_to_matrix_jacobian, tests_run);
@@ -307,6 +309,7 @@ int main(void)
     mu_run_test(test_wsum_hess_multiply_sparse_random, tests_run);
     mu_run_test(test_wsum_hess_multiply_1, tests_run);
     mu_run_test(test_wsum_hess_multiply_2, tests_run);
+    mu_run_test(test_wsum_hess_multiply_duplicate_gathers, tests_run);
     mu_run_test(test_wsum_hess_left_matmul, tests_run);
     mu_run_test(test_wsum_hess_left_matmul_matrix, tests_run);
     mu_run_test(test_wsum_hess_left_matmul_exp_composite, tests_run);

tests/jacobian_tests/affine/test_index.h

@@ -116,6 +116,32 @@ const char *test_index_jacobian_repeated(void)
     return 0;
 }
 
+const char *test_index_jacobian_duplicates_exceed_source_nnz(void)
+{
+    /* x[[0,0,1,2,2,3]] on a size-4 variable: the gathered Jacobian has 6
+     * nonzeros, MORE than the source identity's 4. Pins the allocator sizing
+     * the output by the true per-row sum instead of the source nnz
+     * (cvxpy#3442 — undersized buffers overflowed the heap on the extra rows). */
+    double u[4] = {1.0, 2.0, 3.0, 4.0};
+    int indices[6] = {0, 0, 1, 2, 2, 3};
+    expr *var = new_variable(4, 1, 0, 4);
+    expr *idx = new_index(var, 1, 6, indices, 6);
+    idx->forward(idx, u);
+    jacobian_init(idx);
+    idx->eval_jacobian(idx);
+
+    double expected_x[6] = {1.0, 1.0, 1.0, 1.0, 1.0, 1.0};
+    int expected_p[7] = {0, 1, 2, 3, 4, 5, 6};
+    int expected_i[6] = {0, 0, 1, 2, 2, 3};
+
+    mu_assert("vals fail", cmp_values(idx->jacobian, expected_x, 6));
+    mu_assert("sparsity fail",
+              cmp_sparsity(idx->jacobian, expected_p, expected_i, 6, 6));
+
+    free_expr(idx);
+    return 0;
+}
+
 const char *test_sum_of_index(void)
 {
     /* sum(x[0, 2]) = x[0] + x[2]

tests/jacobian_tests/bivariate_full_dom/test_elementwise_mult.h

@@ -118,6 +118,38 @@ const char *test_jacobian_elementwise_mult_3(void)
     return 0;
 }
 
+const char *test_jacobian_elementwise_mult_duplicate_gathers(void)
+{
+    // var = (a, b) where a is 4 x 1 at offset 0 and b is 4 x 1 at offset 4.
+    // We compute the jacobian of a[idx] * b[idx] with DUPLICATED gather
+    // indices idx = [0, 0, 1, 2, 2, 3] (cvxpy#3442): each gather Jacobian has
+    // 6 nnz, more than its source variable's 4.
+    double u_vals[8] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
+    int indices[6] = {0, 0, 1, 2, 2, 3};
+    expr *a = new_variable(4, 1, 0, 8);
+    expr *b = new_variable(4, 1, 4, 8);
+    expr *ga = new_index(a, 1, 6, indices, 6);
+    expr *gb = new_index(b, 1, 6, indices, 6);
+    expr *node = new_elementwise_mult(ga, gb);
+
+    node->forward(node, u_vals);
+    jacobian_init(node);
+    node->eval_jacobian(node);
+
+    double fwd[6] = {5.0, 5.0, 12.0, 21.0, 21.0, 32.0};
+    mu_assert("forward fail", cmp_double_array(node->value, fwd, 6));
+
+    /* row k: b[idx[k]] at col idx[k], a[idx[k]] at col 4 + idx[k] */
+    double vals[12] = {5.0, 1.0, 5.0, 1.0, 6.0, 2.0, 7.0, 3.0, 7.0, 3.0, 8.0, 4.0};
+    int rows[7] = {0, 2, 4, 6, 8, 10, 12};
+    int cols[12] = {0, 4, 0, 4, 1, 5, 2, 6, 2, 6, 3, 7};
+
+    mu_assert("vals fail", cmp_values(node->jacobian, vals, 12));
+    mu_assert("sparsity fail", cmp_sparsity(node->jacobian, rows, cols, 6, 12));
+    free_expr(node);
+    return 0;
+}
+
 const char *test_jacobian_elementwise_mult_4(void)
 {
     // var = (z, x, w, y) where z is 2 x 1, x is 3 x 1, w is 2 x 1, y is 3 x 1

tests/wsum_hess/bivariate_full_dom/test_multiply.h

@@ -188,6 +188,40 @@ const char *test_wsum_hess_multiply_linear_ops(void)
     return 0;
 }
 
+const char *test_wsum_hess_multiply_duplicate_gathers(void)
+{
+    /* Hessian of w' (a[idx] * b[idx]) with DUPLICATED gather indices
+     * idx = [0, 0, 1, 2, 2, 3] (cvxpy#3442). a at offset 0, b at offset 4.
+     * d²/da[p] db[p] accumulates w_k over every k with idx[k] == p:
+     * H[p, 4+p] = {1+2, 3, 4+5, 6} = {3, 3, 9, 6}, plus the symmetric block. */
+    double u_vals[8] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0};
+    double w[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+    int indices[6] = {0, 0, 1, 2, 2, 3};
+
+    expr *a = new_variable(4, 1, 0, 8);
+    expr *b = new_variable(4, 1, 4, 8);
+    expr *ga = new_index(a, 1, 6, indices, 6);
+    expr *gb = new_index(b, 1, 6, indices, 6);
+    expr *node = new_elementwise_mult(ga, gb);
+
+    node->forward(node, u_vals);
+    jacobian_init(node);
+    node->eval_jacobian(node);
+    wsum_hess_init(node);
+    node->eval_wsum_hess(node, w);
+
+    int expected_p[9] = {0, 1, 2, 3, 4, 5, 6, 7, 8};
+    int expected_i[8] = {4, 5, 6, 7, 0, 1, 2, 3};
+    double expected_x[8] = {3.0, 3.0, 9.0, 6.0, 3.0, 3.0, 9.0, 6.0};
+
+    mu_assert("sparsity fail",
+              cmp_sparsity(node->wsum_hess, expected_p, expected_i, 8, 8));
+    mu_assert("vals fail", cmp_values(node->wsum_hess, expected_x, 8));
+
+    free_expr(node);
+    return 0;
+}
+
 const char *test_wsum_hess_multiply_2(void)
 {
     // Total 12 variables: [?, ?, ?, y0, y1, y2, ?, ?, x0, x1, x2, ?]