From 78b9963cb22c78f777a8301a199eea243a911448 Mon Sep 17 00:00:00 2001
From: govindchari <govindchari1@gmail.com>
Date: Sat, 8 Nov 2025 17:50:43 -0800
Subject: [PATCH] WIP

---
 algebra/builtin/builtin_linalg.c | 137 ++++++++++------
 algebra/builtin/builtin_types.h  |  36 ++++
 devtools/run_tests.sh            |   2 +-
 include/common_linalg.h          |  69 +++++++-
 include/kkt.h                    |   6 +
 include/qoco_linalg.h            | 128 +++++++++++----
 include/structs.h                |  30 ++--
 src/common_linalg.c              |  78 ++++++++-
 src/kkt.c                        | 273 ++++++++++++++++++++++---------
 src/qoco_api.c                   | 176 ++++++++++----------
 10 files changed, 666 insertions(+), 269 deletions(-)
 create mode 100644 algebra/builtin/builtin_types.h

diff --git a/algebra/builtin/builtin_linalg.c b/algebra/builtin/builtin_linalg.c
index 153ff3e..9906900 100644
--- a/algebra/builtin/builtin_linalg.c
+++ b/algebra/builtin/builtin_linalg.c
@@ -8,11 +8,12 @@
  * This source code is licensed under the BSD 3-Clause License
  */
 
-#include "qoco_linalg.h"
+#include "builtin_types.h"
 
-QOCOCscMatrix* new_qoco_csc_matrix(const QOCOCscMatrix* A)
+QOCOMatrix* new_qoco_matrix(const QOCOCscMatrix* A)
 {
-  QOCOCscMatrix* M = qoco_malloc(sizeof(QOCOCscMatrix));
+  QOCOMatrix* M = qoco_malloc(sizeof(QOCOMatrix));
+  QOCOCscMatrix* Mcsc = qoco_malloc(sizeof(QOCOCscMatrix));
 
   if (A) {
     QOCOInt m = A->m;
@@ -27,61 +28,88 @@ QOCOCscMatrix* new_qoco_csc_matrix(const QOCOCscMatrix* A)
     copy_arrayi(A->i, i, nnz);
     copy_arrayi(A->p, p, n + 1);
 
-    M->m = m;
-    M->n = n;
-    M->nnz = nnz;
-    M->x = x;
-    M->i = i;
-    M->p = p;
+    Mcsc->m = m;
+    Mcsc->n = n;
+    Mcsc->nnz = nnz;
+    Mcsc->x = x;
+    Mcsc->i = i;
+    Mcsc->p = p;
   }
   else {
-    M->m = 0;
-    M->n = 0;
-    M->nnz = 0;
-    M->x = NULL;
-    M->i = NULL;
-    M->p = NULL;
+    Mcsc->m = 0;
+    Mcsc->n = 0;
+    Mcsc->nnz = 0;
+    Mcsc->x = NULL;
+    Mcsc->i = NULL;
+    Mcsc->p = NULL;
   }
 
+  M->csc = M;
+
   return M;
 }
 
-void free_qoco_csc_matrix(QOCOCscMatrix* A)
+QOCOVectorf* new_qoco_vectorf(const QOCOFloat* x, QOCOInt n)
 {
-  free(A->x);
-  free(A->i);
-  free(A->p);
-  free(A);
+  QOCOVectorf* v = qoco_malloc(sizeof(QOCOVectorf));
+  QOCOFloat* vdata = qoco_malloc(sizeof(QOCOFloat) * n);
+  copy_arrayf(x, vdata, n);
+
+  v->len = n;
+  v->data = vdata;
 }
 
-void copy_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n)
+void free_qoco_matrix(QOCOMatrix* A)
 {
-  qoco_assert(x || n == 0);
-  qoco_assert(y || n == 0);
+  free_qoco_csc_matrix(A);
+  qoco_free(A);
+}
 
-  for (QOCOInt i = 0; i < n; ++i) {
-    y[i] = x[i];
-  }
+void free_qoco_vectorf(QOCOVectorf* x)
+{
+  qoco_free(x->data);
+  qoco_free(x);
 }
 
-void copy_and_negate_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n)
+QOCOMatrix* create_transposed_matrix(const QOCOMatrix* A, QOCOInt* AtoAt)
 {
-  qoco_assert(x || n == 0);
-  qoco_assert(y || n == 0);
+  QOCOMatrix* At = qoco_malloc(sizeof(QOCOMatrix));
+  QOCOCscMatrix* Atcsc = create_transposed_csc_matrix(A->csc, AtoAt);
+  At->csc = Atcsc;
+  return At;
+}
 
-  for (QOCOInt i = 0; i < n; ++i) {
-    y[i] = -x[i];
-  }
+QOCOMatrix* regularize_P(QOCOInt num_diagP, QOCOMatrix* P, QOCOFloat reg,
+                         QOCOInt* nzadded_idx)
+{
+  regularize_P_csc(num_diagP, P->csc, reg, nzadded_idx);
+  qoco_free(P);
 }
 
-void copy_arrayi(const QOCOInt* x, QOCOInt* y, QOCOInt n)
+void unregularize(QOCOMatrix* M, QOCOFloat lambda)
 {
-  qoco_assert(x || n == 0);
-  qoco_assert(y || n == 0);
+  unregularize_csc(M->csc, lambda);
+}
 
-  for (QOCOInt i = 0; i < n; ++i) {
-    y[i] = x[i];
-  }
+QOCOMatrix* construct_identity(QOCOInt n, QOCOFloat lambda)
+{
+  QOCOMatrix* M = qoco_malloc(sizeof(QOCOMatrix));
+  M->csc = construct_identity_csc(n, lambda);
+}
+
+void scale_matrix(QOCOFloat a, QOCOMatrix* M)
+{
+  scale_arrayf(M->csc->x, M->csc->x, a, M->csc->nnz);
+}
+
+void row_col_scale(QOCOMatrix* M, QOCOVectorf* E, QOCOVectorf* D)
+{
+  row_col_scale_csc(M->csc, E->data, D->data);
+}
+
+void update_matrix(QOCOMatrix* M, QOCOFloat* Mnew)
+{
+  copy_arrayf(Mnew, M->csc->x, M->csc->nnz);
 }
 
 QOCOFloat qoco_dot(const QOCOFloat* u, const QOCOFloat* v, QOCOInt n)
@@ -96,24 +124,33 @@ QOCOFloat qoco_dot(const QOCOFloat* u, const QOCOFloat* v, QOCOInt n)
   return x;
 }
 
-QOCOInt max_arrayi(const QOCOInt* x, QOCOInt n)
+void ew_product(QOCOVectorf* x, QOCOVectorf* y, QOCOVectorf* z)
 {
-  qoco_assert(x || n == 0);
+  qoco_assert(x->len == y->len);
+  qoco_assert(x->len == z->len);
+  ew_product_arrayf(x->data, y->data, z->data, x->len);
+}
 
-  QOCOInt max = -QOCOInt_MAX;
-  for (QOCOInt i = 0; i < n; ++i) {
-    max = qoco_max(max, x[i]);
-  }
-  return max;
+void ew_product_vec_array(QOCOVectorf* x, QOCOFloat* y, QOCOVectorf* z)
+{
+  qoco_assert(x->len == y->len);
+  qoco_assert(x->len == z->len);
+  ew_product_arrayf(x->data, y, z->data, x->len);
 }
 
-void scale_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOFloat s, QOCOInt n)
+void scale_vectorf(QOCOFloat a, QOCOVectorf* u)
 {
-  qoco_assert(x || n == 0);
-  qoco_assert(y || n == 0);
+  scale_arrayf(u->data, u->data, a, u->len);
+}
 
-  for (QOCOInt i = 0; i < n; ++i) {
-    y[i] = s * x[i];
+void copy_vectorf(QOCOVectorf* src, QOCOFloat* dest, QOCOInt dest_idx,
+                  QOCOInt negate)
+{
+  if (negate) {
+    copy_arrayf(src->data, &dest[dest_idx], src->len);
+  }
+  else {
+    copy_and_negate_arrayf(src->data, &dest[dest_idx], src->len);
   }
 }
 
diff --git a/algebra/builtin/builtin_types.h b/algebra/builtin/builtin_types.h
new file mode 100644
index 0000000..eba90f4
--- /dev/null
+++ b/algebra/builtin/builtin_types.h
@@ -0,0 +1,36 @@
+/**
+ * @file builtin_types.h
+ * @author Govind M. Chari <govindchari1@gmail.com>
+ *
+ * @section LICENSE
+ *
+ * Copyright (c) 2025, Govind M. Chari
+ * This source code is licensed under the BSD 3-Clause License
+ *
+ * @section DESCRIPTION
+ *
+ * Defines the vector and matrices for builtin linear algebra.
+ */
+
+#ifndef BUILTIN_TYPES_H
+#define BUILTIN_TYPES_H
+
+#include "definitions.h"
+#include "qoco_linalg.h"
+#include "common_linalg.h"
+
+struct QOCOVectori_ {
+  QOCOInt* data;
+  QOCOInt  len;
+};
+
+struct QOCOVectorf_ {
+  QOCOFloat* data;
+  QOCOInt    len;
+};
+
+struct QOCOMatrix_ {
+  QOCOCscMatrix* csc;
+};
+
+#endif /* ifndef BUILTIN_TYPES_H */
diff --git a/devtools/run_tests.sh b/devtools/run_tests.sh
index 450c81e..08d4a2e 100755
--- a/devtools/run_tests.sh
+++ b/devtools/run_tests.sh
@@ -1 +1 @@
-export CXX=/usr/local/bin/clang++ && export CC=/usr/local/bin/clang && cd build && cmake -DQOCO_BUILD_TYPE:STR=Release -DENABLE_TESTING:BOOL=True -DBUILD_QOCO_BENCHMARK_RUNNER:BOOL=True .. && make && ctest --verbose && cd ..
\ No newline at end of file
+export CXX=/usr/bin/clang++ && export CC=/usr/bin/clang && cd build && cmake -DQOCO_BUILD_TYPE:STR=Release -DENABLE_TESTING:BOOL=True -DBUILD_QOCO_BENCHMARK_RUNNER:BOOL=True .. && make && ctest --verbose && cd ..
\ No newline at end of file
diff --git a/include/common_linalg.h b/include/common_linalg.h
index ac3180a..5cea309 100644
--- a/include/common_linalg.h
+++ b/include/common_linalg.h
@@ -17,6 +17,14 @@
 #include "definitions.h"
 #include "qoco_linalg.h"
 
+/**
+ * @brief Frees all the internal arrays and the pointer to the QOCOCscMatrix.
+ * Should only be used if QOCOCscMatrix and all internal arrays were malloc'ed.
+ *
+ * @param A Pointer to QOCOCscMatrix.
+ */
+void free_qoco_csc_matrix(QOCOCscMatrix* A);
+
 /**
  * @brief Allocates a new csc matrix that is lambda * I.
  *
@@ -24,7 +32,7 @@
  * @param lambda Scaling factor for identity.
  * @return Pointer to new constructed matrix.
  */
-QOCOCscMatrix* construct_identity(QOCOInt n, QOCOFloat lambda);
+QOCOCscMatrix* construct_identity_csc(QOCOInt n, QOCOFloat lambda);
 
 /**
  * @brief Counts the number of diagonal elements in upper triangular CSC matrix
@@ -46,7 +54,7 @@ QOCOInt count_diag(QOCOCscMatrix* M);
  * @param nzadded_idx Indices of elements of M->x that are added.
  * @return P + reg * I.
  */
-QOCOCscMatrix* regularize_P(QOCOInt num_diagP, QOCOCscMatrix* P, QOCOFloat reg,
+QOCOCscMatrix* regularize_P_csc(QOCOInt num_diagP, QOCOCscMatrix* P, QOCOFloat reg,
                             QOCOInt* nzadded_idx);
 
 /**
@@ -57,7 +65,7 @@ QOCOCscMatrix* regularize_P(QOCOInt num_diagP, QOCOCscMatrix* P, QOCOFloat reg,
  * @param M Matrix.
  * @param lambda Regularization.
  */
-void unregularize(QOCOCscMatrix* M, QOCOFloat lambda);
+void unregularize_csc(QOCOCscMatrix* M, QOCOFloat lambda);
 
 /**
  * @brief Computes the infinity norm of each column (or equivalently row) of a
@@ -83,17 +91,44 @@ void row_inf_norm(const QOCOCscMatrix* M, QOCOFloat* norm);
  * @param A Input matrix.
  * @param AtoAt Mapping from A to At.
  */
-QOCOCscMatrix* create_transposed_matrix(const QOCOCscMatrix* A, QOCOInt* AtoAt);
+QOCOCscMatrix* create_transposed_csc_matrix(const QOCOCscMatrix* A, QOCOInt* AtoAt);
 
 /**
  * @brief Scales the rows of M by E and columns of M by D.
  * M = diag(E) * M * diag(S)
  *
- * @param M An m by n sparse matrix.
+ * @param M An m by n QOCOCscMatrix matrix.
  * @param E Vector of length m.
  * @param D Vector of length m.
  */
-void row_col_scale(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D);
+void row_col_scale_csc(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D);
+
+/**
+ * @brief Copies array of QOCOFloats from x to array y.
+ *
+ * @param x Source array.
+ * @param y Destination array.
+ * @param n Length of arrays.
+ */
+void copy_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n);
+
+/**
+ * @brief Copies and negates array of QOCOFloats from x to array y.
+ *
+ * @param x Source array.
+ * @param y Destination array.
+ * @param n Length of arrays.
+ */
+void copy_and_negate_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n);
+
+/**
+ * @brief Copies array of QOCOInts from x to array y.
+ *
+ * @param x Source array.
+ * @param y Destination array.
+ * @param n Length of arrays.
+ */
+void copy_arrayi(const QOCOInt* x, QOCOInt* y, QOCOInt n);
 
 /**
  * @brief Computes elementwise product z = x .* y
@@ -103,7 +138,27 @@ void row_col_scale(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D);
  * @param z Output array.
  * @param n Length of arrays.
  */
-void ew_product(QOCOFloat* x, const QOCOFloat* y, QOCOFloat* z, QOCOInt n);
+void ew_product_arrayf(QOCOFloat* x, QOCOFloat* y, QOCOFloat* z, QOCOInt n);
+
+/**
+ * @brief Computes maximum element of array of QOCOInts.
+ *
+ * @param x Input array.
+ * @param n Length of array.
+ * @return Maximum element of x.
+ */
+QOCOInt max_arrayi(const QOCOInt* x, QOCOInt n);
+
+/**
+ * @brief Scales array x by s and stores result in y.
+ * y = s * x
+ *
+ * @param x Input array.
+ * @param y Output array.
+ * @param s Scaling factor.
+ * @param n Length of arrays.
+ */
+void scale_arrayf(QOCOFloat* x, QOCOFloat* y, QOCOFloat s, QOCOInt n);
 
 /**
  * @brief Inverts permutation vector p and stores inverse in pinv.
diff --git a/include/kkt.h b/include/kkt.h
index 7af88fc..454c18a 100644
--- a/include/kkt.h
+++ b/include/kkt.h
@@ -21,6 +21,12 @@
 #include "qoco_linalg.h"
 #include "structs.h"
 
+QOCOCscMatrix* create_kkt(QOCOCscMatrix* P, QOCOCscMatrix* A, QOCOCscMatrix* G,
+                          QOCOCscMatrix* At, QOCOCscMatrix* Gt, QOCOFloat static_reg, QOCOInt n,
+                          QOCOInt m, QOCOInt p, QOCOInt l, QOCOInt nsoc,
+                          QOCOInt* q, QOCOInt* PregtoKKT, QOCOInt* AttoKKT,
+                          QOCOInt* GttoKKT, QOCOInt* nt2kkt, QOCOInt* ntdiag2kkt);
+
 /**
  * @brief Allocate memory for KKT matrix.
  *
diff --git a/include/qoco_linalg.h b/include/qoco_linalg.h
index 8bbe7f2..39b8812 100644
--- a/include/qoco_linalg.h
+++ b/include/qoco_linalg.h
@@ -16,6 +16,10 @@
 #define QOCO_LINALG_H
 #include "definitions.h"
 
+typedef struct QOCOMatrix_ QOCOMatrix;
+typedef struct QOCOVectorf_ QOCOVectorf;
+typedef struct QOCOVectori_ QOCOVectorf_;
+
 /**
  * @brief Compressed sparse column format matrices.
  *
@@ -42,47 +46,102 @@ typedef struct {
 } QOCOCscMatrix;
 
 /**
- * @brief Allocates a new csc matrix and copies A to it.
+ * @brief Allocates a new QOCOMatrix and copies A to it.
  *
  * @param A Matrix to copy.
  * @return Pointer to new constructed matrix.
  */
-QOCOCscMatrix* new_qoco_csc_matrix(const QOCOCscMatrix* A);
+QOCOMatrix* new_qoco_matrix(const QOCOCscMatrix* A);
+
+/**
+ * @brief Frees QOCOMatrix.
+ *
+ * @param A Matrix to free.
+ */
+void free_qoco_matrix(QOCOMatrix* A);
+
+/**
+ * @brief Allocates a new QOCOVectorf and copies A to it.
+ *
+ * @param x vector to copy.
+ * @param n length of vector.
+ * @return Pointer to new constructed vector.
+ */
+QOCOVectorf* new_qoco_vectorf(const QOCOFloat* x, QOCOInt n);
 
 /**
- * @brief Frees all the internal arrays and the pointer to the QOCOCscMatrix.
- * Should only be used if QOCOCscMatrix and all internal arrays were malloc'ed.
+ * @brief Frees QOCOVectorf.
  *
- * @param A Pointer to QOCOCscMatrix.
+ * @param x Vector to free.
  */
-void free_qoco_csc_matrix(QOCOCscMatrix* A);
+void free_qoco_vectorf(QOCOVectorf* x);
 
 /**
- * @brief Copies array of QOCOFloats from x to array y.
+ * @brief Allocates a new QOCOMatrix and At to it.
  *
- * @param x Source array.
- * @param y Destination array.
- * @param n Length of arrays.
+ * @param A Matrix to transpose.
+ * @param AtoAt Mapping from indices of A to At.
+ * @return Pointer to transposed matrix.
+ */
+QOCOMatrix* create_transposed_matrix(const QOCOMatrix* A, QOCOInt* AtoAt);
+
+/**
+ * @brief Adds reg * I to a QOCOMatrix. Called on P prior to construction
+ * of KKT system in qoco_setup().
+ *
+ * @param num_diagP Number of new element added to diagonal.
+ * @param P Matrix to be regularized.
+ * @param reg Regularization factor.
+ * @param nzadded_idx Indices of elements of M->x that are added.
+ * @return P + reg * I.
  */
-void copy_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n);
+QOCOMatrix* regularize_P(QOCOInt num_diagP, QOCOMatrix* P, QOCOFloat reg,
+                         QOCOInt* nzadded_idx);
 
 /**
- * @brief Copies and negates array of QOCOFloats from x to array y.
+ * @brief Subtracts lambda * I to a QOCOMatrix. Called on P when updating
+ * matrix data in update_matrix_data(). This function does not allocate and must
+ * be called after regularize.
  *
- * @param x Source array.
- * @param y Destination array.
- * @param n Length of arrays.
+ * @param M Matrix.
+ * @param lambda Regularization.
  */
-void copy_and_negate_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n);
+void unregularize(QOCOMatrix* M, QOCOFloat lambda);
 
 /**
- * @brief Copies array of QOCOInts from x to array y.
+ * @brief Allocates a new QOCOMatrix that is lambda * I.
  *
- * @param x Source array.
- * @param y Destination array.
- * @param n Length of arrays.
+ * @param n Size of identity matrix.
+ * @param lambda Scaling factor for identity.
+ * @return Pointer to new constructed matrix.
  */
-void copy_arrayi(const QOCOInt* x, QOCOInt* y, QOCOInt n);
+QOCOMatrix* construct_identity(QOCOInt n, QOCOFloat lambda);
+
+/**
+ * @brief Computes M = a * M.
+ *
+ * @param a Scaling factor.
+ * @param M QOCOMatrix to scale.
+ */
+void scale_matrix(QOCOFloat a, QOCOMatrix* M);
+
+/**
+ * @brief Scales the rows of M by E and columns of M by D.
+ * M = diag(E) * M * diag(S)
+ *
+ * @param M An m by n QOCOMatrix matrix.
+ * @param E Vector of length m.
+ * @param D Vector of length m.
+ */
+void row_col_scale(QOCOMatrix* M, QOCOVectorf* E, QOCOVectorf* D);
+
+/**
+ * @brief Updates all the nonzero elements of M to Mnew.
+ *
+ * @param M Matrix to update.
+ * @param Mnew Vector of new data.
+ */
+void update_matrix(QOCOMatrix* M, QOCOFloat* Mnew);
 
 /**
  * @brief Computes dot product of u and v.
@@ -95,24 +154,27 @@ void copy_arrayi(const QOCOInt* x, QOCOInt* y, QOCOInt n);
 QOCOFloat qoco_dot(const QOCOFloat* u, const QOCOFloat* v, QOCOInt n);
 
 /**
- * @brief Computes maximum element of array of QOCOInts.
+ * @brief Computes elementwise product z = x .* y
  *
- * @param x Input array.
- * @param n Length of array.
- * @return Maximum element of x.
+ * @param x Input vector.
+ * @param y Input vector.
+ * @param z Output vector.
  */
-QOCOInt max_arrayi(const QOCOInt* x, QOCOInt n);
+void ew_product(QOCOVectorf* x, QOCOVectorf* y, QOCOVectorf* z);
 
 /**
- * @brief Scales array x by s and stores result in y.
- * y = s * x
+ * @brief Computes elementwise product z = x .* y
  *
- * @param x Input array.
- * @param y Output array.
- * @param s Scaling factor.
- * @param n Length of arrays.
+ * @param x Input vector.
+ * @param y Input array.
+ * @param z Output vector.
  */
-void scale_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOFloat s, QOCOInt n);
+void ew_product_vec_array(QOCOVectorf* x, QOCOFloat* y, QOCOVectorf* z);
+
+void scale_vectorf(QOCOFloat a, QOCOVectorf* u);
+
+void copy_vectorf(QOCOVectorf* src, QOCOFloat* dest, QOCOInt dest_idx,
+                  QOCOInt negate);
 
 /**
  * @brief Computes z = a * x + y.
diff --git a/include/structs.h b/include/structs.h
index fd0ff93..f1f204e 100644
--- a/include/structs.h
+++ b/include/structs.h
@@ -25,25 +25,25 @@
  */
 typedef struct {
   /** Quadratic cost term. */
-  QOCOCscMatrix* P; //
+  QOCOMatrix* P;
 
   /** Linear cost term. */
-  QOCOFloat* c;
+  QOCOVectorf* c;
 
   /** Affine equality constraint matrix. */
-  QOCOCscMatrix* A;
+  QOCOMatrix* A;
 
   /** Transpose of A (used in Ruiz for fast row norm calculations of A). */
-  QOCOCscMatrix* At;
+  QOCOMatrix* At;
 
   /** Affine equality constraint offset. */
-  QOCOFloat* b;
+  QOCOVectorf* b;
 
   /** Conic constraint matrix. */
-  QOCOCscMatrix* G;
+  QOCOMatrix* G;
 
   /** Transpose of G (used in Ruiz for fast row norm calculations of G). */
-  QOCOCscMatrix* Gt;
+  QOCOMatrix* Gt;
 
   /** Mapping from A to At. */
   QOCOInt* AtoAt;
@@ -52,7 +52,7 @@ typedef struct {
   QOCOInt* GtoGt;
 
   /** Conic constraint offset. */
-  QOCOFloat* h;
+  QOCOVectorf* h;
 
   /** Dimension of non-negative orthant in cone C. */
   QOCOInt l;
@@ -123,25 +123,25 @@ typedef struct {
   QOCOCscMatrix* K;
 
   /** Diagonal of scaling matrix. */
-  QOCOFloat* delta;
+  QOCOVectorf* delta;
 
   /** Diagonal of scaling matrix. */
-  QOCOFloat* Druiz;
+  QOCOVectorf* Druiz;
 
   /** Diagonal of scaling matrix. */
-  QOCOFloat* Eruiz;
+  QOCOVectorf* Eruiz;
 
   /** Diagonal of scaling matrix. */
-  QOCOFloat* Fruiz;
+  QOCOVectorf* Fruiz;
 
   /** Inverse of Druiz. */
-  QOCOFloat* Dinvruiz;
+  QOCOVectorf* Dinvruiz;
 
   /** Inverse of Eruiz. */
-  QOCOFloat* Einvruiz;
+  QOCOVectorf* Einvruiz;
 
   /** Inverse of Fruiz. */
-  QOCOFloat* Finvruiz;
+  QOCOVectorf* Finvruiz;
 
   /** Cost scaling factor. */
   QOCOFloat k;
diff --git a/src/common_linalg.c b/src/common_linalg.c
index d7634aa..c1f24e8 100644
--- a/src/common_linalg.c
+++ b/src/common_linalg.c
@@ -10,7 +10,15 @@
 
 #include "common_linalg.h"
 
-QOCOCscMatrix* construct_identity(QOCOInt n, QOCOFloat lambda)
+void free_qoco_csc_matrix(QOCOCscMatrix* A)
+{
+  free(A->x);
+  free(A->i);
+  free(A->p);
+  free(A);
+}
+
+QOCOCscMatrix* construct_identity_csc(QOCOInt n, QOCOFloat lambda)
 {
   QOCOCscMatrix* M = qoco_malloc(sizeof(QOCOCscMatrix));
   QOCOFloat* x = qoco_malloc(n * sizeof(QOCOFloat));
@@ -49,8 +57,8 @@ QOCOInt count_diag(QOCOCscMatrix* M)
   return count;
 }
 
-QOCOCscMatrix* regularize_P(QOCOInt num_diagP, QOCOCscMatrix* P, QOCOFloat reg,
-                            QOCOInt* nzadded_idx)
+QOCOCscMatrix* regularize_P_csc(QOCOInt num_diagP, QOCOCscMatrix* P,
+                                QOCOFloat reg, QOCOInt* nzadded_idx)
 {
   QOCOInt n = P->n;
 
@@ -109,7 +117,7 @@ QOCOCscMatrix* regularize_P(QOCOInt num_diagP, QOCOCscMatrix* P, QOCOFloat reg,
   return Preg;
 }
 
-void unregularize(QOCOCscMatrix* M, QOCOFloat lambda)
+void unregularize_csc(QOCOCscMatrix* M, QOCOFloat lambda)
 {
   // Iterate over each column.
   for (QOCOInt col = 0; col < M->n; col++) {
@@ -165,7 +173,8 @@ void row_inf_norm(const QOCOCscMatrix* M, QOCOFloat* norm)
   }
 }
 
-QOCOCscMatrix* create_transposed_matrix(const QOCOCscMatrix* A, QOCOInt* AtoAt)
+QOCOCscMatrix* create_transposed_csc_matrix(const QOCOCscMatrix* A,
+                                            QOCOInt* AtoAt)
 {
   QOCOCscMatrix* B = qoco_malloc(sizeof(QOCOCscMatrix));
   B->m = A->n;
@@ -198,7 +207,9 @@ QOCOCscMatrix* create_transposed_matrix(const QOCOCscMatrix* A, QOCOInt* AtoAt)
       int dest_pos = B->p[row] + temp[row];
       B->i[dest_pos] = j;       // Column index becomes row index
       B->x[dest_pos] = A->x[i]; // Value remains the same
-      AtoAt[i] = dest_pos;
+      if (AtoAt) {
+        AtoAt[i] = dest_pos;
+      }
       temp[row]++;
     }
   }
@@ -210,7 +221,7 @@ QOCOCscMatrix* create_transposed_matrix(const QOCOCscMatrix* A, QOCOInt* AtoAt)
   return B;
 }
 
-void row_col_scale(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D)
+void row_col_scale_csc(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D)
 {
   for (QOCOInt j = 0; j < M->n; ++j) {
     for (QOCOInt i = M->p[j]; i < M->p[j + 1]; ++i) {
@@ -219,7 +230,58 @@ void row_col_scale(const QOCOCscMatrix* M, QOCOFloat* E, QOCOFloat* D)
   }
 }
 
-void ew_product(QOCOFloat* x, const QOCOFloat* y, QOCOFloat* z, QOCOInt n)
+void copy_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n)
+{
+  qoco_assert(x || n == 0);
+  qoco_assert(y || n == 0);
+
+  for (QOCOInt i = 0; i < n; ++i) {
+    y[i] = x[i];
+  }
+}
+
+void copy_and_negate_arrayf(const QOCOFloat* x, QOCOFloat* y, QOCOInt n)
+{
+  qoco_assert(x || n == 0);
+  qoco_assert(y || n == 0);
+
+  for (QOCOInt i = 0; i < n; ++i) {
+    y[i] = -x[i];
+  }
+}
+
+void copy_arrayi(const QOCOInt* x, QOCOInt* y, QOCOInt n)
+{
+  qoco_assert(x || n == 0);
+  qoco_assert(y || n == 0);
+
+  for (QOCOInt i = 0; i < n; ++i) {
+    y[i] = x[i];
+  }
+}
+
+QOCOInt max_arrayi(const QOCOInt* x, QOCOInt n)
+{
+  qoco_assert(x || n == 0);
+
+  QOCOInt max = -QOCOInt_MAX;
+  for (QOCOInt i = 0; i < n; ++i) {
+    max = qoco_max(max, x[i]);
+  }
+  return max;
+}
+
+void scale_arrayf(QOCOFloat* x, QOCOFloat* y, QOCOFloat s, QOCOInt n)
+{
+  qoco_assert(x || n == 0);
+  qoco_assert(y || n == 0);
+
+  for (QOCOInt i = 0; i < n; ++i) {
+    y[i] = s * x[i];
+  }
+}
+
+void ew_product_arrayf(QOCOFloat* x, QOCOFloat* y, QOCOFloat* z, QOCOInt n)
 {
   for (QOCOInt i = 0; i < n; ++i) {
     z[i] = x[i] * y[i];
diff --git a/src/kkt.c b/src/kkt.c
index 4a5f124..a606a61 100644
--- a/src/kkt.c
+++ b/src/kkt.c
@@ -11,92 +11,88 @@
 #include "kkt.h"
 #include "qoco_utils.h"
 
-void allocate_kkt(QOCOWorkspace* work)
+QOCOCscMatrix* create_kkt(QOCOCscMatrix* P, QOCOCscMatrix* A, QOCOCscMatrix* G,
+                          QOCOCscMatrix* At, QOCOCscMatrix* Gt, QOCOFloat static_reg, QOCOInt n,
+                          QOCOInt m, QOCOInt p, QOCOInt l, QOCOInt nsoc,
+                          QOCOInt* q, QOCOInt* PregtoKKT, QOCOInt* AttoKKT,
+                          QOCOInt* GttoKKT, QOCOInt* nt2kkt, QOCOInt* ntdiag2kkt)
 {
-  work->kkt->K = qoco_malloc(sizeof(QOCOCscMatrix));
+  QOCOCscMatrix* KKT = qoco_malloc(sizeof(QOCOCscMatrix));
 
   // Number of nonzeros in second-order cone part of NT scaling.
   QOCOInt Wsoc_nnz = 0;
-  for (QOCOInt i = 0; i < work->data->nsoc; ++i) {
-    Wsoc_nnz += work->data->q[i] * work->data->q[i] - work->data->q[i];
+  for (QOCOInt i = 0; i < nsoc; ++i) {
+    Wsoc_nnz += q[i] * q[i] - q[i];
   }
   Wsoc_nnz /= 2;
+  QOCOInt Wnnz = m + Wsoc_nnz;
 
-  work->Wnnz = work->data->m + Wsoc_nnz;
-  work->kkt->Wnnz = work->data->m + Wsoc_nnz;
-  work->kkt->K->m = work->data->n + work->data->m + work->data->p;
-  work->kkt->K->n = work->data->n + work->data->m + work->data->p;
-  work->kkt->K->nnz = work->data->P->nnz + work->data->A->nnz +
-                      work->data->G->nnz + work->Wnnz + work->data->p;
+  KKT->m = n + m + p;
+  KKT->n = n + m + p;
+  KKT->nnz = P->nnz + A->nnz + G->nnz + Wnnz + p;
 
-  work->kkt->K->x = qoco_calloc(work->kkt->K->nnz, sizeof(QOCOFloat));
-  work->kkt->K->i = qoco_calloc(work->kkt->K->nnz, sizeof(QOCOInt));
-  work->kkt->K->p = qoco_calloc((work->kkt->K->n + 1), sizeof(QOCOInt));
-}
+  KKT->x = qoco_calloc(KKT->nnz, sizeof(QOCOFloat));
+  KKT->i = qoco_calloc(KKT->nnz, sizeof(QOCOInt));
+  KKT->p = qoco_calloc((KKT->n + 1), sizeof(QOCOInt));
 
-void construct_kkt(QOCOSolver* solver)
-{
-  QOCOWorkspace* work = solver->work;
   QOCOInt nz = 0;
   QOCOInt col = 0;
   // Add P block
-  for (QOCOInt k = 0; k < work->data->P->nnz; ++k) {
-    work->kkt->PregtoKKT[k] = nz;
-    work->kkt->K->x[nz] = work->data->P->x[k];
-    work->kkt->K->i[nz] = work->data->P->i[k];
+  for (QOCOInt k = 0; k < P->nnz; ++k) {
+    PregtoKKT[k] = nz;
+    KKT->x[nz] = P->x[k];
+    KKT->i[nz] = P->i[k];
     nz += 1;
   }
-  for (QOCOInt k = 0; k < work->data->P->n + 1; ++k) {
-    work->kkt->K->p[col] = work->data->P->p[k];
+  for (QOCOInt k = 0; k < P->n + 1; ++k) {
+    KKT->p[col] = P->p[k];
     col += 1;
   }
 
   // Add A^T block
-  for (QOCOInt Atcol = 0; Atcol < work->data->At->n; ++Atcol) {
+  for (QOCOInt Atcol = 0; Atcol < At->n; ++Atcol) {
     QOCOInt nzadded = 0;
-    for (QOCOInt k = work->data->At->p[Atcol]; k < work->data->At->p[Atcol + 1];
-         ++k) {
+    for (QOCOInt k = At->p[Atcol]; k < At->p[Atcol + 1]; ++k) {
       // If the nonzero is in row i of A then add
-      work->kkt->AttoKKT[k] = nz;
-      work->kkt->K->x[nz] = work->data->At->x[k];
-      work->kkt->K->i[nz] = work->data->At->i[k];
+      AttoKKT[k] = nz;
+      KKT->x[nz] = At->x[k];
+      KKT->i[nz] = At->i[k];
       nz += 1;
       nzadded += 1;
     }
 
     // Add -e * Id regularization.
-    work->kkt->K->x[nz] = -solver->settings->kkt_static_reg;
-    work->kkt->K->i[nz] = work->data->n + Atcol;
+    KKT->x[nz] = -static_reg;
+    KKT->i[nz] = n + Atcol;
     nz += 1;
     nzadded += 1;
-    work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded;
+    KKT->p[col] = KKT->p[col - 1] + nzadded;
     col += 1;
   }
 
   // Add non-negative orthant part of G^T.
   QOCOInt nz_nt = 0;
   QOCOInt diag = 0;
-  for (QOCOInt Gtcol = 0; Gtcol < work->data->l; ++Gtcol) {
+  for (QOCOInt Gtcol = 0; Gtcol < l; ++Gtcol) {
 
     // Counter for number of nonzeros from G added to this column of KKT matrix
     QOCOInt nzadded = 0;
-    for (QOCOInt k = work->data->Gt->p[Gtcol]; k < work->data->Gt->p[Gtcol + 1];
-         ++k) {
-      work->kkt->GttoKKT[k] = nz;
-      work->kkt->K->x[nz] = work->data->Gt->x[k];
-      work->kkt->K->i[nz] = work->data->Gt->i[k];
+    for (QOCOInt k = Gt->p[Gtcol]; k < Gt->p[Gtcol + 1]; ++k) {
+      GttoKKT[k] = nz;
+      KKT->x[nz] = Gt->x[k];
+      KKT->i[nz] = Gt->i[k];
       nz += 1;
       nzadded += 1;
     }
 
     // Add -Id to NT block.
-    work->kkt->K->x[nz] = -1.0;
-    work->kkt->K->i[nz] = work->data->n + work->data->p + Gtcol;
-    work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded + 1;
+    KKT->x[nz] = -1.0;
+    KKT->i[nz] = n + p + Gtcol;
+    KKT->p[col] = KKT->p[col - 1] + nzadded + 1;
 
     // Mapping from NT matrix entries to KKT matrix entries.
-    work->kkt->nt2kkt[nz_nt] = nz;
-    work->kkt->ntdiag2kkt[diag] = nz;
+    nt2kkt[nz_nt] = nz;
+    ntdiag2kkt[diag] = nz;
     diag++;
     nz_nt += 1;
 
@@ -105,51 +101,190 @@ void construct_kkt(QOCOSolver* solver)
   }
 
   // Add second-order cone parts of G^T.
-  QOCOInt idx = work->data->l;
-  for (QOCOInt c = 0; c < work->data->nsoc; ++c) {
-    for (QOCOInt Gtcol = idx; Gtcol < idx + work->data->q[c]; ++Gtcol) {
+  QOCOInt idx = l;
+  for (QOCOInt c = 0; c < nsoc; ++c) {
+    for (QOCOInt Gtcol = idx; Gtcol < idx + q[c]; ++Gtcol) {
       // Loop over columns of G
 
       // Counter for number of nonzeros from G added to this column of KKT
       // matrix
       QOCOInt nzadded = 0;
-      for (QOCOInt k = work->data->Gt->p[Gtcol];
-           k < work->data->Gt->p[Gtcol + 1]; ++k) {
-        work->kkt->GttoKKT[k] = nz;
-        work->kkt->K->x[nz] = work->data->Gt->x[k];
-        work->kkt->K->i[nz] = work->data->Gt->i[k];
+      for (QOCOInt k = Gt->p[Gtcol]; k < Gt->p[Gtcol + 1]; ++k) {
+        GttoKKT[k] = nz;
+        KKT->x[nz] = Gt->x[k];
+        KKT->i[nz] = Gt->i[k];
         nz += 1;
         nzadded += 1;
       }
 
       // Add NT block.
-      for (QOCOInt i = idx; i < idx + work->data->q[c]; i++) {
+      for (QOCOInt i = idx; i < idx + q[c]; i++) {
         // Only add upper triangular part.
-        if (i + work->data->n + work->data->p <= col - 1) {
+        if (i + n + p <= col - 1) {
           // Add -1 if element is on main diagonal and 0 otherwise.
-          if (i + work->data->n + work->data->p == col - 1) {
-            work->kkt->K->x[nz] = -1.0;
-            work->kkt->ntdiag2kkt[diag] = nz;
+          if (i + n + p == col - 1) {
+            KKT->x[nz] = -1.0;
+            ntdiag2kkt[diag] = nz;
             diag++;
           }
           else {
-            work->kkt->K->x[nz] = 0.0;
+            KKT->x[nz] = 0.0;
           }
-          work->kkt->K->i[nz] = work->data->n + work->data->p + i;
-          work->kkt->nt2kkt[nz_nt] = nz;
+          KKT->i[nz] = n + p + i;
+          nt2kkt[nz_nt] = nz;
           nz_nt += 1;
           nz += 1;
           nzadded += 1;
         }
       }
-      work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded;
+      KKT->p[col] = KKT->p[col - 1] + nzadded;
       // Mapping from NT matrix entries to KKT matrix entries.
       col += 1;
     }
-    idx += work->data->q[c];
+    idx += q[c];
   }
+  return KKT;
 }
 
+// void allocate_kkt(QOCOWorkspace* work)
+// {
+//   work->kkt->K = qoco_malloc(sizeof(QOCOCscMatrix));
+
+//   // Number of nonzeros in second-order cone part of NT scaling.
+//   QOCOInt Wsoc_nnz = 0;
+//   for (QOCOInt i = 0; i < work->data->nsoc; ++i) {
+//     Wsoc_nnz += work->data->q[i] * work->data->q[i] - work->data->q[i];
+//   }
+//   Wsoc_nnz /= 2;
+
+//   work->Wnnz = work->data->m + Wsoc_nnz;
+//   work->kkt->Wnnz = work->data->m + Wsoc_nnz;
+//   work->kkt->K->m = work->data->n + work->data->m + work->data->p;
+//   work->kkt->K->n = work->data->n + work->data->m + work->data->p;
+//   work->kkt->K->nnz = work->data->P->nnz + work->data->A->nnz +
+//                       work->data->G->nnz + work->Wnnz + work->data->p;
+
+//   work->kkt->K->x = qoco_calloc(work->kkt->K->nnz, sizeof(QOCOFloat));
+//   work->kkt->K->i = qoco_calloc(work->kkt->K->nnz, sizeof(QOCOInt));
+//   work->kkt->K->p = qoco_calloc((work->kkt->K->n + 1), sizeof(QOCOInt));
+// }
+
+// void construct_kkt(QOCOSolver* solver)
+// {
+//   QOCOWorkspace* work = solver->work;
+//   QOCOInt nz = 0;
+//   QOCOInt col = 0;
+//   // Add P block
+//   for (QOCOInt k = 0; k < work->data->P->nnz; ++k) {
+//     work->kkt->PregtoKKT[k] = nz;
+//     work->kkt->K->x[nz] = work->data->P->x[k];
+//     work->kkt->K->i[nz] = work->data->P->i[k];
+//     nz += 1;
+//   }
+//   for (QOCOInt k = 0; k < work->data->P->n + 1; ++k) {
+//     work->kkt->K->p[col] = work->data->P->p[k];
+//     col += 1;
+//   }
+
+//   // Add A^T block
+//   for (QOCOInt Atcol = 0; Atcol < work->data->At->n; ++Atcol) {
+//     QOCOInt nzadded = 0;
+//     for (QOCOInt k = work->data->At->p[Atcol]; k < work->data->At->p[Atcol + 1];
+//          ++k) {
+//       // If the nonzero is in row i of A then add
+//       work->kkt->AttoKKT[k] = nz;
+//       work->kkt->K->x[nz] = work->data->At->x[k];
+//       work->kkt->K->i[nz] = work->data->At->i[k];
+//       nz += 1;
+//       nzadded += 1;
+//     }
+
+//     // Add -e * Id regularization.
+//     work->kkt->K->x[nz] = -solver->settings->kkt_static_reg;
+//     work->kkt->K->i[nz] = work->data->n + Atcol;
+//     nz += 1;
+//     nzadded += 1;
+//     work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded;
+//     col += 1;
+//   }
+
+//   // Add non-negative orthant part of G^T.
+//   QOCOInt nz_nt = 0;
+//   QOCOInt diag = 0;
+//   for (QOCOInt Gtcol = 0; Gtcol < work->data->l; ++Gtcol) {
+
+//     // Counter for number of nonzeros from G added to this column of KKT matrix
+//     QOCOInt nzadded = 0;
+//     for (QOCOInt k = work->data->Gt->p[Gtcol]; k < work->data->Gt->p[Gtcol + 1];
+//          ++k) {
+//       work->kkt->GttoKKT[k] = nz;
+//       work->kkt->K->x[nz] = work->data->Gt->x[k];
+//       work->kkt->K->i[nz] = work->data->Gt->i[k];
+//       nz += 1;
+//       nzadded += 1;
+//     }
+
+//     // Add -Id to NT block.
+//     work->kkt->K->x[nz] = -1.0;
+//     work->kkt->K->i[nz] = work->data->n + work->data->p + Gtcol;
+//     work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded + 1;
+
+//     // Mapping from NT matrix entries to KKT matrix entries.
+//     work->kkt->nt2kkt[nz_nt] = nz;
+//     work->kkt->ntdiag2kkt[diag] = nz;
+//     diag++;
+//     nz_nt += 1;
+
+//     nz += 1;
+//     col += 1;
+//   }
+
+//   // Add second-order cone parts of G^T.
+//   QOCOInt idx = work->data->l;
+//   for (QOCOInt c = 0; c < work->data->nsoc; ++c) {
+//     for (QOCOInt Gtcol = idx; Gtcol < idx + work->data->q[c]; ++Gtcol) {
+//       // Loop over columns of G
+
+//       // Counter for number of nonzeros from G added to this column of KKT
+//       // matrix
+//       QOCOInt nzadded = 0;
+//       for (QOCOInt k = work->data->Gt->p[Gtcol];
+//            k < work->data->Gt->p[Gtcol + 1]; ++k) {
+//         work->kkt->GttoKKT[k] = nz;
+//         work->kkt->K->x[nz] = work->data->Gt->x[k];
+//         work->kkt->K->i[nz] = work->data->Gt->i[k];
+//         nz += 1;
+//         nzadded += 1;
+//       }
+
+//       // Add NT block.
+//       for (QOCOInt i = idx; i < idx + work->data->q[c]; i++) {
+//         // Only add upper triangular part.
+//         if (i + work->data->n + work->data->p <= col - 1) {
+//           // Add -1 if element is on main diagonal and 0 otherwise.
+//           if (i + work->data->n + work->data->p == col - 1) {
+//             work->kkt->K->x[nz] = -1.0;
+//             work->kkt->ntdiag2kkt[diag] = nz;
+//             diag++;
+//           }
+//           else {
+//             work->kkt->K->x[nz] = 0.0;
+//           }
+//           work->kkt->K->i[nz] = work->data->n + work->data->p + i;
+//           work->kkt->nt2kkt[nz_nt] = nz;
+//           nz_nt += 1;
+//           nz += 1;
+//           nzadded += 1;
+//         }
+//       }
+//       work->kkt->K->p[col] = work->kkt->K->p[col - 1] + nzadded;
+//       // Mapping from NT matrix entries to KKT matrix entries.
+//       col += 1;
+//     }
+//     idx += work->data->q[c];
+//   }
+// }
+
 void initialize_ipm(QOCOSolver* solver)
 {
 
@@ -184,18 +319,9 @@ void initialize_ipm(QOCOSolver* solver)
   solver->work->a = 1.0;
 
   // Construct rhs of KKT system.
-  idx = 0;
-  for (idx = 0; idx < solver->work->data->n; ++idx) {
-    solver->work->kkt->rhs[idx] = -solver->work->data->c[idx];
-  }
-  for (QOCOInt i = 0; i < solver->work->data->p; ++i) {
-    solver->work->kkt->rhs[idx] = solver->work->data->b[i];
-    idx += 1;
-  }
-  for (QOCOInt i = 0; i < solver->work->data->m; ++i) {
-    solver->work->kkt->rhs[idx] = solver->work->data->h[i];
-    idx += 1;
-  }
+  copy_vectorf(solver->work->data->c, solver->work->kkt->rhs, 0, 1);
+  copy_vectorf(solver->work->data->b, solver->work->kkt->rhs, solver->work->data->n, 0);
+  copy_vectorf(solver->work->data->h, solver->work->kkt->rhs, solver->work->data->n + solver->work->data->n, 0);
 
   // Factor KKT matrix.
   solver->linsys->linsys_factor(solver->linsys_data, solver->work->data->n,
@@ -276,6 +402,7 @@ void compute_kkt_residual(QOCOSolver* solver)
         work->kkt->kktres[idx] +
         (work->data->c[idx] - solver->settings->kkt_static_reg * work->x[idx]);
   }
+  qoco_ax
 
   // Add -b.
   for (QOCOInt i = 0; i < work->data->p; ++i) {
diff --git a/src/qoco_api.c b/src/qoco_api.c
index 0464485..b3ca8f5 100644
--- a/src/qoco_api.c
+++ b/src/qoco_api.c
@@ -51,16 +51,15 @@ QOCOInt qoco_setup(QOCOSolver* solver, QOCOInt n, QOCOInt m, QOCOInt p,
   solver->work->data->m = m;
   solver->work->data->n = n;
   solver->work->data->p = p;
-  solver->work->data->A = new_qoco_csc_matrix(A);
-  solver->work->data->G = new_qoco_csc_matrix(G);
-  solver->work->data->c = qoco_malloc(n * sizeof(QOCOFloat));
-  solver->work->data->b = qoco_malloc(p * sizeof(QOCOFloat));
-  solver->work->data->h = qoco_malloc(m * sizeof(QOCOFloat));
+  solver->work->data->A = new_qoco_matrix(A);
+  solver->work->data->G = new_qoco_matrix(G);
+
   solver->work->data->q = qoco_malloc(nsoc * sizeof(QOCOInt));
 
-  copy_arrayf(c, solver->work->data->c, n);
-  copy_arrayf(b, solver->work->data->b, p);
-  copy_arrayf(h, solver->work->data->h, m);
+  solver->work->data->c = new_qoco_vectorf(c, n);
+  solver->work->data->b = new_qoco_vectorf(b, p);
+  solver->work->data->h = new_qoco_vectorf(h, m);
+
   copy_arrayi(q, solver->work->data->q, nsoc);
 
   solver->work->data->l = l;
@@ -68,7 +67,7 @@ QOCOInt qoco_setup(QOCOSolver* solver, QOCOInt n, QOCOInt m, QOCOInt p,
 
   // Copy P.
   if (P) {
-    solver->work->data->P = new_qoco_csc_matrix(P);
+    solver->work->data->P = new_qoco_matrix(P);
   }
   else {
     solver->work->data->P = NULL;
@@ -83,10 +82,8 @@ QOCOInt qoco_setup(QOCOSolver* solver, QOCOInt n, QOCOInt m, QOCOInt p,
   solver->work->kkt->Dinvruiz = qoco_malloc(n * sizeof(QOCOFloat));
   solver->work->kkt->Einvruiz = qoco_malloc(p * sizeof(QOCOFloat));
   solver->work->kkt->Finvruiz = qoco_malloc(m * sizeof(QOCOFloat));
-  solver->work->data->AtoAt =
-      qoco_malloc(solver->work->data->A->nnz * sizeof(QOCOInt));
-  solver->work->data->GtoGt =
-      qoco_malloc(solver->work->data->G->nnz * sizeof(QOCOInt));
+  solver->work->data->AtoAt = qoco_malloc(A->nnz * sizeof(QOCOInt));
+  solver->work->data->GtoGt = qoco_malloc(G->nnz * sizeof(QOCOInt));
 
   solver->work->data->At = create_transposed_matrix(solver->work->data->A,
                                                     solver->work->data->AtoAt);
@@ -96,36 +93,60 @@ QOCOInt qoco_setup(QOCOSolver* solver, QOCOInt n, QOCOInt m, QOCOInt p,
 
   // Regularize P.
   solver->work->kkt->Pnzadded_idx = qoco_calloc(n, sizeof(QOCOInt));
+
+  QOCOCscMatrix* Preg_csc;
   if (P) {
     QOCOInt num_diagP = count_diag(P);
     solver->work->kkt->Pnum_nzadded = n - num_diagP;
-    QOCOCscMatrix* Preg = regularize_P(num_diagP, solver->work->data->P,
-                                       solver->settings->kkt_static_reg,
-                                       solver->work->kkt->Pnzadded_idx);
+    QOCOMatrix* Preg = regularize_P(num_diagP, solver->work->data->P,
+                                    solver->settings->kkt_static_reg,
+                                    solver->work->kkt->Pnzadded_idx);
+    Preg_csc =
+        regularize_P_csc(num_diagP, P, solver->settings->kkt_static_reg, NULL);
     solver->work->data->P = Preg;
   }
   else {
     solver->work->data->P =
         construct_identity(n, solver->settings->kkt_static_reg);
     solver->work->kkt->Pnum_nzadded = n;
+    Preg_csc = construct_identity_csc(n, solver->settings->kkt_static_reg);
   }
 
   // Allocate KKT struct.
-  allocate_kkt(solver->work);
+  // allocate_kkt(solver->work);
   solver->work->kkt->nt2kkt = qoco_calloc(solver->work->Wnnz, sizeof(QOCOInt));
   solver->work->kkt->ntdiag2kkt = qoco_calloc(m, sizeof(QOCOInt));
   solver->work->kkt->PregtoKKT =
-      qoco_calloc(solver->work->data->P->nnz, sizeof(QOCOInt));
-  solver->work->kkt->AttoKKT =
-      qoco_calloc(solver->work->data->A->nnz, sizeof(QOCOInt));
-  solver->work->kkt->GttoKKT =
-      qoco_calloc(solver->work->data->G->nnz, sizeof(QOCOInt));
+      qoco_calloc(P->nnz + solver->work->kkt->Pnum_nzadded, sizeof(QOCOInt));
+  solver->work->kkt->AttoKKT = qoco_calloc(A->nnz, sizeof(QOCOInt));
+  solver->work->kkt->GttoKKT = qoco_calloc(G->nnz, sizeof(QOCOInt));
   solver->work->kkt->rhs = qoco_malloc((n + m + p) * sizeof(QOCOFloat));
   solver->work->kkt->kktres = qoco_malloc((n + m + p) * sizeof(QOCOFloat));
   solver->work->kkt->xyz = qoco_malloc((n + m + p) * sizeof(QOCOFloat));
   solver->work->kkt->xyzbuff1 = qoco_malloc((n + m + p) * sizeof(QOCOFloat));
   solver->work->kkt->xyzbuff2 = qoco_malloc((n + m + p) * sizeof(QOCOFloat));
-  construct_kkt(solver);
+  // construct_kkt(solver);
+
+  // Number of nonzeros in second-order cone part of NT scaling.
+  QOCOInt Wsoc_nnz = 0;
+  for (QOCOInt i = 0; i < nsoc; ++i) {
+    Wsoc_nnz += q[i] * q[i] - q[i];
+  }
+  Wsoc_nnz /= 2;
+  QOCOInt Wnnz = m + Wsoc_nnz;
+  solver->work->Wnnz = Wnnz;
+  solver->work->kkt->Wnnz = Wnnz;
+
+  QOCOCscMatrix* At_csc = create_transposed_csc_matrix(A, NULL);
+  QOCOCscMatrix* Gt_csc = create_transposed_csc_matrix(A, NULL);
+  solver->work->kkt->K = create_kkt(
+      Preg_csc, A, G, At_csc, Gt_csc, solver->settings->kkt_static_reg, n, m, p,
+      l, nsoc, q, solver->work->kkt->PregtoKKT, solver->work->kkt->AttoKKT,
+      solver->work->kkt->GttoKKT, solver->work->kkt->nt2kkt,
+      solver->work->kkt->ntdiag2kkt);
+  free_qoco_csc_matrix(At_csc);
+  free_qoco_csc_matrix(Gt_csc);
+  free_qoco_csc_matrix(Preg_csc);
 
   solver->linsys = &backend;
 
@@ -243,23 +264,18 @@ void update_vector_data(QOCOSolver* solver, QOCOFloat* cnew, QOCOFloat* bnew,
 
   // Update cost vector.
   if (cnew) {
-    for (QOCOInt i = 0; i < data->n; ++i) {
-      data->c[i] = solver->work->kkt->k * solver->work->kkt->Druiz[i] * cnew[i];
-    }
+    ew_product_vec_array(solver->work->kkt->Druiz, cnew, data->c);
+    scale_vectorf(solver->work->kkt->k, data->c);
   }
 
   // Update equality constraint vector.
   if (bnew) {
-    for (QOCOInt i = 0; i < data->p; ++i) {
-      data->b[i] = solver->work->kkt->Eruiz[i] * bnew[i];
-    }
+    ew_product_vec_array(solver->work->kkt->Eruiz, bnew, data->b);
   }
 
   // Update conic constraint vector.
   if (hnew) {
-    for (QOCOInt i = 0; i < data->m; ++i) {
-      data->h[i] = solver->work->kkt->Fruiz[i] * hnew[i];
-    }
+    ew_product_vec_array(solver->work->kkt->Fruiz, hnew, data->h);
   }
 }
 
@@ -274,12 +290,12 @@ void update_matrix_data(QOCOSolver* solver, QOCOFloat* Pxnew, QOCOFloat* Axnew,
   unregularize(data->P, solver->settings->kkt_static_reg);
 
   // Unequilibrate P.
-  scale_arrayf(data->P->x, data->P->x, kkt->kinv, data->P->nnz);
+  scale_matrix(kkt->kinv, data->P);
   row_col_scale(data->P, kkt->Dinvruiz, kkt->Dinvruiz);
 
   // Unequilibrate c.
-  scale_arrayf(data->c, data->c, kkt->kinv, data->n);
-  ew_product(data->c, kkt->Dinvruiz, data->c, data->n);
+  scale_vectorf(kkt->kinv, data->c);
+  ew_product(kkt->Dinvruiz, data->c, data->c);
 
   // Unequilibrate A.
   row_col_scale(data->A, kkt->Einvruiz, kkt->Dinvruiz);
@@ -288,38 +304,34 @@ void update_matrix_data(QOCOSolver* solver, QOCOFloat* Pxnew, QOCOFloat* Axnew,
   row_col_scale(data->G, kkt->Finvruiz, kkt->Dinvruiz);
 
   // Unequilibrate b.
-  ew_product(data->b, kkt->Einvruiz, data->b, data->p);
+  ew_product(kkt->Einvruiz, data->b, data->b);
 
   // Unequilibrate h.
-  ew_product(data->h, kkt->Finvruiz, data->h, data->m);
+  ew_product(kkt->Finvruiz, data->h, data->h);
 
   // Update P and avoid nonzeros that were added for regularization.
-  if (Pxnew) {
-    QOCOInt avoid =
-        kkt->Pnum_nzadded > 0 ? kkt->Pnzadded_idx[0] : data->P->nnz + 1;
-    QOCOInt offset = 0;
-    for (QOCOInt i = 0; i < data->P->nnz - kkt->Pnum_nzadded; ++i) {
-      if (i == avoid) {
-        offset++;
-        avoid = offset > kkt->Pnum_nzadded ? kkt->Pnzadded_idx[offset]
-                                           : data->P->nnz + 1;
-      }
-      data->P->x[i + offset] = Pxnew[i];
-    }
-  }
+  // if (Pxnew) {
+  //   QOCOInt avoid =
+  //       kkt->Pnum_nzadded > 0 ? kkt->Pnzadded_idx[0] : data->P->nnz + 1;
+  //   QOCOInt offset = 0;
+  //   for (QOCOInt i = 0; i < data->P->nnz - kkt->Pnum_nzadded; ++i) {
+  //     if (i == avoid) {
+  //       offset++;
+  //       avoid = offset > kkt->Pnum_nzadded ? kkt->Pnzadded_idx[offset]
+  //                                          : data->P->nnz + 1;
+  //     }
+  //     data->P->x[i + offset] = Pxnew[i];
+  //   }
+  // }
 
   // Update A.
   if (Axnew) {
-    for (QOCOInt i = 0; i < data->A->nnz; ++i) {
-      data->A->x[i] = Axnew[i];
-    }
+    update_matrix(data->A, Axnew);
   }
 
   // Update G.
   if (Gxnew) {
-    for (QOCOInt i = 0; i < data->G->nnz; ++i) {
-      data->G->x[i] = Gxnew[i];
-    }
+    update_matrix(data->G, Gxnew);
   }
 
   // Equilibrate new matrix data.
@@ -330,24 +342,24 @@ void update_matrix_data(QOCOSolver* solver, QOCOFloat* Pxnew, QOCOFloat* Axnew,
 
   solver->linsys->linsys_update_data(solver->linsys_data, solver->work->data);
 
-  // Update P in KKT matrix.
-  for (QOCOInt i = 0; i < data->P->nnz; ++i) {
-    solver->work->kkt->K->x[solver->work->kkt->PregtoKKT[i]] = data->P->x[i];
-  }
-
-  // Update A in KKT matrix.
-  for (QOCOInt i = 0; i < data->A->nnz; ++i) {
-    solver->work->kkt->K
-        ->x[solver->work->kkt->AttoKKT[solver->work->data->AtoAt[i]]] =
-        data->A->x[i];
-  }
-
-  // Update G in KKT matrix.
-  for (QOCOInt i = 0; i < data->G->nnz; ++i) {
-    solver->work->kkt->K
-        ->x[solver->work->kkt->GttoKKT[solver->work->data->GtoGt[i]]] =
-        data->G->x[i];
-  }
+  // // Update P in KKT matrix.
+  // for (QOCOInt i = 0; i < data->P->nnz; ++i) {
+  //   solver->work->kkt->K->x[solver->work->kkt->PregtoKKT[i]] = data->P->x[i];
+  // }
+
+  // // Update A in KKT matrix.
+  // for (QOCOInt i = 0; i < data->A->nnz; ++i) {
+  //   solver->work->kkt->K
+  //       ->x[solver->work->kkt->AttoKKT[solver->work->data->AtoAt[i]]] =
+  //       data->A->x[i];
+  // }
+
+  // // Update G in KKT matrix.
+  // for (QOCOInt i = 0; i < data->G->nnz; ++i) {
+  //   solver->work->kkt->K
+  //       ->x[solver->work->kkt->GttoKKT[solver->work->data->GtoGt[i]]] =
+  //       data->G->x[i];
+  // }
 }
 
 QOCOInt qoco_solve(QOCOSolver* solver)
@@ -419,17 +431,17 @@ QOCOInt qoco_cleanup(QOCOSolver* solver)
 {
 
   // Free problem data.
-  free_qoco_csc_matrix(solver->work->data->P);
-  free_qoco_csc_matrix(solver->work->data->A);
-  free_qoco_csc_matrix(solver->work->data->G);
-  free_qoco_csc_matrix(solver->work->data->At);
-  free_qoco_csc_matrix(solver->work->data->Gt);
+  free_qoco_matrix(solver->work->data->P);
+  free_qoco_matrix(solver->work->data->A);
+  free_qoco_matrix(solver->work->data->G);
+  free_qoco_matrix(solver->work->data->At);
+  free_qoco_matrix(solver->work->data->Gt);
 
   qoco_free(solver->work->data->AtoAt);
   qoco_free(solver->work->data->GtoGt);
-  qoco_free(solver->work->data->b);
-  qoco_free(solver->work->data->c);
-  qoco_free(solver->work->data->h);
+  free_qoco_vectorf(solver->work->data->b);
+  free_qoco_vectorf(solver->work->data->c);
+  free_qoco_vectorf(solver->work->data->h);
   qoco_free(solver->work->data->q);
   qoco_free(solver->work->data);