1 files changed, 15 insertions, 11 deletions

diff --git a/ell_grid_solve.c b/ell_grid_solve.c
index 13d9072..7dcd41e 100644
--- a/ell_grid_solve.c
+++ b/ell_grid_solve.c
@@ -73,10 +73,10 @@ typedef struct EGSExactInternal {
 } EGSExactInternal;
 
 struct EGSInternal {
-    NDArray *u_base;
     NDArray *rhs_base;
     NDArray *residual_base;
 
+    NDArray *u_next_base;
     NDArray *u_next;
     int      u_next_valid;
 
@@ -331,9 +331,12 @@ static int solve_relax_step(EGSContext *ctx, int export_res)
     int u_next_valid = priv->u_next_valid;
 
     if (u_next_valid) {
-        NDArray *tmp = ctx->u;
+        NDArray *tmp      = ctx->u;
+        NDArray *tmp_base = ctx->u_base;
         ctx->u       = priv->u_next;
+        ctx->u_base  = priv->u_next_base;
         priv->u_next = tmp;
+        priv->u_next_base = tmp_base;
         priv->u_next_valid = 0;
     }
 
@@ -930,26 +933,26 @@ static int arrays_alloc(EGSContext *ctx, const size_t domain_size[2])
         size_padded[0] * MG2D_DIFF_COEFF_NB,
         size_padded[1],
     };
-    const size_t size_u[2] = {
-        size_padded[0] * 2,
-        size_padded[1],
-    };
     const Slice slice[2] = { SLICE(ghosts, -ghosts, 1),
                              SLICE(ghosts, -ghosts, 1) };
     int ret;
 
-    ret = mg2di_ndarray_alloc(&priv->u_base, 2, size_u, NDARRAY_ALLOC_ZERO);
+    ret = mg2di_ndarray_alloc(&ctx->u_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&ctx->u, priv->u_base,
+    ret = mg2di_ndarray_slice(&ctx->u, ctx->u_base,
                               (Slice [2]){ SLICE(ghosts, size_padded[0] - ghosts, 1),
                                            SLICE(ghosts, -ghosts, 1) });
     if (ret < 0)
         return ret;
 
-    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_base,
-                              (Slice [2]){ SLICE(size_padded[0] + ghosts, size_u[0] - ghosts, 1),
+    ret = mg2di_ndarray_alloc(&priv->u_next_base, 2, size_padded, NDARRAY_ALLOC_ZERO);
+    if (ret < 0)
+        return ret;
+
+    ret = mg2di_ndarray_slice(&priv->u_next, priv->u_next_base,
+                              (Slice [2]){ SLICE(ghosts, size_padded[0] - ghosts, 1),
                                            SLICE(ghosts, -ghosts, 1) });
     if (ret < 0)
         return ret;
@@ -1079,8 +1082,9 @@ void mg2di_egs_free(EGSContext **pctx)
     exact_arrays_free(ctx);
 
     mg2di_ndarray_free(&ctx->u);
+    mg2di_ndarray_free(&ctx->u_base);
     mg2di_ndarray_free(&ctx->priv->u_next);
-    mg2di_ndarray_free(&ctx->priv->u_base);
+    mg2di_ndarray_free(&ctx->priv->u_next_base);
 
     mg2di_ndarray_free(&ctx->rhs);
     mg2di_ndarray_free(&ctx->priv->rhs_base);