license import { getGlslDifferences } from './glsl_version' ; import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; import * as shader_util from './shader_compiler_util' ; export class EncodeMatrixProgram implements GPGPUProgram { variableNames = [ 'A' ] ; userCode : string ; outputShape : number [ ] ; enableShapeUniforms : boolean ; customUniforms = [ { name : 'texShape' , type : 'ivec2' as const } ] ; constructor ( outputShape , inputIsUnsignedByte = false ) { const glsl = getGlslDifferences ( ) ; this . outputShape = outputShape ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; let output = `result` ; if ( inputIsUnsignedByte ) { output = `floor(result * 255. + 0.5)` ; } this . userCode = `
      ${ this . enableShapeUniforms ? shader_util . getFlatIndexFrom3DOutput ( ) : shader_util . getFlatIndexFrom3D ( outputShape ) }

      void main() {
        ivec3 coords = getOutputCoords();

        int flatIndex = getFlatIndex(coords);
        int offset = imod(flatIndex, 4);

        flatIndex = idiv(flatIndex, 4, 1.);

        int r = flatIndex / texShape[1];
        int c = imod(flatIndex, texShape[1]);
        vec2 uv = (vec2(c, r) + halfCR) / vec2(texShape[1], texShape[0]);
        vec4 values = ${ glsl . texture2D }(A, uv);

        float result;

        if(offset == 0) {
          result = values[0];
        } else if(offset == 1) {
          result = values[1];
        } else if(offset == 2) {
          result = values[2];
        } else {
          result = values[3];
        }

        ${ glsl . output } = vec4(${ output }, 0., 0., 0.);
      }
    ` ; } }