3 files changed, 42 insertions, 21 deletions

diff --git a/src/shader/code/collide.glsl b/src/shader/code/collide.glsl
index 67e762b..90af535 100644
--- a/src/shader/code/collide.glsl
+++ b/src/shader/code/collide.glsl
@@ -2,34 +2,43 @@ static const std::string COLLIDE_SHADER_CODE = R"(
 #version 430
 
 layout (local_size_x = 1, local_size_y = 1) in;
+
 layout (std430, binding=1) buffer bufferCollide{ float collideCells[]; };
 layout (std430, binding=2) buffer bufferStream{  float streamCells[]; };
 layout (std430, binding=3) buffer bufferFluid{   float fluidCells[]; };
 
+uniform uint nX;
+uniform uint nY;
+
+const uint  q     = 9;
+const float omega = 0.6;
+
+const float velocityDisplayScalar = 500.;
+
 float get(uint x, uint y, int i, int j) {
-	return collideCells[9*128*y + 9*x + (i+1)*3 + j+1];
+	return collideCells[q*nX*y + q*x + (i+1)*3 + j+1];
 }
 
 void set(uint x, uint y, int i, int j, float v) {
-	collideCells[9*128*y + 9*x + (i+1)*3 + j+1] = v;
+	collideCells[q*nX*y + q*x + (i+1)*3 + j+1] = v;
 }
 
 void setFluid(uint x, uint y, vec2 v, float d) {
-	fluidCells[3*128*y + 3*x + 0] = float(x)-64. + 10.*v.x;
-	fluidCells[3*128*y + 3*x + 1] = float(y)-64. + 10.*v.y;
-	fluidCells[3*128*y + 3*x + 2] = d;
+	fluidCells[3*nX*y + 3*x + 0] = float(x)-nX/2 + velocityDisplayScalar*v.x;
+	fluidCells[3*nX*y + 3*x + 1] = float(y)-nY/2 + velocityDisplayScalar*v.y;
+	fluidCells[3*nX*y + 3*x + 2] = d;
 }
 
 float density(uint x, uint y) {
-	return collideCells[9*128*y + 9*x + 0]
-	     + collideCells[9*128*y + 9*x + 1]
-	     + collideCells[9*128*y + 9*x + 2]
-	     + collideCells[9*128*y + 9*x + 3]
-	     + collideCells[9*128*y + 9*x + 4]
-	     + collideCells[9*128*y + 9*x + 5]
-	     + collideCells[9*128*y + 9*x + 6]
-	     + collideCells[9*128*y + 9*x + 7]
-	     + collideCells[9*128*y + 9*x + 8];
+	return collideCells[q*nX*y + q*x + 0]
+	     + collideCells[q*nX*y + q*x + 1]
+	     + collideCells[q*nX*y + q*x + 2]
+	     + collideCells[q*nX*y + q*x + 3]
+	     + collideCells[q*nX*y + q*x + 4]
+	     + collideCells[q*nX*y + q*x + 5]
+	     + collideCells[q*nX*y + q*x + 6]
+	     + collideCells[q*nX*y + q*x + 7]
+	     + collideCells[q*nX*y + q*x + 8];
 }
 
 vec2 velocity(uint x, uint y, float d) {
@@ -77,8 +86,6 @@ void main() {
 	const uint x = gl_GlobalInvocationID.x;
 	const uint y = gl_GlobalInvocationID.y;
 
-	const float omega = 0.6;
-
 	const float d = density(x,y);
 	const vec2  v = velocity(x,y,d);
 
diff --git a/src/shader/code/stream.glsl b/src/shader/code/stream.glsl
index 7dd5ba8..147b49b 100644
--- a/src/shader/code/stream.glsl
+++ b/src/shader/code/stream.glsl
@@ -2,29 +2,43 @@ static const std::string STREAM_SHADER_CODE = R"(
 #version 430
 
 layout (local_size_x = 1, local_size_y = 1) in;
+
 layout (std430, binding=1) buffer bufferCollide{ float collideCells[]; };
 layout (std430, binding=2) buffer bufferStream{  float streamCells[]; };
 
+uniform uint nX;
+uniform uint nY;
+
+const uint q = 9;
+
 float get(uint x, uint y, int i, int j) {
-	return collideCells[9*128*y + 9*x + (i+1)*3 + j+1];
+	return collideCells[q*nX*y + q*x + (i+1)*3 + j+1];
 }
 
 void set(uint x, uint y, int i, int j, float v) {
-	streamCells[9*128*y + 9*x + (i+1)*3 + j+1] = v;
+	streamCells[q*nX*y + q*x + (i+1)*3 + j+1] = v;
 }
 
 void main() {
 	const uint x = gl_GlobalInvocationID.x;
 	const uint y = gl_GlobalInvocationID.y;
 
-	if ( x != 0 && x != 128-1 && y != 0 && y != 128-1 ) {
+	if ( x != 0 && x != nX-1 && y != 0 && y != nY-1 ) {
 		for ( int i = -1; i <= 1; ++i ) {
 			for ( int j = -1; j <= 1; ++j ) {
 				set(x+i,y+j,i,j, get(x,y,i,j));
 			}
 		}
 	} else {
-
+		for ( int i = -1; i <= 1; ++i ) {
+			for ( int j = -1; j <= 1; ++j ) {
+				if ( (x > 0 || i >= 0) && x+i <= nX-1 && (y > 0 || j >= 0) && y+j <= nY-1 ) {
+					set(x+i,y+j,i,j, get(x,y,i,j));
+				} else {
+					set(x,y,i*(-1),j*(-1), get(x,y,i,j));
+				}
+			}
+		}
 	}
 }
 )";
diff --git a/src/shader/code/vertex.glsl b/src/shader/code/vertex.glsl
index 4c307d8..0eccad4 100644
--- a/src/shader/code/vertex.glsl
+++ b/src/shader/code/vertex.glsl
@@ -3,6 +3,6 @@ uniform mat4 MVP;
 
 void main() {
 	gl_Position = MVP * vec4(gl_Vertex.xy, 0.0, 1.0);
-	gl_FrontColor = vec4(1., 0., 0., 0.);
+	gl_FrontColor = gl_Vertex.z * vec4(1., 0., 0., 0.);
 }
 )";