const wgslcompute = `
struct Data {
    numbers: array<f32>
};

@group(0) @binding(0) var<storage, read_write> data : Data;

fn compareAndSwap(i: u32, j: u32, dir: bool) {
    if ((data.numbers[i] > data.numbers[j]) == dir) {
        let temp = data.numbers[i];
        data.numbers[i] = data.numbers[j];
        data.numbers[j] = temp;
    }
}

@compute @workgroup_size(64)
fn main( @builtin(global_invocation_id) id: vec3<u32>) 
{
    var k: u32 = 2;
    let datalength = arrayLength( &data.numbers );
    
    while ( k <= datalength ) 
    {
        var j: u32 = k >> 1;
        while (j > 0) {
            let idx = id.x;
            let ixj = idx ^ j;
            if (ixj > idx) {
                let direction = ((idx & k) == 0);
                compareAndSwap(idx, ixj, direction);
            }
            j = j >> 1;
            workgroupBarrier();
        }
        k = k << 1;
    }
}
`;

let div = document.createElement('div');
document.body.appendChild( div );
div.style['font-size'] = '20pt';
function log( s )
{
  console.log( s );
  let args = [...arguments].join(' ');
  div.innerHTML += args + '<br>';
}

log('WebGPU Compute Example');

if (!navigator.gpu) { log("WebGPU is not supported (or is it disabled? flags/settings)"); return; }

const adapter = await navigator.gpu.requestAdapter();
const device  = await adapter.requestDevice();


const arraySize = 128;
const data0 = new Float32Array(arraySize);
for (let i = 0; i < arraySize; i++) {
  data0[i] = Math.round( Math.random() * 100.0 );
}

const gbuffer0 = device.createBuffer({ size:  data0.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC });
device.queue.writeBuffer(gbuffer0, 0, data0);

const resultBuffer = device.createBuffer({ size: data0.byteLength, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });

const readbackBuffer = device.createBuffer({ size: data0.byteLength, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });


const bindGroupLayout = device.createBindGroupLayout({
  entries: [ {binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: {type: "storage"} } ]
});

const bindGroup = device.createBindGroup({
  layout: bindGroupLayout,
  entries: [{ binding: 0,  resource: {  buffer: gbuffer0  }  }]
});

const shaderModule = device.createShaderModule( {  code: wgslcompute } );
  
const computePipeline = device.createComputePipeline({
  layout: device.createPipelineLayout({bindGroupLayouts: [bindGroupLayout]}),
  compute: { module: shaderModule,
             entryPoint: 'main'  }
});

const commandEncoder = device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(computePipeline);
passEncoder.setBindGroup(0, bindGroup);
passEncoder.dispatchWorkgroups(arraySize / 64);
passEncoder.end();

commandEncoder.copyBufferToBuffer(gbuffer0, 0, resultBuffer, 0, data0.byteLength);
commandEncoder.copyBufferToBuffer(resultBuffer, 0, readbackBuffer, 0, data0.byteLength);

const commands = commandEncoder.finish();
device.queue.submit([commands]);

await readbackBuffer.mapAsync(GPUMapMode.READ);
const sortedArray = new Float32Array(readbackBuffer.getMappedRange());
console.log(sortedArray);

log('Sorted array:')
const arr = Array.from( sortedArray );
for (let i=0; i<arr.length; i++)
{
      // display over multiple lines
    log( arr.splice( 0, 10 ) );  
}