<!DOCTYPE html>
<html>
<head>
<meta name="viewport" content="width=1000">
<title>WebGPU Cube demo</title>
<script src="scripts/gl-matrix-min.js"></script>
<link rel="stylesheet" href="css/style.css"/>
</head>
<body>
<canvas></canvas>
<script>
const positionAttributeNum  = 0;
const texCoordsAttributeNum = 1;

const transformBindingNum   = 0;
const textureBindingNum     = 1;
const samplerBindingNum     = 2;

const vertexBufferIndex     = 0;
const bindGroupIndex        = 0;

const shader = `
#include <metal_stdlib>

using namespace metal;

struct Vertex
{
    float4 position [[attribute(${positionAttributeNum})]];
    float2 texCoords [[attribute(${texCoordsAttributeNum})]];
};

struct FragmentData
{
    float4 position [[position]];
    float2 texCoords;
};

struct Uniform
{
    device float4x4* modelViewProjectionMatrix [[id(${transformBindingNum})]];
};

struct SampledTexture
{
    texture2d<float> faceTexture [[id(${textureBindingNum})]];
    sampler faceSampler [[id(${samplerBindingNum})]];
};

vertex FragmentData vertex_main(Vertex vertexIn [[stage_in]],
                            const device Uniform& uniforms [[buffer(${bindGroupIndex})]])
{
    FragmentData output;
    output.position = uniforms.modelViewProjectionMatrix[0] * vertexIn.position;
    output.texCoords = vertexIn.texCoords;

    return output;
}

fragment float4 fragment_main(FragmentData data [[stage_in]],
                            const device SampledTexture& args [[buffer(${bindGroupIndex})]])
{
    float4 color = args.faceTexture.sample(args.faceSampler, data.texCoords);
    if (color.a < 1)
        discard_fragment();

    return color;
}
`

let device, swapChain, verticesBuffer, bindGroupLayout, pipeline, renderPassDescriptor, queue, textureViewBinding, samplerBinding;
let projectionMatrix = mat4.create();

const texCoordsOffset = 4 * 4;
const vertexSize = 4 * 6;
function createVerticesArray() {
    return new Float32Array([
        // float4 position, float2 texCoords
        1, -1, 1, 1, 1, 0, 
        -1, -1, 1, 1, 0, 0, 
        -1, -1, -1, 1, 0, 1, 
        1, -1, -1, 1, 1, 1,
        1, -1, 1, 1, 1, 0,
        -1, -1, -1, 1, 0, 1,

        1, 1, 1, 1, 0, 0,
        1, -1, 1, 1, 0, 1,
        1, -1, -1, 1, 1, 1,
        1, 1, -1, 1, 1, 0,
        1, 1, 1, 1, 0, 0,
        1, -1, -1, 1, 1, 1,

        -1, 1, 1, 1, 0, 1,
        1, 1, 1, 1, 1, 1,
        1, 1, -1, 1, 1, 0,
        -1, 1, -1, 1, 0, 0,
        -1, 1, 1, 1, 0, 1,
        1, 1, -1, 1, 1, 0,

        -1, -1, 1, 1, 1, 1,
        -1, 1, 1, 1, 1, 0,
        -1, 1, -1, 1, 0, 0,
        -1, -1, -1, 1, 0, 1,
        -1, -1, 1, 1, 1, 1,
        -1, 1, -1, 1, 0, 0,

        1, 1, 1, 1, 1, 0,
        -1, 1, 1, 1, 0, 0,
        -1, -1, 1, 1, 0, 1,
        -1, -1, 1, 1, 0, 1,
        1, -1, 1, 1, 1, 1,
        1, 1, 1, 1, 1, 0,

        1, -1, -1, 1, 0, 1,
        -1, -1, -1, 1, 1, 1,
        -1, 1, -1, 1, 1, 0,
        1, 1, -1, 1, 0, 0,
        1, -1, -1, 1, 0, 1,
        -1, 1, -1, 1, 1, 0,
    ]);
}

async function init() {
    const adapter = await navigator.gpu.requestAdapter();
    device = await adapter.requestDevice();

    const canvas = document.querySelector('canvas');
    let canvasSize = canvas.getBoundingClientRect();
    canvas.width = canvasSize.width;
    canvas.height = canvasSize.height;

    const aspect = Math.abs(canvas.width / canvas.height);
    mat4.perspective(projectionMatrix, (2 * Math.PI) / 5, aspect, 1, 100.0);

    const context = canvas.getContext('gpu');

    const swapChainDescriptor = { 
        device: device,
        format: "bgra8unorm"
    };
    swapChain = context.configureSwapChain(swapChainDescriptor);

    // WebKit WebGPU accepts only MSL for now.
    const shaderModuleDescriptor = { code: shader };
    const shaderModule = device.createShaderModule(shaderModuleDescriptor);

    const verticesArray = createVerticesArray();
    const verticesBufferDescriptor = { 
        size: verticesArray.byteLength, 
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.TRANSFER_DST
    };
    verticesBuffer = device.createBuffer(verticesBufferDescriptor);
    verticesBuffer.setSubData(0, verticesArray.buffer);

    // Input state. Model will change soon to adopt one of https://github.com/kainino0x/gpuweb/pull/2/'s ideas.
    const positionAttributeDescriptor = {
        shaderLocation: positionAttributeNum,  // [[attribute(0)]].
        inputSlot: vertexBufferIndex,       // Used as vertex buffer index in Metal.
        offset: 0,
        format: "float4"
    };
    const texCoordsAttributeDescriptor = {
        shaderLocation: texCoordsAttributeNum,
        inputSlot: vertexBufferIndex,
        offset: texCoordsOffset,
        format: "float2"
    }
    const vertexBufferDescriptor = {
        inputSlot: vertexBufferIndex,
        stride: vertexSize,
        stepMode: "vertex"
    };
    const inputStateDescriptor = {
        indexFormat: "uint32",
        attributes: [positionAttributeDescriptor, texCoordsAttributeDescriptor], 
        inputs: [vertexBufferDescriptor]
    };

    // Texture

    // Load texture image
    const image = new Image();
    const imageLoadPromise = new Promise(resolve => { 
        image.onload = () => resolve(); 
        image.src = "resources/safari-alpha.png"
    });
    await Promise.resolve(imageLoadPromise);

    const textureSize = {
        width: image.width,
        height: image.height,
        depth: 1
    };

    const textureDescriptor = {
        size: textureSize,
        arrayLayerCount: 1,
        mipLevelCount: 1,
        sampleCount: 1,
        dimension: "2d",
        format: "rgba8unorm",
        usage: GPUTextureUsage.TRANSFER_DST | GPUTextureUsage.SAMPLED
    };
    const texture = device.createTexture(textureDescriptor);

    // Texture data 
    const canvas2d = document.createElement('canvas');
    canvas2d.width = image.width;
    canvas2d.height = image.height;
    const context2d = canvas2d.getContext('2d');
    context2d.drawImage(image, 0, 0);

    const imageData = context2d.getImageData(0, 0, image.width, image.height);

    const textureDataBufferDescriptor = {
        size: imageData.data.length,
        usage: GPUBufferUsage.TRANSFER_SRC | GPUBufferUsage.TRANSFER_DST
    };
    const textureDataBuffer = device.createBuffer(textureDataBufferDescriptor);
    textureDataBuffer.setSubData(0, imageData.data.buffer);

    const dataCopyView = {
        buffer: textureDataBuffer,
        offset: 0,
        rowPitch: image.width * 4,
        imageHeight: 0
    };
    const textureCopyView = {
        texture: texture,
        mipLevel: 0,
        arrayLayer: 0,
        origin: { x: 0, y: 0, z: 0 }
    };

    const blitCommandEncoder = device.createCommandEncoder();
    blitCommandEncoder.copyBufferToTexture(dataCopyView, textureCopyView, textureSize);

    queue = device.getQueue();

    queue.submit([blitCommandEncoder.finish()]);

    // Bind group binding layout
    const transformBufferBindGroupLayoutBinding = {
        binding: transformBindingNum, // id[[(0)]]
        visibility: GPUShaderStageBit.VERTEX,
        type: "uniform-buffer"
    };

    const textureBindGroupLayoutBinding = {
        binding: textureBindingNum,
        visibility: GPUShaderStageBit.FRAGMENT,
        type: "sampled-texture"
    };
    textureViewBinding = {
        binding: textureBindingNum,
        resource: texture.createDefaultView()
    };

    const samplerBindGroupLayoutBinding = {
        binding: samplerBindingNum,
        visibility: GPUShaderStageBit.FRAGMENT,
        type: "sampler"
    };
    samplerBinding = {
        binding: samplerBindingNum,
        resource: device.createSampler({})
    };

    const bindGroupLayoutDescriptor = { 
        bindings: [transformBufferBindGroupLayoutBinding, textureBindGroupLayoutBinding, samplerBindGroupLayoutBinding] 
    };
    bindGroupLayout = device.createBindGroupLayout(bindGroupLayoutDescriptor);

    // Pipeline
    const depthStateDescriptor = {
        depthWriteEnabled: true,
        depthCompare: "less"
    };

    const pipelineLayoutDescriptor = { bindGroupLayouts: [bindGroupLayout] };
    const pipelineLayout = device.createPipelineLayout(pipelineLayoutDescriptor);
    const vertexStageDescriptor = {
        module: shaderModule,
        entryPoint: "vertex_main"
    };
    const fragmentStageDescriptor = {
        module: shaderModule,
        entryPoint: "fragment_main"
    };
    const colorState = {
        format: "bgra8unorm",
        alphaBlend: {
            srcFactor: "src-alpha",
            dstFactor: "one-minus-src-alpha",
            operation: "add"
        },
        colorBlend: {
            srcFactor: "src-alpha",
            dstFactor: "one-minus-src-alpha",
            operation: "add"
        },
        writeMask: GPUColorWriteBits.ALL
    };
    const pipelineDescriptor = {
        layout: pipelineLayout,

        vertexStage: vertexStageDescriptor,
        fragmentStage: fragmentStageDescriptor,

        primitiveTopology: "triangle-list",
        colorStates: [colorState],
        depthStencilState: depthStateDescriptor,
        inputState: inputStateDescriptor
    };
    pipeline = device.createRenderPipeline(pipelineDescriptor);

    let colorAttachment = {
        // attachment is acquired in render loop.
        loadOp: "clear",
        storeOp: "store",
        clearColor: { r: 0.5, g: 1.0, b: 1.0, a: 1.0 } // GPUColor
    };

    // Depth stencil texture

    // GPUExtent3D
    const depthSize = {
        width: canvas.width,
        height: canvas.height,
        depth: 1
    };

    const depthTextureDescriptor = {
        size: depthSize,
        arrayLayerCount: 1,
        mipLevelCount: 1,
        sampleCount: 1,
        dimension: "2d",
        format: "depth32float-stencil8",
        usage: GPUTextureUsage.OUTPUT_ATTACHMENT
    };

    const depthTexture = device.createTexture(depthTextureDescriptor);

    // GPURenderPassDepthStencilAttachmentDescriptor
    const depthAttachment = {
        attachment: depthTexture.createDefaultView(),
        depthLoadOp: "clear",
        depthStoreOp: "store",
        clearDepth: 1.0
    };

    renderPassDescriptor = { 
        colorAttachments: [colorAttachment],
        depthStencilAttachment: depthAttachment
    };

    render();
}

/* Transform Buffers and Bindings */
const transformSize = 4 * 16;
function updateTransformArray(array) {
    let viewMatrix = mat4.create();
    mat4.translate(viewMatrix, viewMatrix, vec3.fromValues(0, 0, -5));
    let now = Date.now() / 1000;
    mat4.rotate(viewMatrix, viewMatrix, 1, vec3.fromValues(Math.sin(now), Math.cos(now), 0));
    let modelViewProjectionMatrix = mat4.create();
    mat4.multiply(modelViewProjectionMatrix, projectionMatrix, viewMatrix);
    for (let i = 0; i < 16; i++) {
        array[i] = modelViewProjectionMatrix[i];
    }
}

const transformBufferDescriptor = {
    size: transformSize,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.MAP_WRITE
};

function createBindGroupDescriptor(transformBuffer, textureViewBinding, samplerBinding) {
    const transformBufferBinding = {
        buffer: transformBuffer,
        offset: 0,
        size: transformSize
    };
    const transformBufferBindGroupBinding = {
        binding: transformBindingNum,
        resource: transformBufferBinding
    };
    return {
        layout: bindGroupLayout,
        bindings: [transformBufferBindGroupBinding, textureViewBinding, samplerBinding]
    };
}

let mappedGroups = [];

function render() {
    if (mappedGroups.length == 0) {
        const buffer = device.createBuffer(transformBufferDescriptor);
        buffer.mapWriteAsync().then(arrayBuffer => {
            const group = device.createBindGroup(createBindGroupDescriptor(buffer, textureViewBinding, samplerBinding));

            let mappedGroup = { buffer: buffer, arrayBuffer: arrayBuffer, bindGroup: group };
            drawCommands(mappedGroup);
        });
    } else
        drawCommands(mappedGroups.shift());
}

function drawCommands(mappedGroup) {
    updateTransformArray(new Float32Array(mappedGroup.arrayBuffer));
    mappedGroup.buffer.unmap();

    const commandEncoder = device.createCommandEncoder();
    renderPassDescriptor.colorAttachments[0].attachment = swapChain.getCurrentTexture().createDefaultView();
    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
    // Encode drawing commands.
    passEncoder.setPipeline(pipeline);
    // Vertex attributes
    passEncoder.setVertexBuffers(vertexBufferIndex, [verticesBuffer], [0]);
    // Bind groups
    passEncoder.setBindGroup(bindGroupIndex, mappedGroup.bindGroup);
    passEncoder.draw(36, 1, 0, 0);
    passEncoder.endPass();

    queue.submit([commandEncoder.finish()]);

    // Ready the current buffer for update after GPU is done with it.
    mappedGroup.buffer.mapWriteAsync().then((arrayBuffer) => {
        mappedGroup.arrayBuffer = arrayBuffer;
        mappedGroups.push(mappedGroup);
    });

    requestAnimationFrame(render);
}

init();
</script>
</body>
</html>