\#4: port sort to shader preprocessor, remove old odd-even sort

src/shader/sortOddEvenMerge.ts

@@ -0,0 +1,57 @@
+import {compile, glsl} from '../glslpp';
+import {
+  compareWidth,
+  inputSampler,
+  n,
+  resolution,
+  stageWidth,
+} from './uniforms';
+
+export const sortOddEvenMergeFs = compile(glsl`
+out ivec4 outResult;
+
+void main() {
+  // self index
+  int i = int(gl_FragCoord.y) * ${resolution}.x + int(gl_FragCoord.x);
+
+  // is this the left side of the pair (0) or the right (1)?
+  bool isLeft;
+  if (${compareWidth} < ${stageWidth}) {
+    // this is a merge pass; add an offset
+    isLeft = (i / ${compareWidth} + 1) % 2 == 0;
+  } else {
+    // first pass
+    isLeft = (i / ${compareWidth}) % 2 == 0;
+  }
+  int direction = isLeft ? 1 : -1;
+  
+  // pair index
+  int j = i + direction * ${compareWidth};
+
+  // read value of self and pair
+  ivec2 texCoordi = ivec2(i % ${resolution}.x, i / ${resolution}.x);
+  ivec2 texCoordj = ivec2(j % ${resolution}.x, j / ${resolution}.x);
+  // values consist of (key, value)
+  ivec2 vi = texelFetch(${inputSampler}, texCoordi, 0).rg;
+  ivec2 vj = texelFetch(${inputSampler}, texCoordj, 0).rg;
+
+  ivec2 result;
+  int sw2 = ${stageWidth} * 2;
+  int lower = (i / sw2) * sw2;
+  int upper = lower + sw2;
+  if (j < lower || j >= upper || j >= ${n}) {
+    // pair is not in bounds; simply copy own value
+    result = vi;
+  } else {
+    if (isLeft) {
+      // self is on the left side of the pair; pick <
+      result = vi.x < vj.x ? vi : vj;
+    } else {
+      // self is on the right side of the pair; pick >
+      result = vi.x > vj.x ? vi : vj;
+    }
+  }
+
+  outResult = ivec4(result, 0, 0);
+}
+`);

src/shader/uniforms.ts

@@ -3,6 +3,7 @@ import {GLSLUniform} from '../glslpp';
 export const cellResolution = new GLSLUniform('cellResolution', 'ivec2');
 export const cellSize = new GLSLUniform('cellSize', 'vec2');
 export const collisionDistance = new GLSLUniform('collisionDistance', 'float');
+export const compareWidth = new GLSLUniform('compareWidth', 'int');
 export const densitySampler = new GLSLUniform('densitySampler', 'sampler2D');
 export const dt = new GLSLUniform('dt', 'float');
 export const eta = new GLSLUniform('eta', 'float');
@@ -11,6 +12,7 @@ export const fPressureSampler = new GLSLUniform(
   'sampler2D'
 );
 export const hSmoothing = new GLSLUniform('hSmoothing', 'float');
+export const inputSampler = new GLSLUniform('inputSampler', 'isampler2D');
 export const keyParticleResolution = new GLSLUniform(
   'keyParticleResolution',
   'ivec2'
@@ -20,6 +22,7 @@ export const keyParticleSampler = new GLSLUniform(
   'isampler2D'
 );
 export const massSampler = new GLSLUniform('massSampler', 'sampler2D');
+export const n = new GLSLUniform('n', 'int');
 export const neighborsTableSampler = new GLSLUniform(
   'neighborsTableSampler',
   'sampler2D'
@@ -37,6 +40,7 @@ export const projection = new GLSLUniform('projection', 'mat4');
 export const resolution = new GLSLUniform('resolution', 'ivec2');
 export const restDensity = new GLSLUniform('restDensity', 'float');
 export const sigma = new GLSLUniform('sigma', 'float');
+export const stageWidth = new GLSLUniform('stageWidth', 'int');
 export const stiffness = new GLSLUniform('stiffness', 'float');
 export const velocityGuessSampler = new GLSLUniform(
   'velocityGuessSampler',

src/sortGPU.ts

@@ -5,17 +5,9 @@ import {
   PingPongTexture,
 } from './gl';
 import {memoize, nextPowerOf2, shuffle, time} from './util';
-import sortEvenOddFrag from './sortEvenOdd.frag.glsl';
-import sortOddEvenMergeFrag from './sortOddEvenMerge.frag.glsl';
+import {sortOddEvenMergeFs} from './shader/sortOddEvenMerge';
 import {getCopyVertexVert, getQuadVAO} from './gpuUtil';
 
-const getSortEvenOddFrag = memoize((gl: WebGL2RenderingContext) =>
-  createShader(gl, {source: sortEvenOddFrag, type: gl.FRAGMENT_SHADER})
-);
-const getSortEvenOddProgram = memoize((gl: WebGL2RenderingContext) =>
-  createProgram(gl, {shaders: [getCopyVertexVert(gl), getSortEvenOddFrag(gl)]})
-);
-
 const copyToTextureInt = (
   gl: WebGL2RenderingContext,
   src: Int32Array,
@@ -42,43 +34,8 @@ const copyFromTextureInt = (
   gl.bindFramebuffer(gl.FRAMEBUFFER, null);
 };
 
-export const sortEvenOdd = (
-  gl: WebGL2RenderingContext,
-  texture: PingPongTexture,
-  width: number,
-  height: number
-) => {
-  const program = getSortEvenOddProgram(gl);
-
-  gl.useProgram(program);
-  gl.bindVertexArray(getQuadVAO(gl));
-
-  gl.viewport(0, 0, width, height);
-  gl.uniform2i(gl.getUniformLocation(program, 'resolution'), width, height);
-  const inputSamplerLoc = gl.getUniformLocation(program, 'inputSampler');
-  const oddStepLoc = gl.getUniformLocation(program, 'oddStep');
-
-  for (let i = 0; i < width * height; ++i) {
-    gl.activeTexture(gl.TEXTURE0);
-    gl.bindTexture(gl.TEXTURE_2D, texture.read.texture);
-    gl.uniform1i(inputSamplerLoc, 0);
-    gl.uniform1i(oddStepLoc, i % 2);
-
-    gl.bindFramebuffer(gl.FRAMEBUFFER, texture.write.framebuffer);
-
-    gl.drawArrays(gl.TRIANGLE_FAN, 0, 4);
-
-    texture.swap();
-  }
-
-  gl.bindFramebuffer(gl.FRAMEBUFFER, null);
-  gl.bindTexture(gl.TEXTURE_2D, null);
-  gl.bindVertexArray(null);
-  gl.useProgram(null);
-};
-
 const getSortOddEvenMergeFrag = memoize((gl: WebGL2RenderingContext) =>
-  createShader(gl, {source: sortOddEvenMergeFrag, type: gl.FRAGMENT_SHADER})
+  createShader(gl, {source: sortOddEvenMergeFs, type: gl.FRAGMENT_SHADER})
 );
 const getSortOddEvenMergeProgram = memoize((gl: WebGL2RenderingContext) =>
   createProgram(gl, {