Use sp_free instead of free

include/utils/Vec_macros.h

@@ -36,6 +36,7 @@
 #ifndef VEC_MACROS_H
 #define VEC_MACROS_H
 
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -60,12 +61,12 @@
     PSLP_UNUSED static TYPE_NAME##Vec *TYPE_NAME##Vec_new(int capacity)             \
     {                                                                               \
         assert(capacity > 0);                                                       \
-        TYPE_NAME##Vec *vec = (TYPE_NAME##Vec *) malloc(sizeof(TYPE_NAME##Vec));    \
+        TYPE_NAME##Vec *vec = (TYPE_NAME##Vec *) sp_malloc(sizeof(TYPE_NAME##Vec)); \
         if (vec == NULL) return NULL;                                               \
-        vec->data = (TYPE *) malloc(capacity * sizeof(TYPE));                       \
+        vec->data = (TYPE *) sp_malloc(capacity * sizeof(TYPE));                    \
         if (vec->data == NULL)                                                      \
         {                                                                           \
-            free(vec);                                                              \
+            sp_free(vec);                                                           \
             return NULL;                                                            \
         }                                                                           \
                                                                                     \
@@ -76,8 +77,8 @@
                                                                                     \
     static inline void TYPE_NAME##Vec_free(TYPE_NAME##Vec *vec)                     \
     {                                                                               \
-        free(vec->data);                                                            \
-        free(vec);                                                                  \
+        sp_free(vec->data);                                                         \
+        sp_free(vec);                                                               \
     }                                                                               \
                                                                                     \
     static inline void TYPE_NAME##Vec_clear_no_resize(TYPE_NAME##Vec *vec)          \
@@ -91,8 +92,8 @@
         {                                                                           \
             vec->capacity *= 2;                                                     \
             assert(vec->capacity > 0);                                              \
-            TYPE *temp =                                                            \
-                (TYPE *) realloc(vec->data, (size_t) vec->capacity * sizeof(TYPE)); \
+            TYPE *temp = (TYPE *) sp_realloc(vec->data, (size_t) vec->capacity *    \
+                                                            sizeof(TYPE));          \
             if (temp == NULL)                                                       \
             {                                                                       \
                 TYPE_NAME##Vec_free(vec);                                           \
@@ -117,8 +118,8 @@
                 new_capacity *= 2;                                                  \
             }                                                                       \
                                                                                     \
-            TYPE *temp =                                                            \
-                (TYPE *) realloc(vec->data, (size_t) new_capacity * sizeof(TYPE));  \
+            TYPE *temp = (TYPE *) sp_realloc(vec->data,                             \
+                                             (size_t) new_capacity * sizeof(TYPE)); \
             if (temp == NULL)                                                       \
             {                                                                       \
                 TYPE_NAME##Vec_free(vec);                                           \

include/utils/tracked_alloc.h

@@ -77,4 +77,22 @@ static inline void sp_free(void *ptr)
     }
 }
 
+static inline void *sp_realloc(void *ptr, size_t size)
+{
+    size_t old_block_size = ptr ? TRACKED_BLOCK_SIZE(ptr) : 0;
+    void *new_ptr = realloc(ptr, size);
+    if (new_ptr)
+    {
+        g_allocated_bytes =
+            g_allocated_bytes - old_block_size + TRACKED_BLOCK_SIZE(new_ptr);
+        if (g_allocated_bytes > g_peak_bytes)
+        {
+            g_peak_bytes = g_allocated_bytes;
+        }
+    }
+    /* realloc returning NULL means failure — original block is still live
+       per C standard, so no counter adjustment. */
+    return new_ptr;
+}
+
 #endif /* TRACKED_ALLOC_H */

src/atoms/affine/diag_mat.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdlib.h>
 
@@ -28,13 +29,13 @@ expr *new_diag_mat(expr *child)
     assert(child->d1 == child->d2);
     int n = child->d1;
 
-    int *indices = (int *) malloc(n * sizeof(int));
+    int *indices = (int *) sp_malloc(n * sizeof(int));
     for (int i = 0; i < n; i++)
     {
         indices[i] = i * (n + 1);
     }
 
     expr *node = new_index(child, n, 1, indices, n);
-    free(indices);
+    sp_free(indices);
     return node;
 }

src/atoms/affine/hstack.c

@@ -174,7 +174,7 @@ static void free_type_data(expr *node)
 
     free_CSR_matrix(hnode->CSR_work);
     hnode->CSR_work = NULL;
-    free(hnode->args);
+    sp_free(hnode->args);
     hnode->args = NULL;
 }
 

src/atoms/affine/index.c

@@ -39,7 +39,7 @@ static bool check_for_duplicates(const int *indices, int n_idxs, int max_idx)
         }
         seen[indices[i]] = true;
     }
-    free(seen);
+    sp_free(seen);
     return has_dup;
 }
 
@@ -139,7 +139,7 @@ static void free_type_data(expr *node)
     index_expr *idx = (index_expr *) node;
     if (idx->indices)
     {
-        free(idx->indices);
+        sp_free(idx->indices);
         idx->indices = NULL;
     }
 }

src/atoms/affine/left_matmul.c

@@ -98,7 +98,7 @@ static void free_type_data(expr *node)
     free_matrix(lnode->AT);
     free_CSC_matrix(lnode->Jchild_CSC);
     free_CSC_matrix(lnode->J_CSC);
-    free(lnode->csc_to_csr_work);
+    sp_free(lnode->csc_to_csr_work);
     free_expr(lnode->param_source);
 
     lnode->A = NULL;

src/atoms/affine/right_matmul.c

@@ -76,7 +76,7 @@ expr *new_right_matmul(expr *param_node, expr *u, const CSR_matrix *A)
     expr *node = new_transpose(left_matmul);
 
     free_CSR_matrix(AT);
-    free(work_transpose);
+    sp_free(work_transpose);
     return node;
 }
 
@@ -107,7 +107,7 @@ expr *new_right_matmul_dense(expr *param_node, expr *u, int m, int n,
         double *AT = (double *) sp_malloc(n * m * sizeof(double));
         A_transpose(AT, data, m, n);
         left_matmul_node = new_left_matmul_dense(NULL, u_transpose, n, m, AT);
-        free(AT);
+        sp_free(AT);
     }
 
     return new_transpose(left_matmul_node);

src/atoms/affine/sum.c

@@ -188,7 +188,7 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     sum_expr *snode = (sum_expr *) node;
-    free(snode->idx_map);
+    sp_free(snode->idx_map);
 }
 
 expr *new_sum(expr *child, int axis)

src/atoms/affine/trace.c

@@ -142,7 +142,7 @@ static void free_type_data(expr *node)
     if (node)
     {
         trace_expr *tnode = (trace_expr *) node;
-        free(tnode->idx_map);
+        sp_free(tnode->idx_map);
     }
 }
 

src/atoms/affine/upper_tri.c

@@ -16,6 +16,7 @@
  * limitations under the License.
  */
 #include "atoms/affine.h"
+#include "utils/tracked_alloc.h"
 #include <assert.h>
 #include <stdlib.h>
 
@@ -40,7 +41,7 @@ expr *new_upper_tri(expr *child)
     int *indices = NULL;
     if (n_elems > 0)
     {
-        indices = (int *) malloc(n_elems * sizeof(int));
+        indices = (int *) sp_malloc(n_elems * sizeof(int));
         int k = 0;
         for (int i = 0; i < n; i++)
         {
@@ -53,6 +54,6 @@ expr *new_upper_tri(expr *child)
     }
 
     expr *node = new_index(child, n_elems, 1, indices, n_elems);
-    free(indices);
+    sp_free(indices);
     return node;
 }

src/atoms/affine/vstack.c

@@ -42,7 +42,7 @@ expr *new_vstack(expr **args, int n_args, int n_vars)
     }
 
     expr *hstacked = new_hstack(transposed, n_args, n_vars);
-    free(transposed);
+    sp_free(transposed);
 
     return new_transpose(hstacked);
 }

src/atoms/bivariate_full_dom/matmul.c

@@ -113,13 +113,13 @@ static void free_matmul_data(expr *node)
     free_CSR_matrix(mnode->B);
     free_CSR_matrix(mnode->BJg);
     free_CSC_matrix(mnode->BJg_CSC);
-    free(mnode->BJg_csc_work);
+    sp_free(mnode->BJg_csc_work);
     free_CSR_matrix(mnode->C);
     free_CSR_matrix(mnode->CT);
-    free(mnode->idx_map_C);
-    free(mnode->idx_map_CT);
-    free(mnode->idx_map_Hf);
-    free(mnode->idx_map_Hg);
+    sp_free(mnode->idx_map_C);
+    sp_free(mnode->idx_map_CT);
+    sp_free(mnode->idx_map_Hf);
+    sp_free(mnode->idx_map_Hg);
 }
 
 static bool is_affine(const expr *node)

src/atoms/bivariate_full_dom/multiply.c

@@ -303,10 +303,10 @@ static void free_type_data(expr *node)
     elementwise_mult_expr *mul_node = (elementwise_mult_expr *) node;
     free_matrix(mul_node->C);
     free_matrix(mul_node->CT);
-    free(mul_node->idx_map_C);
-    free(mul_node->idx_map_CT);
-    free(mul_node->idx_map_Hx);
-    free(mul_node->idx_map_Hy);
+    sp_free(mul_node->idx_map_C);
+    sp_free(mul_node->idx_map_CT);
+    sp_free(mul_node->idx_map_Hx);
+    sp_free(mul_node->idx_map_Hy);
 }
 
 static bool is_affine(const expr *node)

src/atoms/bivariate_restricted_dom/quad_over_lin.c

@@ -105,7 +105,7 @@ static void jacobian_init_impl(expr *node)
         }
         assert(nonzero_cols == jac->nnz);
 
-        free(col_nz);
+        sp_free(col_nz);
 
         /* insert y variable index at correct position */
         insert_idx(y->var_id, jac->i, jac->nnz);

src/atoms/other/prod_axis_one.c

@@ -391,9 +391,9 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     prod_axis *pnode = (prod_axis *) node;
-    free(pnode->num_of_zeros);
-    free(pnode->zero_index);
-    free(pnode->prod_nonzero);
+    sp_free(pnode->num_of_zeros);
+    sp_free(pnode->zero_index);
+    sp_free(pnode->prod_nonzero);
 }
 
 expr *new_prod_axis_one(expr *child)

src/atoms/other/prod_axis_zero.c

@@ -350,9 +350,9 @@ static bool is_affine(const expr *node)
 static void free_type_data(expr *node)
 {
     prod_axis *pnode = (prod_axis *) node;
-    free(pnode->num_of_zeros);
-    free(pnode->zero_index);
-    free(pnode->prod_nonzero);
+    sp_free(pnode->num_of_zeros);
+    sp_free(pnode->zero_index);
+    sp_free(pnode->prod_nonzero);
 }
 
 /* TODO: refactor to remove diagonal entry as nonzero since it's always zero */

src/expr.c

@@ -75,25 +75,25 @@ void free_expr(expr *node)
     }
 
     /* free value array and derivative matrices */
-    free(node->value);
+    sp_free(node->value);
     free_matrix(node->jacobian);
     free_matrix(node->wsum_hess);
 
     /* free workspace */
     if (node->work)
     {
-        free(node->work->dwork);
-        free(node->work->iwork);
+        sp_free(node->work->dwork);
+        sp_free(node->work->iwork);
         free_CSC_matrix(node->work->jacobian_csc);
-        free(node->work->csc_work);
-        free(node->work->local_jac_diag);
+        sp_free(node->work->csc_work);
+        sp_free(node->work->local_jac_diag);
         free_matrix(node->work->hess_term1);
         free_matrix(node->work->hess_term2);
-        free(node->work);
+        sp_free(node->work);
     }
 
     /* free the node itself */
-    free(node);
+    sp_free(node);
 }
 
 void jacobian_init(expr *node)

src/old-code/linear_op.c

@@ -55,7 +55,7 @@ static void free_type_data(expr *node)
     linear_op_expr *lin_node = (linear_op_expr *) node;
     if (lin_node->b != NULL)
     {
-        free(lin_node->b);
+        sp_free(lin_node->b);
         lin_node->b = NULL;
     }
 }

src/old-code/old_permuted_dense.c

@@ -58,8 +58,8 @@ matrix *BTA_pd_csr_alloc(const permuted_dense *B, const CSR_matrix *A)
 
     matrix *C =
         new_permuted_dense(B->base.n, p, B->n0, s_A, B->col_perm, col_active, NULL);
-    free(col_active);
-    free(seen);
+    sp_free(col_active);
+    sp_free(seen);
 
     /* Upgrade `dwork` (currently sized for the Y-role at m0_C * n0_C = B->n0 *
        s_A) to fit the gather buffer A_sub_dense used by BTA_csr_pd /
@@ -69,7 +69,7 @@ matrix *BTA_pd_csr_alloc(const permuted_dense *B, const CSR_matrix *A)
     size_t gather_size = B->m0 * s_A;
     if (gather_size > C_pd->kernel_dwork_size)
     {
-        free(C_pd->kernel_dwork);
+        sp_free(C_pd->kernel_dwork);
         C_pd->kernel_dwork_size = gather_size;
         C_pd->kernel_dwork = (double *) sp_calloc(gather_size, sizeof(double));
     }
@@ -191,8 +191,8 @@ matrix *BTA_csr_pd_alloc(const CSR_matrix *B_csr, const permuted_dense *A)
 
     matrix *C =
         new_permuted_dense(q, A->base.n, r_B, A->n0, row_active, A->col_perm, NULL);
-    free(row_active);
-    free(seen);
+    sp_free(row_active);
+    sp_free(seen);
 
     /* Upgrade `dwork` (currently sized for the Y-role at m0_C * n0_C = r_B *
        A->n0) to fit the gather buffer B_sub_dense used by BTA_csr_pd /
@@ -201,7 +201,7 @@ matrix *BTA_csr_pd_alloc(const CSR_matrix *B_csr, const permuted_dense *A)
     size_t gather_size = A->m0 * r_B;
     if (gather_size > C_pd->kernel_dwork_size)
     {
-        free(C_pd->kernel_dwork);
+        sp_free(C_pd->kernel_dwork);
         C_pd->kernel_dwork_size = gather_size;
         C_pd->kernel_dwork = (double *) sp_calloc(gather_size, sizeof(double));
     }