Update `fmpi_new_domain()` and `fmpi_domain_set_inner()` for 2D support

src/fmpi_domain.c

@@ -84,10 +84,32 @@ struct fmpi_domain fmpi_new_domain(
         .stencil = stencil,
         .part_cnt = proc_cnt
     };
-    domain.inner = fmpi_partition_block_1d(ctx, data);
+    switch(data->dim_cnt) {
+        case 1: {
+            domain.inner = fmpi_partition_block_1d(ctx, data);
+            break;
+        }
+        case 2: {
+            domain.inner = fmpi_partition_block_2d(ctx, data);
+            break;
+        }
+        default: break;
+    }
     if(stencil.type != FMPI_STENCIL_NONE) {
-        domain.halo = fmpi_halo_1d(ctx, domain.inner, stencil);
-        domain.inner.raw = (char *)domain.halo.raw + (data->type.size * stencil.length);
+        switch(data->dim_cnt) {
+            case 1: {
+                domain.halo = fmpi_halo_1d(ctx, domain.inner, stencil);
+                domain.inner.raw = (char *)domain.halo.raw + (data->type.size * stencil.length);
+                break;
+            }
+            case 2: {
+                domain.halo = fmpi_halo_2d(ctx, domain.inner, stencil);
+                const size_t offset = stencil.length * (domain.halo.dim_len[0] + 1) * data->type.size;
+                domain.inner.raw = (char *)domain.halo.raw + offset;
+                break;
+            }
+            default: break;
+        }
     }
     return domain;
 }
@@ -101,7 +123,20 @@ void fmpi_domain_set_inner(
     assert(ctx != NULL);
     assert(domain != NULL);
     assert(data != NULL);
-    memcpy(domain->inner.raw, data, domain->inner.size);
+    const size_t dim_cnt = domain->inner.dim_cnt;
+    if(dim_cnt == 1) {
+        memcpy(domain->inner.raw, data, domain->inner.size);
+    } else if(dim_cnt == 2) {
+        const size_t type_size = domain->inner.type.size;
+        const size_t dim_len_y = domain->inner.dim_len[1];
+        const size_t src_size_x =  domain->inner.dim_len[0] * type_size;
+        for(size_t i = 0; i < dim_len_y; i++) {
+            const size_t dest_offset = i * domain->halo.dim_len[0] * type_size;
+            void * const dest = (char *)domain->inner.raw + dest_offset;
+            const void * const src = (const char *)data + (i * src_size_x);
+            memcpy(dest, src, src_size_x);
+        }
+    }
 }
 /*------------------------------------------------------------------------------
     fmpi_halo_exchange_1d()
@@ -256,13 +291,13 @@ static struct fmpi_data fmpi_halo_1d(
         );
     }
     const size_t rank = (size_t)ctx->mpi->rank;
-    // Left boundary
+    // First rank
     if(rank == 0) {
         memcpy(((char *)halo.raw + type_size), data.raw, (data.size + type_size));
-    // Right boundary
+    // Last rank
     } else if(rank == (size_t)(ctx->mpi->size - 1)) {
         memcpy(halo.raw, ((const char *)(data.raw) - type_size), (data.size + type_size));
-    // Middle
+    // Middle rank
     } else {
         memcpy(halo.raw, ((const char *)(data.raw) - type_size), size);
     }