\#4: port density update to shader preprocessor

src/shader/cellKey.ts

@@ -0,0 +1,11 @@
+import {glsl, GLSLDefinition} from '../glslpp';
+import {cellResolution} from './uniforms';
+
+export const cellKey = new GLSLDefinition(
+  'cellKey',
+  glsl`
+int cellKey(ivec2 cellPos) {
+  return cellPos.y * ${cellResolution}.x + cellPos.x;
+}
+`
+);

src/shader/foreachNeighbor.ts

@@ -0,0 +1,40 @@
+import {glsl, GLSLDefinition} from '../glslpp';
+import {cellKey} from './cellKey';
+import {posToCell} from './posToCell';
+import {
+  cellResolution,
+  hSmoothing,
+  keyParticleResolution,
+  keyParticleSampler,
+  neighborsTableSampler,
+  resolution,
+} from './uniforms';
+
+export const foreachNeighbor = new GLSLDefinition(
+  'FOREACH_NEIGHBOR',
+  glsl`
+#define FOREACH_NEIGHBOR(TEXCOORD_NAME, BLOCK) \
+  { \
+    int tableSize = ${cellResolution}.x * ${cellResolution}.y; \
+    ivec2 c0 = ${posToCell}(ownPos - vec2(${hSmoothing})); \
+    ivec2 c1 = ${posToCell}(ownPos + vec2(${hSmoothing})); \
+    for (int cx = c0.x; cx <= c1.x; ++cx) { \
+      for (int cy = c0.y; cy <= c1.y; ++cy) { \
+        int hash = ${cellKey}(ivec2(cx, cy)); \
+        ivec2 ntTexCoord = ivec2(hash % ${cellResolution}.x, hash / ${cellResolution}.x); \
+        ivec2 startCount = ivec2(texelFetch(${neighborsTableSampler}, ntTexCoord, 0).xy); \
+        int start = startCount.x; \
+        int count = startCount.y; \
+        for (int j = start; j < start + count; ++j) { \
+          ivec2 kpTexCoord = ivec2(j % ${keyParticleResolution}.x, j / ${keyParticleResolution}.x); \
+          int neighborIndex = texelFetch(${keyParticleSampler}, kpTexCoord, 0).y; \
+          ivec2 TEXCOORD_NAME = ivec2(neighborIndex % ${resolution}.x, neighborIndex / ${resolution}.x); \
+          { \
+            BLOCK \
+          } \
+        } \
+      } \
+    } \
+  }
+`
+);

src/shader/kernel.ts

@@ -0,0 +1,19 @@
+import {glsl, GLSLDefinition} from '../glslpp';
+import {hSmoothing, sigma} from './uniforms';
+
+export const W = new GLSLDefinition(
+  'W',
+  glsl`
+float W(vec2 dx) {
+  // cubic spline kernel
+  float q = sqrt(pow(dx.x, 2.) + pow(dx.y, 2.)) / ${hSmoothing};
+  if (0. <= q && q <= 0.5) {
+    return ${sigma} * (6. * (pow(q, 3.) - pow(q, 2.)) + 1.);
+  } else if (0.5 < q && q <= 1.) {
+    return ${sigma} * (2. * pow(1. - q, 3.));
+  } else {
+    return 0.;
+  }
+}
+`
+);

src/shader/posToCell.ts

@@ -0,0 +1,11 @@
+import {glsl, GLSLDefinition} from '../glslpp';
+import {cellResolution, cellSize} from './uniforms';
+
+export const posToCell = new GLSLDefinition(
+  'posToCell',
+  glsl`
+ivec2 posToCell(vec2 pos) {
+  return max(ivec2(0), min(${cellResolution} - ivec2(1), ivec2(floor(pos / ${cellSize}))));
+}
+`
+);

src/shader/uniforms.ts

@@ -1,6 +1,27 @@
 import {GLSLUniform} from '../glslpp';
 
+export const neighborsTableSampler = new GLSLUniform(
+  'neighborsTableSampler',
+  'sampler2D'
+);
+export const massSampler = new GLSLUniform('massSampler', 'sampler2D');
 export const positionSampler = new GLSLUniform('positionSampler', 'sampler2D');
 export const velocitySampler = new GLSLUniform('velocitySampler', 'sampler2D');
 
+export const keyParticleSampler = new GLSLUniform(
+  'keyParticleSampler',
+  'isampler2D'
+);
+
+export const cellSize = new GLSLUniform('cellSize', 'vec2');
+
+export const cellResolution = new GLSLUniform('cellResolution', 'ivec2');
+export const keyParticleResolution = new GLSLUniform(
+  'keyParticleResolution',
+  'ivec2'
+);
+export const resolution = new GLSLUniform('resolution', 'ivec2');
+
 export const dt = new GLSLUniform('dt', 'float');
+export const hSmoothing = new GLSLUniform('hSmoothing', 'float');
+export const sigma = new GLSLUniform('sigma', 'float');

src/shader/updateDensity.ts

@@ -0,0 +1,26 @@
+import {compile, glsl} from '../glslpp';
+import {foreachNeighbor} from './foreachNeighbor';
+import {W} from './kernel';
+import {massSampler, positionSampler, resolution} from './uniforms';
+
+export const updateDensityFs = compile(glsl`
+out vec4 densityOut;
+
+void main() {
+  int particleIndex = int(gl_FragCoord.y) * ${resolution}.x + int(gl_FragCoord.x);
+  ivec2 ownTexCoord = ivec2(gl_FragCoord.xy);
+  float ownMass = texelFetch(${massSampler}, ownTexCoord, 0).x;
+  vec2 ownPos = texelFetch(${positionSampler}, ownTexCoord, 0).xy;
+
+  float density = ownMass * ${W}(vec2(0.));
+
+  ${foreachNeighbor}(neighborTexCoord, {
+    vec2 neighborPos = texelFetch(${positionSampler}, neighborTexCoord, 0).xy;
+    float neighborMass = texelFetch(${massSampler}, neighborTexCoord, 0).x;
+    vec2 dx = neighborPos - ownPos;
+    density += neighborMass * ${W}(dx);
+  });
+
+  densityOut = vec4(density, 0, 0, 0);
+}
+`);

src/simulationGPU.ts

@@ -5,7 +5,7 @@ import {advectParticlesFs} from './shader/advectParticles';
 import updateVelocityFrag from './updateVelocity.frag.glsl';
 import {Params} from './params';
 import {getCopyVertexVert, getQuadVAO} from './gpuUtil';
-import updateDensityFrag from './updateDensity.frag.glsl';
+import {updateDensityFs} from './shader/updateDensity';
 import updateVelocityGuessFrag from './updateVelocityGuess.frag.glsl';
 import updateFPressureFrag from './updateFPressure.frag.glsl';
 import {updateNeighborsGPU} from './neighborsGPU';
@@ -153,7 +153,7 @@ export const copyStateFromGPU = (
 };
 
 const getUpdateDensityFrag = memoize((gl: WebGL2RenderingContext) =>
-  createShader(gl, {source: updateDensityFrag, type: gl.FRAGMENT_SHADER})
+  createShader(gl, {source: updateDensityFs, type: gl.FRAGMENT_SHADER})
 );
 
 const getUpdateDensityProgram = memoize((gl: WebGL2RenderingContext) =>