diff --git a/src/main/c/densitystep.c b/src/main/c/densitystep.c
index 9b4f49a..b25e82f 100644
--- a/src/main/c/densitystep.c
+++ b/src/main/c/densitystep.c
@@ -40,8 +40,8 @@ JNIEXPORT void JNICALL Java_electrosphere_FluidSim_solveDiffuseDensity
   jobject this,
   jint N,
   jint chunk_mask,
-  jobjectArray jrx,
-  jobjectArray jrx0,
+  float ** d,
+  float ** d0,
   jobjectArray jru,
   jobjectArray jrv,
   jobjectArray jrw,
@@ -51,8 +51,8 @@ JNIEXPORT void JNICALL Java_electrosphere_FluidSim_solveDiffuseDensity
     float a=dt*DIFFUSION_CONST*N*N*N;
     float c=1+6*a;
     int i, j, k, l, m;
-    float * x = GET_ARR(env,jrx,CENTER_LOC);
-    float * x0 = GET_ARR(env,jrx0,CENTER_LOC);
+    float * x = GET_ARR_RAW(env,d,CENTER_LOC);
+    float * x0 = GET_ARR_RAW(env,d0,CENTER_LOC);
     
     __m256 aScalar = _mm256_set1_ps(a);
     __m256 cScalar = _mm256_set1_ps(c);
diff --git a/src/main/c/fluidsim.c b/src/main/c/fluidsim.c
index c29d9ff..28d00bd 100644
--- a/src/main/c/fluidsim.c
+++ b/src/main/c/fluidsim.c
@@ -772,7 +772,7 @@ JNIEXPORT void JNICALL Java_electrosphere_FluidSim_simulate(
                 u0 = currentChunk->ju0;
                 v0 = currentChunk->jv0;
                 w0 = currentChunk->jw0;
-                Java_electrosphere_FluidSim_solveDiffuseDensity(env,chunkJRaw,DIM,chunkMask,jd,jd0,u,v,w,DIFFUSION_CONSTANT,VISCOSITY_CONSTANT,timestep);
+                Java_electrosphere_FluidSim_solveDiffuseDensity(env,chunkJRaw,DIM,chunkMask,currentChunk->d,currentChunk->d0,u,v,w,DIFFUSION_CONSTANT,VISCOSITY_CONSTANT,timestep);
             }
             for(int i = 0; i < numChunks; i++){
                 Chunk * currentChunk = chunks[i];
@@ -786,7 +786,7 @@ JNIEXPORT void JNICALL Java_electrosphere_FluidSim_simulate(
                 u0 = currentChunk->ju0;
                 v0 = currentChunk->jv0;
                 w0 = currentChunk->jw0;
-                Java_electrosphere_FluidSim_setBoundsToNeighbors(env,chunkJRaw,DIM,chunkMask,0,jd);
+                setBoundsToNeighborsRaw(env,chunkJRaw,DIM,chunkMask,0,currentChunk->d);
             }
         }
     }
diff --git a/src/main/c/includes/mainFunctions.h b/src/main/c/includes/mainFunctions.h
index 740f66d..03df600 100644
--- a/src/main/c/includes/mainFunctions.h
+++ b/src/main/c/includes/mainFunctions.h
@@ -73,7 +73,7 @@ JNIEXPORT void JNICALL Java_electrosphere_FluidSim_addDensity
  * Signature: (II[Ljava/nio/ByteBuffer;[Ljava/nio/ByteBuffer;[Ljava/nio/ByteBuffer;[Ljava/nio/ByteBuffer;[Ljava/nio/ByteBuffer;FFF)V
  */
 JNIEXPORT void JNICALL Java_electrosphere_FluidSim_solveDiffuseDensity
-  (JNIEnv *, jobject, jint, jint, jobjectArray, jobjectArray, jobjectArray, jobjectArray, jobjectArray, jfloat, jfloat, jfloat);
+  (JNIEnv *, jobject, jint, jint, float **, float **, jobjectArray, jobjectArray, jobjectArray, jfloat, jfloat, jfloat);
 
 /*
  * Class:     electrosphere_FluidSim