DEV Community: Zend

Triangles On Web Ch2 Creating the Life Game

Zend — Sun, 24 Nov 2024 17:00:50 +0000

We will build a cellular automaton that simulates the game of life. The game of life is a zero-player game, meaning that its evolution is determined by its initial state, requiring no further input. One interacts with the Game of Life by creating an initial configuration and observing how it evolves.

In the last chapter, we learnt about WebGPU basics, and now, we will use that knowledge to create a simple game of life simulation. We will use the GPU to perform the calculations and render the game of life.

The Game of Life

The Game of Life is a cellular automaton devised by the British mathematician John Horton Conway in 1970.

Basically, there is a gird of any size, and each node in the grid can either be empty or alive. The game evolves in steps, and at each step, the following rules are applied to each node:

If the cell is alive, then it stays alive if it has either 2 or 3 live neighbors, otherwise it dies.
If the cell is empty, then it springs to life only in the case that it has 3 live neighbors, otherwise it remains empty.

By neighboring, we mean the 8 cells surrounding the cell.

Clean Up the Code

Before we start, let's clean up the code from the last chapter. It's just a bit refactoring.

const requestDevice = async (): Promise<[GPUAdapter, GPUDevice] | null> => {
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) {
        console.error('WebGPU not supported');
        return null;
    }
    const device = await adapter.requestDevice();
    console.log(device);
    return [adapter, device];
}

const getContext = async (device: GPUDevice): Promise<[GPUCanvasContext, GPUTextureFormat]> => {
    const canvas = document.getElementById('app') as HTMLCanvasElement;
    const context = canvas.getContext("webgpu")!;
    const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
    context.configure({
        device: device,
        format: canvasFormat,
    });
    return [context, canvasFormat];
}

const getShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "shader",
        code: `
            @vertex
            fn vertexMain(@location(0) pos: vec2f) -> @builtin(position) vec4f {
                return vec4f(pos, 0, 1);
            }
            @fragment
            fn fragmentMain() -> @location(0) vec4<f32> {
                return vec4<f32>(1.0, 1.0, 0.0, 1.0);
            }
        `
    });
}

const getBuffer = async (device: GPUDevice): Promise<[GPUBuffer, GPUVertexBufferLayout]> => {
    const vertices = [
        0.0, 0.0,
        0.0, 1.0,
        1.0, 0.0,

        1.0, 1.0,
        0.0, 1.0,
        1.0, 0.0,
    ]
    const vertexBuffer = device.createBuffer({
        size: vertices.length * 4,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
        mappedAtCreation: true,
    });
    new Float32Array(vertexBuffer.getMappedRange()).set(vertices);
    const vertexBufferLayout: GPUVertexBufferLayout = {
        arrayStride: 8,
        attributes: [{
            format: "float32x2",
            offset: 0,
            shaderLocation: 0,
        }],
    };
    vertexBuffer.unmap();
    return [vertexBuffer, vertexBufferLayout];
}


const main = async () => {
    const [_, device] = (await requestDevice())!;
    const [context, canvasFormat] = await getContext(device);
    const shaderModule = await getShaderModule(device);
    const [vertexBuffer, vertexBufferLayout] = await getBuffer(device);
    const pipeline = device.createRenderPipeline({
        label: "Cell pipeline",
        layout: "auto",
        vertex: {
            module: shaderModule,
            entryPoint: "vertexMain",
            buffers: [vertexBufferLayout]
        },
        fragment: {
            module: shaderModule,
            entryPoint: "fragmentMain",
            targets: [{
                format: canvasFormat
            }]
        }
    });
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginRenderPass({
        colorAttachments: [{
            view: context.getCurrentTexture().createView(),
            loadOp: "clear",
            clearValue: { r: 0.1, g: 0.3, b: 0.8, a: 1.0 },
            storeOp: "store",
        }]
    });
    pass.setPipeline(pipeline);
    pass.setVertexBuffer(0, vertexBuffer);
    pass.draw(6);
    pass.end();
    const commandBuffer = encoder.finish();
    device.queue.submit([commandBuffer]);
}

main()

Draw the Grid

We will start by drawing all the cells on the screen, which forms a grid.

First we generate all the points for each grid cell. A grid is a 2D array of cells, and each cell is a square. We will draw the grid as a set of triangles. Each cell will be represented by two triangles.

const grid_length = 64;
const getBuffer = async (device: GPUDevice): Promise<[GPUBuffer, GPUVertexBufferLayout]> => {
    const vertexBuffer = device.createBuffer({
        size: grid_length * grid_length * 4 * 12,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
        mappedAtCreation: true,
    });
    const vertices = new Float32Array(grid_length * grid_length * 12);
    // starts from -1 to 1
    const step = 2 / grid_length;
    const padding = step / 4;
    for(let i = 0; i < grid_length; i++) {
        for(let j = 0; j < grid_length; j++) {
            const top_left_x = -1 + i * step;
            const top_left_y = 1 - j * step;
            const bottom_right_x = top_left_x + step;
            const bottom_right_y = top_left_y - step;
            const index = (i * grid_length + j) * 12;

            vertices[index] = top_left_x + padding;
            vertices[index + 1] = top_left_y - padding;

            vertices[index + 2] = bottom_right_x - padding;
            vertices[index + 3] = top_left_y - padding;

            vertices[index + 4] = top_left_x + padding;
            vertices[index + 5] = bottom_right_y + padding;

            vertices[index + 6] = top_left_x + padding;
            vertices[index + 7] = bottom_right_y + padding;

            vertices[index + 8] = bottom_right_x - padding;
            vertices[index + 9] = top_left_y - padding;

            vertices[index + 10] = bottom_right_x - padding;
            vertices[index + 11] = bottom_right_y + padding;
        }
    }
    new Float32Array(vertexBuffer.getMappedRange()).set(vertices);
    const vertexBufferLayout: GPUVertexBufferLayout = {
        arrayStride: 8,
        attributes: [{
            format: "float32x2",
            offset: 0,
            shaderLocation: 0,
        }],
    };
    vertexBuffer.unmap();
    return [vertexBuffer, vertexBufferLayout];
}

Then change the draw pass to the correct number,

pass.draw(grid_length * grid_length * 6);

Now, you should see a grid of cells on the screen.

Handle the States

Now we need to deal with the states of each cell. Here is our first issue- how can we pass the states of each cell to the GPU? Previously, we learnt about passing data to GPU as shader parameters. But in this case, we can't do that since every cell needs access to the state of all other cells. We can't pass the state of all cells as a parameter to the shader.

The solution is a bind group. A bind group is a collection of resources that are bound to the pipeline. We can directly create an array and bind it to a buffer, then pass it to the shader.

We can simply use u32 to represent the state of each cell. We will use 0 to represent an empty cell and 1 to represent a live cell, please note that this is definitely not the most economical way to represent the state of a cell, but it's the simplest way for now.

To create a bind group, first we initialize the data for it. Here, we will use a random pattern, where for each cell, there is a 25% chance that it will be alive.

const states = new Uint32Array(grid_length * grid_length);
for(let i = 0; i < grid_length * grid_length; i++) {
    states[i] = Math.random() < 0.25 ? 1 : 0;
}

Then we create a buffer and bind group for it. The buffer for vertex shader input is vertex buffer, and for now, we just need a buffer that stores the states of each cell, so we use a storage buffer.

const statesStorageBuffer = device.createBuffer({
    size: grid_length * grid_length * 4,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
});

Previously, we used the map function to get the mapped range of the buffer, but this time, we can use device.queue.writeBuffer, which is another way to write data to the buffer. For mapped range, it is a complete map and thus suitable for constantly updating, small data, while writeBuffer is suitable for writing data once.

device.queue.writeBuffer(statesStorageBuffer, 0, states.buffer);

The second parameter is an offset, obviously, it should be a zero here.

To access the buffer in the shader, we need to create a bind group layout. A bind group layout is a description of the resources that a bind group will contain.

const bindGroup = device.createBindGroup({
  label: "bind group",
  layout: pipeline.getBindGroupLayout(0),
  entries: [{
    binding: 0,
    resource: { buffer: statesStorageBuffer },
  }],
});

Here we create a bind group with a single entry, which is a buffer that stores the states of each cell.

Typically, you need to manually create a layout. But we created pipeline with layout: "auto", so we can just use pipeline.getBindGroupLayout(0) to get the layout. It will look for the shader and create the layout for you.

Then we need to update the render pass to bind the bind group.

pass.setBindGroup(0, bindGroup);

So we create a bind group, which contains a buffer that stores the states of each cell, and bind it to the pipeline.

To access the buffer in the shader, we need to update the shader code.

const getShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "shader",
        code: `
            @group(0) @binding(0) var<storage> states: array<u32>;

            @vertex
            fn vertexMain(@location(0) pos: vec2f) -> @builtin(position) vec4f {
                return vec4f(pos, 0, 1);
            }
            @fragment
            fn fragmentMain() -> @location(0) vec4<f32> {
                return vec4<f32>(1.0, 1.0, 0.0, 1.0);
            }
        `
    });
}

Here, the var<storage> is used to declare a storage buffer, and array<u32> is used to declare an array of u32. The @group(0) @binding(0) is used to bind the buffer to the bind group.

Now we can access the buffer in the shader. However, there is another issue. How can we knew the index of the cell to be rendered?

We just can just pass the index of the cell to the vertex shader. To do so, we can add a new parameter to the vertex shader.

@group(0) @binding(0) var<storage> states: array<u32>;

@vertex
fn vertexMain(@location(0) pos: vec2f, @location(1) cell: vec2f) -> @builtin(position) vec4f {

    return vec4f(pos, 0, 1);
}

Then we need to update the vertex buffer layout.

const vertexBufferLayout: GPUVertexBufferLayout = {
    arrayStride: 12,
    attributes: [{
        format: "float32x2",
        offset: 0,
        shaderLocation: 0,
    }, {
        format: "float32x2",
        offset: 8,
        shaderLocation: 1,
    }],
};

However, a simpler way would be using vertex index. The vertex index is the index of the vertex in the vertex buffer. We can use the vertex index to calculate the cell index.

@group(0) @binding(0) var<storage> states: array<u32>;

@vertex
fn vertexMain(@location(0) pos: vec2f, @builtin(vertex_index) vertexIndex: u32) -> @builtin(position) vec4f {
    let cell_index = vertexIndex / 6;
    return vec4f(pos, 0, 1);
}

But the problem is, vertex shader can only take parameters and return position, so we can not discard certain vertices in vertex shader. We need to use the fragment shader to do that. But how we can pass the state to the fragment shader?

Actually, you can just use a struct. So long that the returned value of the vertex shader is a struct with a field of type @builtin(position) vec4f, the fragment shader can access the struct.

struct VertexOutput {
    @builtin(position) pos: vec4<f32>,
    @location(0) @interpolate(flat) cell: u32,
};

@group(0) @binding(0) var<storage> states: array<u32>;

@vertex
fn vertexMain(@location(0) pos: vec2f, @builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
    let cell_index = vertexIndex / 6;
    var output: VertexOutput;
    output.pos = vec4f(gridPos, 0, 1);
    output.cell = cell;
    return output;
}

Please note that, fragment shader is responsible for rendering every pixel, but vertex shader output only the vertex. When the vertices are passed, they are first ensemble into a primitive, then the fragment shader is called for each pixel in the primitive. This is why we need a @interpolate(flat) for the cell. This tells the pipeline that when dealing with the value of pixels that are not vertices, it should use the flatten value of the primitive- that is, the value of the vertex.

Previously, we passed no parameters to the fragment shader. Actually, it just accepts the output from the vertex shader, so if you want the position, if vertex only return position, then you can just use pos in the fragment shader.

@fragment
fn fragmentMain(@builtin(position) pos: vec4f) -> vec4<f32> {
    return vec4<f32>(1.0, 1.0, 0.0, 1.0);
}

But now, since we have a struct, we can access the field of the struct in the fragment shader.

In the fragment, you can discard some nodes based on the state. The discarded nodes will not be rendered. However, for example, if in a triangle, one vertex is discarded, the whole triangle will be discarded.

@fragment
fn fragmentMain(input: VertexOutput) -> @location(0) vec4<f32> {
    if(states[input.cell] == 0u) {
        discard;
    }
    return vec4<f32>(1.0, 1.0, 0.0, 1.0);
}

Now, you should only see part of the grid, where the cells are alive initially.

Update the States

However, our cells are not alive- we have to update the states of each cell.

To run computation on GPU, instead of vertex shader and fragment shader, we can use compute shader. A compute shader is a shader stage that is used to perform general-purpose computation on the GPU. It is used to perform calculations that are not necessarily related to rendering.

For computation shader, we need to create a new shader module.

const getComputationShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "computation shader",
        code: `
            @group(0) @binding(0) var<storage> states: array<u32>;

            @compute
            @workgroup_size(${workgroup_size}, ${workgroup_size})
            fn computeMain() {

            }
        `
    });
}

We have something unique for the compute shader- @workgroup_size. The workgroup size is the number of threads in a workgroup. The workgroup size is a 3D vector, and the typical maximum size is 256x256x64, may vary depending on the device, in my Mac, it's only capable of up to 16.

If you want to do, for example, computation over grid_length by grid_length cells, if it exceeds the workgroup size, you can make the workgroup size smaller. When the actual computation is done, it's done by a workgroup_size times workgroup_size every time. It slides over the grid. global_invocation_id is still the actual index of the cell over the whole grid.

By omitting the parameters, it is the same as setting that dimension to 1.

Then the same pipeline and pass,

const computationShaderModule = await getComputationShaderModule(device);
const computationPipeline = device.createComputePipeline({
    label: "computation pipeline",
    layout: "auto",
    compute: {
        module: computationShaderModule,
        entryPoint: "computeMain",
    }
});
const computationEncoder = device.createCommandEncoder();
const computationPass = computationEncoder.beginComputePass();

Before going further, we need to introduce a two-buffer system for updating states. Basically, since rendering is a long-lasting process, we can't update the states while rendering. So we need to use two buffers, one for the current state and one for the next state. After each frame, we swap the buffers.

Now we extract our code that is responsible for rendering, we isolate them and pass an extra parameter, the step.

const render = async (step: number) => {
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginRenderPass({
        colorAttachments: [{
            view: context.getCurrentTexture().createView(),
            loadOp: "clear",
            clearValue: { r: 0.1, g: 0.3, b: 0.6, a: 1.0 },
            storeOp: "store",
        }]
    });
    const bindGroup = device.createBindGroup({
        label: "bind group",
        layout: pipeline.getBindGroupLayout(0),
        entries: [{
            binding: 0,
            resource: { buffer: statesStorageBuffer },
        }],
    });
    pass.setPipeline(pipeline);
    pass.setBindGroup(0, bindGroup);
    pass.setVertexBuffer(0, vertexBuffer);
    pass.draw(grid_length * grid_length * 6);
    pass.end();
    const commandBuffer = encoder.finish();
    device.queue.submit([commandBuffer]);
}
render(0);

Instead of one buffer, we use two buffer to allow switching between two.

const statesStorageBuffer = [
    device.createBuffer({
        size: grid_length * grid_length * 4,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    }),
    device.createBuffer({
        size: grid_length * grid_length * 4,
        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    })
]
device.queue.writeBuffer(statesStorageBuffer[0], 0, states.buffer);

Then, when rendering, we need to bind the correct buffer. Since we will be swapping the buffers, we need to use the modulo operator to get the correct buffer.

const bindGroup = device.createBindGroup({
    label: "bind group",
    layout: pipeline.getBindGroupLayout(0),
    entries: [{
        binding: 0,
        resource: { buffer: statesStorageBuffer[step % 2] },
    }],
});

And lastly,

let step = 0;
setInterval(() => {
    render(step);
    step++;
    step = step % 2;
}, 1000);

We will see the grid of cells updating every second. The second state is all the cells are dead, since we didn't update the states.

Now let's write the actual computation shader. Now we have two buffers and will be switching between them. So after every rendering, we just use one as the previous state, and write the states to the other buffer, so that the next rendering will use the updated states.

Previously, we could only read from the buffer, in order to write to the buffer in the wsgl, we need to specify read_write in the var declaration.

@group(0) @binding(1) var<storage, read_write> states: array<u32>;

Note that we don't need to change GPU usage, since that is how CPU will interact with the buffer, where as now, buffers are handled completely by GPU except for the first write.

Then we can write the computation shader.

@group(0) @binding(0) var<storage> states: array<u32>;
@group(0) @binding(1) var<storage, read_write> next_states: array<u32>;

@compute
@workgroup_size(${grid_length}, ${grid_length})
fn computeMain(@builtin(global_invocation_id) cell: vec3u) {

}

The compute shader is executed in parallel, and the global_invocation_id is the global index of the thread. It is a vec3u. For example, for the first task executaed, it is 0, 0, 0, then increase for the next task- though please note that they are computed parallelly.

We can just use cell.x, cell.y as the index of the cell. That is, for task x, y, z, we deal with the state of the cell at x, y.

We can just use cell.x, cell.y as the index of the cell.

const getComputationShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "computation shader",
        code: `
            @group(0) @binding(0) var<storage> states: array<u32>;
            @group(0) @binding(1) var<storage, read_write> next_states: array<u32>;

            @compute
            @workgroup_size(${workgroup_size}, ${workgroup_size})
            fn computeMain(@builtin(global_invocation_id) cell: vec3u) {
                let index = cell.x * ${grid_length} + cell.y;
                var count = 0u;
                for(var i: i32 = -1; i < 2; i++) {
                    for(var j: i32 = -1; j < 2; j++) {
                        if(i == 0 && j == 0) {
                            continue;
                        }
                        let x = i32(cell.x) + i;
                        let y = i32(cell.y) + j;
                        if(x >= 0 && x < ${grid_length} && y >= 0 && y < ${grid_length}) {
                            count += states[x * ${grid_length} + y];
                        }
                    }
                }
                if(states[index] == 1u) {
                    if(count < 2u || count > 3u) {
                        next_states[index] = 0u;
                    } else {
                        next_states[index] = 1u;
                    }
                } else {
                    if(count == 3u) {
                        next_states[index] = 1u;
                    } else {
                        next_states[index] = 0u;
                    }
                }
            }
        `
    });
}

Then the rest of the code, they are the same.

const computationShaderModule = await getComputationShaderModule(device);
const computationEncoder = device.createCommandEncoder();
const computationPass = computationEncoder.beginComputePass();

However, we now want two buffers be shared for the two pipelines, since they will create their respective bind group layout. We need to make sure that the two bind groups are compatible.

const bindGroupLayout = device.createBindGroupLayout({
    entries: [{
        binding: 0,
        visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE | GPUShaderStage.VERTEX,
        buffer: {
            type: "read-only-storage"
        }
    }, {
        binding: 1,
        visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE,
        buffer: {
            type: "storage"
        }
    }]
});
const pipelineLayout = device.createPipelineLayout({
    bindGroupLayouts: [bindGroupLayout]
});
const pipeline = device.createRenderPipeline({
    label: "pipeline",
    layout: pipelineLayout,
    vertex: {
        module: shaderModule,
        entryPoint: "vertexMain",
        buffers: [vertexBufferLayout]
    },
    fragment: {
        module: shaderModule,
        entryPoint: "fragmentMain",
        targets: [{
            format: canvasFormat
        }]
    }
});
const computeShaderModule = await getComputationShaderModule(device);
const computePipeline = device.createComputePipeline({
    layout: pipelineLayout,
    compute: {
        module: computeShaderModule,
        entryPoint: "computeMain"
    }
});

Then we can create the bind group.

const bindGroup = device.createBindGroup({
    label: "bind group",
    layout: bindGroupLayout,
    entries: [{
        binding: 0,
        resource: { buffer: statesStorageBuffer[step % 2] },
    }, {
        binding: 1,
        resource: { buffer: statesStorageBuffer[(step + 1) % 2] },
    }]
});

And lastly, before the encoder finishes, we need to dispatch the computation.

const computePass = encoder.beginComputePass();
computePass.setPipeline(computePipeline);
computePass.setBindGroup(0, bindGroup);
const factor = Math.floor(grid_length / workgroup_size);
computePass.dispatchWorkgroups(factor, factor)
computePass.end();

Here for the workgroup, if the total computation is x by y by z, and the workgroup size is a by b by c, then the workgroup size is Math.ceil(x / a), Math.ceil(y / b), Math.ceil(z / c). Every time, x by y by z cells are computed.

Now, the complete code is as follows.

const grid_length = 64;
const workgroup_size = 8;

const requestDevice = async (): Promise<[GPUAdapter, GPUDevice] | null> => {
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) {
        console.error('WebGPU not supported');
        return null;
    }
    const device = await adapter.requestDevice();
    console.log(device);
    return [adapter, device];
}

const getContext = async (device: GPUDevice): Promise<[GPUCanvasContext, GPUTextureFormat]> => {
    const canvas = document.getElementById('app') as HTMLCanvasElement;
    const context = canvas.getContext("webgpu")!;
    const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
    context.configure({
        device: device,
        format: canvasFormat,
    });
    return [context, canvasFormat];
}

const getShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "shader",
        code: `
            @group(0) @binding(0) var<storage> states: array<u32>;

            struct VertexOutput {
                @builtin(position) pos: vec4f,
                @location(0) cell: u32,
            };

            @vertex
            fn vertexMain(@location(0) pos: vec2f, @builtin(vertex_index) vertexIndex: u32) -> VertexOutput {
                var output: VertexOutput;
                output.pos = vec4<f32>(pos, 0.0, 1.0);
                output.cell = vertexIndex / 6;
                return output;
            }
            @fragment
            fn fragmentMain(input: VertexOutput) -> vec4<f32> {
                if(states[input.cell] == 0u) {
                    discard;
                }
                return vec4<f32>(1.0, 1.0, 0.0, 1.0);
            }
        `
    });
}

const getVertexBuffer = async (device: GPUDevice): Promise<[GPUBuffer, GPUVertexBufferLayout]> => {
    const vertexBuffer = device.createBuffer({
        size: grid_length * grid_length * 4 * 12,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
        mappedAtCreation: true,
    });
    const vertices = new Float32Array(grid_length * grid_length * 12);
    // starts from -1 to 1
    const step = 2 / grid_length;
    const padding = step / 6;
    for(let i = 0; i < grid_length; i++) {
        for(let j = 0; j < grid_length; j++) {
            const top_left_x = -1 + i * step;
            const top_left_y = 1 - j * step;
            const bottom_right_x = top_left_x + step;
            const bottom_right_y = top_left_y - step;
            const index = (i * grid_length + j) * 12;

            vertices[index] = top_left_x + padding;
            vertices[index + 1] = top_left_y - padding;

            vertices[index + 2] = bottom_right_x - padding;
            vertices[index + 3] = top_left_y - padding;

            vertices[index + 4] = top_left_x + padding;
            vertices[index + 5] = bottom_right_y + padding;

            vertices[index + 6] = top_left_x + padding;
            vertices[index + 7] = bottom_right_y + padding;

            vertices[index + 8] = bottom_right_x - padding;
            vertices[index + 9] = top_left_y - padding;

            vertices[index + 10] = bottom_right_x - padding;
            vertices[index + 11] = bottom_right_y + padding;
        }
    }
    new Float32Array(vertexBuffer.getMappedRange()).set(vertices);
    const vertexBufferLayout: GPUVertexBufferLayout = {
        arrayStride: 8,
        attributes: [{
            format: "float32x2",
            offset: 0,
            shaderLocation: 0,
        }],
    };
    vertexBuffer.unmap();
    return [vertexBuffer, vertexBufferLayout];
}

const getComputationShaderModule = async (device: GPUDevice): Promise<GPUShaderModule> => {
    return device.createShaderModule({
        label: "computation shader",
        code: `
            @group(0) @binding(0) var<storage> states: array<u32>;
            @group(0) @binding(1) var<storage, read_write> next_states: array<u32>;

            @compute
            @workgroup_size(${workgroup_size}, ${workgroup_size})
            fn computeMain(@builtin(global_invocation_id) cell: vec3u) {
                let index = cell.x * ${grid_length} + cell.y;
                var count = 0u;
                for(var i: i32 = -1; i < 2; i++) {
                    for(var j: i32 = -1; j < 2; j++) {
                        if(i == 0 && j == 0) {
                            continue;
                        }
                        let x = i32(cell.x) + i;
                        let y = i32(cell.y) + j;
                        if(x >= 0 && x < ${grid_length} && y >= 0 && y < ${grid_length}) {
                            count += states[x * ${grid_length} + y];
                        }
                    }
                }
                if(states[index] == 1u) {
                    if(count < 2u || count > 3u) {
                        next_states[index] = 0u;
                    } else {
                        next_states[index] = 1u;
                    }
                } else {
                    if(count == 3u) {
                        next_states[index] = 1u;
                    } else {
                        next_states[index] = 0u;
                    }
                }
            }
        `
    });
}


const main = async () => {
    const [_, device] = (await requestDevice())!;
    const [context, canvasFormat] = await getContext(device);
    const shaderModule = await getShaderModule(device);
    const [vertexBuffer, vertexBufferLayout] = await getVertexBuffer(device);
    const states = new Uint32Array(grid_length * grid_length);
    for(let i = 0; i < grid_length * grid_length; i++) {
        states[i] = Math.random() > 0.25 ? 0 : 1;
    }
    const statesStorageBuffer = [
        device.createBuffer({
            size: grid_length * grid_length * 4,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
        }),
        device.createBuffer({
            size: grid_length * grid_length * 4,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
        })
    ]
    device.queue.writeBuffer(statesStorageBuffer[0], 0, states.buffer);
    const bindGroupLayout = device.createBindGroupLayout({
        entries: [{
            binding: 0,
            visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE | GPUShaderStage.VERTEX,
            buffer: {
                type: "read-only-storage"
            }
        }, {
            binding: 1,
            visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE,
            buffer: {
                type: "storage"
            }
        }]
    });
    const pipelineLayout = device.createPipelineLayout({
        bindGroupLayouts: [bindGroupLayout]
    });
    const pipeline = device.createRenderPipeline({
        label: "pipeline",
        layout: pipelineLayout,
        vertex: {
            module: shaderModule,
            entryPoint: "vertexMain",
            buffers: [vertexBufferLayout]
        },
        fragment: {
            module: shaderModule,
            entryPoint: "fragmentMain",
            targets: [{
                format: canvasFormat
            }]
        }
    });
    const computeShaderModule = await getComputationShaderModule(device);
    const computePipeline = device.createComputePipeline({
        layout: pipelineLayout,
        compute: {
            module: computeShaderModule,
            entryPoint: "computeMain"
        }
    });
    const render = async (step: number) => {
        const encoder = device.createCommandEncoder();
        const bindGroup = device.createBindGroup({
            label: "bind group",
            layout: bindGroupLayout,
            entries: [{
                binding: 0,
                resource: { buffer: statesStorageBuffer[step % 2] },
            }, {
                binding: 1,
                resource: { buffer: statesStorageBuffer[(step + 1) % 2] },
            }]
        });

        const pass = encoder.beginRenderPass({
            colorAttachments: [{
                view: context.getCurrentTexture().createView(),
                loadOp: "clear",
                clearValue: { r: 0.1, g: 0.3, b: 0.6, a: 1.0 },
                storeOp: "store",
            }]
        });
        pass.setPipeline(pipeline);
        pass.setBindGroup(0, bindGroup);
        pass.setVertexBuffer(0, vertexBuffer);
        pass.draw(grid_length * grid_length * 6);
        pass.end();

        const computePass = encoder.beginComputePass();
        computePass.setPipeline(computePipeline);
        computePass.setBindGroup(0, bindGroup);
        const factor = Math.floor(grid_length / workgroup_size);
        computePass.dispatchWorkgroups(factor, factor)
        computePass.end();

        const commandBuffer = encoder.finish();
        device.queue.submit([commandBuffer]);
    }
    let step = 0;
    setInterval(() => {
        render(step);
        step++;
        step = step % 2;
    }, 500);
}

main()

Now you should see the game of life simulation running on the screen.

Congratulations! You have successfully created a game of life simulation using WebGPU.

If you are willing to observe the game of life simulation in a more detailed way, you can increase the grid length and workgroup size. However, keep in mind that the larger the grid length and workgroup size, the more computation is required, and it may slow down the simulation. However, GPU is powerful and can handle a large amount of computation. You can build a pure CPU version and see the great difference between the two.

In the next chapter, we will learn about 3D models and how to render them using WebGPU.

Triangles On Web Ch1 Draw Something

Zend — Sun, 24 Nov 2024 16:49:29 +0000

This series introduces WebGPU, and computer graphics in general.

First let's look what we are going to build,

Life Game

3D Rendering

3D Rendering, but with lighting

Rendering 3D Model

Except for the basic knowledge of JS, no prior knowledge is needed.

The tutorial is already finished on my github, along with the source code.

WebGPU is a relatively new API for the GPU. Albeit named as WebGPU, it can actually be considered a layer on top of Vulkan, DirectX 12, and Metal, OpenGL and WebGL. It is designed to be a low-level API, and is intended to be used for high-performance applications, such as games and simulations.

In this chapter, we will draw something on the screen. The first part will refer to the Google Codelabs Tutorial. We will create a life game on the screen.

Starting Point

We will just create an empty vanilla JS project in vite with typescript enabled. Then clear all the extra codes, leaving only the main.ts.

const main = async () => {
    console.log('Hello, world!')
}

main()

Before actual coding, please check if your browser has WebGPU enabled. You can check it on WebGPU Samples.

Chrome now defaults to enabled. On Safari, you should go to developer settings, flag settings and enable WebGPU.

We also need to enable thee types for WebGPU, install @webgpu/types, and in tsc compiler options, add "types": ["@webgpu/types"].

Furthermore, we replace the <div id="app"></div> with <canvas id="app"></canvas> in the index.html.

Drawing a Triangle

There are many boilerplate codes to WebGPU, here is how it looks like.

Requesting Device

First we need access to the GPU. In WebGPU, it is done by the concept of an adapter, which is a bridge between the GPU and the browser.

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

Then we need to request a device from the adapter.

const device = await adapter.requestDevice();
console.log(device);

Configure the Canvas

We draw our triangle on the canvas. We need to get the canvas element and configure it.

const canvas = document.getElementById('app') as HTMLCanvasElement;
const context = canvas.getContext("webgpu")!;
const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
context.configure({
    device: device,
    format: canvasFormat,
});

Here, we use getContext to get relative information about the canvas. By specifying webgpu, we will get a context that is responsible for rendering with WebGPU.

CanvasFormat is actually the color mode, for example, srgb. We usually just use the preferred format.

Lastly, we configure the context with the device and the format.

Understanding GPU Rendering Pipeline

Before diving further into the engineering details, we first must understand how GPU handles rendering.

The GPU rendering pipeline is a series of steps that the GPU takes to render an image.

The application run on GPU is called a shader. The shader is a program that runs on the GPU. The shader has a special programming language that we will discuss later.

The render pipeline has the following steps,

CPU loads the data into the GPU. CPU may removed some invisible objects to save GPU resources.
CPU sets all the colors, textures, and other data that the GPU needs to render the scene.
CPU trigger a draw call to the GPU.
GPU gets the data from the CPU and starts rendering the scene.
GPU run into the geometry process, which processes the vertices of the scene.
In the geometry process, the first step is the vertex shader, which processes the vertices of the scene. It may transform the vertices, change the color of the vertices, or do other things to the vertices.
The next step is the tessellation shader, which processes the vertices of the scene. It performs subdivision of the vertices, whose purpose is to increase the detail of the scene. It also has many procedures but it's too complex to explain here.
The next step is the geometry shader, which processes the vertices of the scene. In contrast to the vertex shader, where the developer could only define how one vertex is transformed, the geometry shader can define how multiple vertices are transformed. It can also create new vertices, which can be used to create new geometry.
The last step of geometry process contains clipping, removing the undue parts that exceed the screen, and culling, removing the invisible parts that are not visible to the camera.
The next step is the rasterization process, which converts the vertices into fragments. A fragment is a pixel that is going to be rendered on the screen.
The next step is iteration of triangles, which iterates over the triangles of the scene.
The next step is the fragment shader, which processes the fragments of the scene. It may change the color of the fragments, change the texture of the fragments, or do other things to the fragments. In this part, the depth test and stencil test are also performed. Depth test means to confer each fragment with the depth value, and the fragment with the smallest depth value will be rendered. Stencil test means to confer each fragment with the stencil value, and the fragment that passes the stencil test will be rendered. The stencil value is decided by the developer.
The next step is the blending process, which blends the fragments of the scene. For example, if two fragments are overlapping, the blending process will blend the two fragments together.
The last step is the output process, which outputs the fragments to the swap chain. The swap chain is a chain of images that are used to render the scene. To put it more simply, it is a buffer that holds the image that is going to be displayed on the screen.

Depending on the primitives, the smallest unit that GPU can render, the pipeline may have different steps. Typically, we use triangles, which signals the GPU to treat every 3 group of vertices as a triangle.

Creating Render Pass

Render Pass is a step of the full GPU rendering. When a render pass is created, the GPU will start rendering the scene, and vice versa when it finishes.

To create a render pass, we need to create an encoder that is responsible for compiling the render pass to GPU codes.

const encoder = device.createCommandEncoder();

Then we create a render pass.

const pass = encoder.beginRenderPass({
  colorAttachments: [{
     view: context.getCurrentTexture().createView(),
     loadOp: "clear",
     storeOp: "store",
  }]
});

Here, we create a render pass with a color attachment. Attachment is a concept in GPU that represents the image that is going to be rendered. An image may have many aspect which the GPU need to process, and each of them is an attachment.

Here we only have one attachment, which is the color attachment. The view is the panel that the GPU will render on, here we set it to the texture of the canvas.

loadOp is the operation that the GPU will do before the render pass, clear means GPU will first clear all the previously data from the last frame, and storeOp is the operation that the GPU will do after the render pass, store means GPU will store the data to the texture.

loadOp can be load, which preserves the data from the last frame, or clear, which clears the data from the last frame. storeOp can be store, which stores the data to the texture, or discard, which discards the data.

Now, just call pass.end() to end the render pass. Now, the command is saved in the command buffer of the GPU.

To get the compiled command, use the following code,

const commandBuffer = encoder.finish();

And, finally, submit the command to the render queue of the GPU.

device.queue.submit([commandBuffer]);

Now, you should see an ugly black canvas.

Based on our stereotypical concepts about 3D, we would expect empty space to be a blue color. We can done that by setting the clear color.

const pass = encoder.beginRenderPass({
    colorAttachments: [{
        view: context.getCurrentTexture().createView(),
        loadOp: "clear",
        clearValue: { r: 0.1, g: 0.3, b: 0.8, a: 1.0 },
        storeOp: "store",
    }]
});

Drawing a Triangle Using Shader

Now, we will draw a triangle on the canvas. We will use a shader to do that. The shader language will be wgsl, WebGPU Shading Language.

Now, suppose we want to draw a triangle with the following coordinates,

(-0.5, -0.5), (0.5, -0.5), (0.0, 0.5)

As we stated before, to complete a render pipeline, we need a vertex shader and a fragment shader.

Vertex Shader

Use the following code to create shader modules.

const cellShaderModule = device.createShaderModule({
  label: "shader",
  code: `
    // Shaders
  `
});

label here is simply a name, which is meant for debugging. code is the actual shader code.

Vertex shader is a function that takes any parameter and returns the position of the vertex. However, contrary to what we might expect, the vertex shader returns a four dimensional vector, not a three dimensional vector. The fourth dimension is the w dimension, which is used for perspective division. We will discuss it later.

Now, you can simply regard a four dimensional vector (x, y, z, w) as a three dimensional vector (x / w, y / w, z / w).

However, there is another problem- how to pass the data to the shader, and how to get the data out from the shader.

To pass the data to the shader, we use the vertexBuffer, a buffer that contains the data of the vertices. We can create a buffer with the following code,

const vertexBuffer = device.createBuffer({
  size: 24,
  usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
  mappedAtCreation: true,
});

Here we create a buffer with a size of 24 bytes, 6 floats, which is the size of the vertices.

usage is the usage of the buffer, which is VERTEX for vertex data. GPUBufferUsage.COPY_DST means this buffer is valid as a copy destination. For all buffer whose data are written by the CPU, we need to set this flag.

The map here means to map the buffer to the CPU, which means the CPU can read and write the buffer. The unmap means to unmap the buffer, which means the CPU can no longer read and write the buffer, and thus the content is available to the GPU.

Now, we can write the data to the buffer.

new Float32Array(vertexBuffer.getMappedRange()).set([
  -0.5, -0.5,
  0.5, -0.5,
  0.0, 0.5,
]);
vertexBuffer.unmap();

Here, we map the buffer to the CPU, and write the data to the buffer. Then we unmap the buffer.

vertexBuffer.getMappedRange() will return the range of the buffer that is mapped to the CPU. We can use it to write the data to the buffer.

However, these are just raw data, and the GPU doesn't know how to interpret them. We need to define the layout of the buffer.

const vertexBufferLayout: GPUVertexBufferLayout = {
    arrayStride: 8,
    attributes: [{
        format: "float32x2",
        offset: 0,
        shaderLocation: 0,
    }],
};

Here, arrayStride is the number of bytes the GPU needs to skip forward in the buffer when it's looking for the next input. For example, if the arrayStride is 8, the GPU will skip 8 bytes to get the next input.

Since here, we use float32x2, the stride is 8 bytes, 4 bytes for each float, and 2 floats for each vertex.

Now we can write the vertex shader.

const shaderModule = device.createShaderModule({
  label: "shader",
  code: `
    @vertex
    fn vertexMain(@location(0) pos: vec2f) -> @builtin(position) vec4f {
        return vec4f(pos, 0, 1);
    }
  `
});

Here, @vertex means this is a vertex shader. @location(0) means the location of the attribute, which is 0, as previously defined. Please note that in the shader language, you are dealing with the layout of the buffer, so whenever you pass a value, you need to pass either a struct, whose fields had defined @location, or just a value with @location.

vec2f is a two dimensional float vector, and vec4f is a four dimensional float vector. Since vertex shader is required to return a vec4f position, we need to annotate that with @builtin(position).

Fragment Shader

Fragment shader, similarly, is something that takes the interpolated vertex output and output the attachments, color in this case. The interpolated means that although only certain pixel on the vertices have decided value, for every other pixel, the values are interpolated, either linear, averaged, or other means. The color of fragment is a four dimensional vector, which is the color of the fragment, respectively red, green, blue, and alpha.

Please note that the color is in the range of 0 to 1, not 0 to 255. In addition that, fragment shader defines the color of every vertex, not the color of the triangle. The color of the triangle is determined by the color of the vertices, by interpolation.

Since we currently does not bother to control the color of the fragment, we can simply return a constant color.

const shaderModule = device.createShaderModule({
  label: "shader",
  code: `
    @vertex
    fn vertexMain(@location(0) pos: vec2f) -> @builtin(position) vec4f {
        return vec4f(pos, 0, 1);
    }
    @fragment
    fn fragmentMain() -> vec4<f32> {
        return vec4<f32>(1.0, 1.0, 0.0, 1.0);
    }
  `
});

Render Pipeline

Then we define the customized render pipeline by replacing the vertex and fragment shader.

const pipeline = device.createRenderPipeline({
    label: "pipeline",
    layout: "auto",
    vertex: {
        module: shaderModule,
        entryPoint: "vertexMain",
        buffers: [vertexBufferLayout]
    },
    fragment: {
        module: shaderModule,
        entryPoint: "fragmentMain",
        targets: [{
            format: canvasFormat
        }]
    }
});

Note that in fragment shader, we need to specify the format of the target, which is the format of the canvas.

Draw Call

Before render pass ends, we add the draw call.

pass.setPipeline(pipeline);
pass.setVertexBuffer(0, vertexBuffer);
pass.draw(3);

Here, in setVertexBuffer, the first parameter is the index of the buffer, in the pipeline definition field buffers, and the second parameter is the buffer itself.

When calling draw, the parameter is the number of vertices to draw. Since we have 3 vertices, we draw 3.

Now, you should see a yellow triangle on the canvas.

Draw Life Game Cells

Now we tweak our codes a bit- since we want to build a life game, so we need to draw squares instead of triangles.

A square is actually two triangles, so we need to draw 6 vertices. The changes here are simple and you don't need a detailed explanation.

const main = async () => {
    const adapter = await navigator.gpu.requestAdapter();
    if (!adapter) {
        console.error('WebGPU not supported');
        return;
    }
    const device = await adapter.requestDevice();
    console.log(device);
    const canvas = document.getElementById('app') as HTMLCanvasElement;
    const context = canvas.getContext("webgpu")!;
    const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
    context.configure({
        device: device,
        format: canvasFormat,
    });
    const vertices = [
        0.0, 0.0,
        0.0, 1.0,
        1.0, 0.0,

        1.0, 1.0,
        0.0, 1.0,
        1.0, 0.0,
    ]
    const vertexBuffer = device.createBuffer({
        size: vertices.length * 4,
        usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
        mappedAtCreation: true,
    });
    new Float32Array(vertexBuffer.getMappedRange()).set(vertices);
    const vertexBufferLayout: GPUVertexBufferLayout = {
        arrayStride: 8,
        attributes: [{
            format: "float32x2",
            offset: 0,
            shaderLocation: 0,
        }],
    };
    const shaderModule = device.createShaderModule({
        label: "shader",
        code: `
            @vertex
            fn vertexMain(@location(0) pos: vec2f) -> @builtin(position) vec4f {
                return vec4f(pos, 0, 1);
            }
            @fragment
            fn fragmentMain() -> vec4<f32> {
                return vec4<f32>(1.0, 1.0, 0.0, 1.0);
            }
        `
    });
    vertexBuffer.unmap();
    const pipeline = device.createRenderPipeline({
        label: "Cell pipeline",
        layout: "auto",
        vertex: {
            module: shaderModule,
            entryPoint: "vertexMain",
            buffers: [vertexBufferLayout]
        },
        fragment: {
            module: shaderModule,
            entryPoint: "fragmentMain",
            targets: [{
                format: canvasFormat
            }]
        }
    });
    const encoder = device.createCommandEncoder();
    const pass = encoder.beginRenderPass({
        colorAttachments: [{
            view: context.getCurrentTexture().createView(),
            loadOp: "clear",
            clearValue: { r: 0.1, g: 0.3, b: 0.8, a: 1.0 },
            storeOp: "store",
        }]
    });
    pass.setPipeline(pipeline);
    pass.setVertexBuffer(0, vertexBuffer);
    pass.draw(vertices.length / 2);
    pass.end();
    const commandBuffer = encoder.finish();
    device.queue.submit([commandBuffer]);
}

Now, you should see a yellow square on the canvas.

Coordinate System

We didn't discuss the coordinate system of the GPU. It is, well, rather simple. The actual coordinate system of the GPU is a right-handed coordinate system, which means the x-axis points to the right, the y-axis points up, and the z-axis points out of the screen.

The range of the coordinate system is from -1 to 1. The origin is at the center of the screen. z-axis is from 0 to 1, 0 is the near plane, and 1 is the far plane. However, z-axis is for depth. When you do 3D rendering, you can not just use z-axis to determine the position of the object, you need to use the perspective division. This is called the NDC, normalized device coordinate.

For example, if you want to draw a square at the top left corner of the screen, the vertices are (-1, 1), (-1, 0), (0, 1), (0, 0), though you need to use two triangles to draw it.

Build a Tiny React Ch04 Hooks

Zend — Sun, 20 Oct 2024 09:35:02 +0000

Prev Part

This tutorial is based on this tutorial, but with JSX, typescript and an easier approach to implement. You can checkout the notes and code on my GitHub repo.

Okay, before diving into hooks, we need a bit wrap up on the last chapter- still something to fix, but the last chapter was too much, so, well, here it is.

Fixing the last chapter

Here are some minor things- not entirely bugs, but better to fix them.

vDOM Comparison

In javascript, two functions are equal only if they are the same, stays unequal even if they have the same procedure, that is,

const a = () => 1;
const b = () => 1;
a === b; // false

So, when it comes to vDOM Comparison, we should skip the function comparison. Here is the fix,

for (let i = 0; i < aKeys.length; i++) {
    const key = aKeys[i]
    if (key === 'key') continue
    if (aProps[key] instanceof Function && bProps[key] instanceof Function) continue
    if (aProps[key] !== bProps[key]) return false
}
for (let i = 0; i < bKeys.length; i++) {
    const key = bKeys[i]
    if (key === 'key') continue
    if (aProps[key] instanceof Function && bProps[key] instanceof Function) continue
    if (aProps[key] !== bProps[key]) return false
}

Handling CSS

Style should be treated as a special property that is attributed to element with .style property. Here is the fix,

export type VDomAttributes = { 
    key?: string | number
    style?: object
    [_: string]: unknown | undefined
}

export function createDom(vDom: VDomNode): HTMLElement | Text {
    if (isElement(vDom)) {
        const element = document.createElement(vDom.tag)
        Object.entries(vDom.props ?? {}).forEach(([name, value]) => {
            if (value === undefined) return
            if (name === 'key') return
            if (name === 'style') {
                Object.entries(value as Record<string, unknown>).forEach(([styleName, styleValue]) => {
                    element.style[styleName as any] = styleValue as any
                })
                return
            }
            if (name.startsWith('on') && value instanceof Function) {
                element.addEventListener(name.slice(2).toLowerCase(), value as EventListener)
            } else {
                element.setAttribute(name, value?.toString() ?? "")
            }
        })
        return element
    } else {
        return document.createTextNode(vDom)
    }
}

Now these side fixes are done, let's move on to the main topic of this chapter- Hooks.

Capsuling vDOM Creation

We previously explicitly called render(vDom, app!), which requires creating vDOM by user, here is a better way to do it.

import { mount, useState, type FuncComponent } from "./runtime";
import { createElement, fragment, VDomAttributes, VDomNode } from "./v-dom";

const App: FuncComponent = (props: VDomAttributes, __: VDomNode[]) => {
    const [cnt, setCnt] = useState(0)
    return <div>
        <button onClick={() => setCnt(cnt() + 1)}>Click me</button>
        <p>Count: {cnt()}</p>
    </div>
}
const app = document.getElementById('app')
mount(App, {}, [], app!)

let reRender: () => void = () => {}
export function mount(app: FuncComponent, props: VDomAttributes, children: VDomNode[], parent: HTMLElement) {
    reRender = () => {
        const vDom = app(props, children) as unknown as VDomNode
        render(vDom, parent)
    }
    reRender()
}

Looks better more or less. Now let's move on to the main topic of this chapter- Hooks.

useState

Okay, let's get to the hook. The first hook we are going to implement is useState. It is a hook that allows us to manage the state of a component. We can have the following signature for useState,

Note that our implementation is slightly different from the original React. We are going to return a getter and a setter function, instead of returning the state directly.

function useState<T>(initialValue: T): [() => T, (newValue: T) => void] {
    // implementation
}

So where will we hook the value at? If we just hide it within the closure itself, the value will be lost when the component is re-rendered. If you insist on doing so, you need to access the space of the outer function, which is not possible in javascript.

So our way is to store it, you guessed it, in fiber. So, let's add a field to the fiber.

interface Fiber {
    parent: Fiber | null
    sibling: Fiber | null
    child: Fiber | null
    vDom: VDomNode
    dom: HTMLElement | Text | null
    alternate: Fiber | null
    committed: boolean
    hooks?: {
        state: unknown[]
    },
    hookIndex?: {
        state: number
    }
}

And we only mount hooks to the root fiber, so we can add the following line to the mount function.

export function render(vDom: VDomNode, parent: HTMLElement) {
    wip = {
        parent: null,
        sibling: null,
        child: null,
        vDom: vDom,
        dom: null,
        committed: false,
        alternate: oldFiber,
        hooks: oldFiber?.hooks ?? {
            state: []
        },
        hookIndex: {
            state: 0
        }
    }
    wipParent = parent
    nextUnitOfWork = wip
}

Hook index will come into use later. Now, hook index resets every time the component is re-rendered, but old hooks are carried over.

Please note that, we rendering component vDOM, only old fiber is accessible, so we can only manipulate that variable. However, it is null at the very beginning, so let's set up a dummy.

export function mount(app: FuncComponent, props: VDomAttributes, children: VDomNode[], parent: HTMLElement) {
    reRender = () => {
        if(!oldFiber) {
            oldFiber = {
                parent: null,
                sibling: null,
                child: null,
                vDom: '' as any,
                dom: null,
                committed: false,
                alternate: null,
                hooks: {
                    state: []
                },
                hookIndex: {
                    state: 0
                }
            }
        }
        const vDom = app(props, children) as unknown as VDomNode
        render(vDom, parent)
        console.log('rerender')
        console.log(vDom)
    }
    reRender()
}

Now we will have big brain time- since the order of each hook call is fixed (you can not use hooks in a loop or a condition, basic React rule, you know why it is now), so we can use safely use hookIndex to access the hook.

export function useState<T>(initialState: T): [T, (newState: T) => void] {
    const hookIndex = oldFiber!.hookIndex!.state
    oldFiber!.hookIndex!.state++
    oldFiber!.hooks!.state[hookIndex] = oldFiber!.hooks!.state[hookIndex] ?? initialState
    console.log(oldFiber!.hooks!.state)
    const state = oldFiber!.hooks!.state[hookIndex]
    const setState = (newState: T) => {
        oldFiber!.hooks!.state[hookIndex] = newState
        reRender!()
    }
    return [state as T, setState]
}

Well, let's try,

import { mount, useState, type FuncComponent } from "./runtime";
import { createElement, fragment, VDomAttributes, VDomNode } from "./v-dom";

const App: FuncComponent = (props: VDomAttributes, __: VDomNode[]) => {
    const [cnt, setCnt] = useState(0)
    return <div>
        <button onClick={() => setCnt(cnt + 1)}>Click me</button>
        <p>Count: {cnt}</p>
    </div>
}
const app = document.getElementById('app')
mount(App, {}, [], app!)

It kind of works- the count increased from zero to one, but it does not increase further.

Well... strange right? Let's see what's going on, debug time.

const App: FuncComponent = (props: VDomAttributes, __: VDomNode[]) => {
    const [cnt, setCnt] = useState(0)
    return <div>
        <button onClick={() => {
            console.log(cnt + 1)
            console.log("FROM COMPONENT")
            setCnt(cnt + 1)
        }}>Click me</button>
        <p>Count: {cnt}</p>
    </div>
}
const app = document.getElementById('app')
mount(App, {}, [], app!)

You will see that, it always log 1. But the web page tells us that it is 1, so it should be 2. What's going on?

For native types, javascript passes by value, so the value is copied, not referenced. In React class component, it requires you to have a state object to address the issues. In functional component, the React addresses with, closure. But if we were to use the latter, it requires a big change in our code. So a simple way to get pass is, using function to get the state, so that the function will always return the latest state.

export function useState<T>(initialState: T): [() => T, (newState: T) => void] {
    const hookIndex = oldFiber!.hookIndex!.state
    oldFiber!.hookIndex!.state++
    oldFiber!.hooks!.state[hookIndex] = oldFiber!.hooks!.state[hookIndex] ?? initialState
    console.log(oldFiber!.hooks!.state)
    const state = () => oldFiber!.hooks!.state[hookIndex] as T
    const setState = (newState: T) => {
        oldFiber!.hooks!.state[hookIndex] = newState
        reRender!()
    }
    return [state, setState]
}

And now, here we got it! It works! We created the useState hook for our tiny React.

Summary

Okay, you may believe that this chapter is too short- hooks are important to react, so why did we only implement useState?

First, many hooks are just variations of useState. Such hook is irrelevant to the component it is called, for example, useMemo. Such things are just trivial works and we got no time to waste.

But, second, the most important reason, is that, for hooks like useEffect, under our current, root-update based frame, they are nearly impossible to do. When a fiber unmount, you can not signal, because we only fetch the global vDOM and update the whole vDOM, whereas, in real React, it's not like so.

In real React, functional components are updated by the parent component, so the parent component can signal the child component to unmount. But in our case, we only update the root component, so we can not signal the child component to unmount.

However, the current small project have basically demonstrated how react works, and I hope you find it helpful in gaining a better understanding of the react framework.

Build a Tiny React Ch3 Updating vDOM

Zend — Sun, 20 Oct 2024 09:32:10 +0000

Next Part Prev Part
This tutorial is based on this tutorial, but with JSX, typescript and an easier approach to implement. You can checkout the notes and code on my GitHub repo.

Now let's talk about the reactivity.

Save the Old Fiber

We need to save the old fiber so that we can compare it with the new fiber. We can do this by adding a field to the fiber. We also need a committed field- which will be useful later.

export interface Fiber {
  type: string
  props: VDomAttributes
  parent: Fiber | null
  child: Fiber | null
  sibling: Fiber | null
  dom: HTMLElement | Text | null
  alternate: Fiber | null
  committed: boolean
}

Then we set the committed state here,

function commit() {
    function commitChildren(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent.dom.appendChild(fiber.dom)
            fiber.committed = true
        }
        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }
        commitChildren(fiber.child)
        commitChildren(fiber.sibling)
    }
    commitChildren(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    wip = null
}

We also need to save the old fiber tree.

let oldFiber: Fiber | null = null

function commit() {
    function commitChildren(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent.dom.appendChild(fiber.dom)
            fiber.committed = true
        }
        commitChildren(fiber.child)
        commitChildren(fiber.sibling)
    }
    commitChildren(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

Now, we need to compare the old fiber with the new fiber during iteration. This is called the reconciliation process.

Reconciliation

We need to compare the old fiber with the new fiber. We first put the old fiber in the initial work.

export function render(vDom: VDomNode, parent: HTMLElement) {
    wip = {
        parent: null,
        sibling: null,
        child: null,
        vDom: vDom,
        dom: null,
        committed: false,
        alternate: oldFiber,
    }
    wipParent = parent
    nextUnitOfWork = wip
}

Then we separate the creation of the new fiber into a new function.

function reconcile(fiber: Fiber, isFragment: boolean) {
    if (isElement(fiber.vDom)) {
        const elements = fiber.vDom.children ?? []
        let index = 0
        let prevSibling = null

        while (index < elements.length) {
            const element = elements[index]
            const newFiber: Fiber = {
                parent: isFragment ? fiber.parent : fiber,
                dom: null,
                sibling: null,
                child: null,
                vDom: element,
                committed: false,
                alternate: null,
            }

            if (index === 0) {
                fiber.child = newFiber
            } else {
                prevSibling!.sibling = newFiber
            }
            prevSibling = newFiber
            index++
        }
    }
}

function performUnitOfWork(nextUnitOfWork: Fiber | null): Fiber | null {
    if(!nextUnitOfWork) {
        return null
    }
    const fiber = nextUnitOfWork
    const isFragment = isElement(fiber.vDom) && fiber.vDom.tag === '' && fiber.vDom.kind === 'fragment'

    if(!fiber.dom && !isFragment) {
        fiber.dom = createDom(fiber.vDom)
    }

    reconcile(fiber, isFragment)

    if (fiber.child) {
        return fiber.child
    }
    let nextFiber: Fiber | null = fiber
    while (nextFiber) {
        if (nextFiber.sibling) {
          return nextFiber.sibling
        }
        nextFiber = nextFiber.parent
    }
    return null
}

However, we need to mount the old fiber onto the new one.

function reconcile(fiber: Fiber, isFragment: boolean) {
    if (isElement(fiber.vDom)) {
        const elements = fiber.vDom.children ?? []
        let index = 0
        let prevSibling = null

        let currentOldFiber = fiber.alternate?.child ?? null
        while (index < elements.length) {
            const element = elements[index]
            const newFiber: Fiber = {
                parent: isFragment ? fiber.parent : fiber,
                dom: null,
                sibling: null,
                child: null,
                vDom: element,
                committed: false,
                alternate: currentOldFiber,
            }

            if (index === 0) {
                fiber.child = newFiber
            } else {
                prevSibling!.sibling = newFiber
            }
            prevSibling = newFiber
            currentOldFiber = currentOldFiber?.sibling ?? null
            index++
        }
    }
}

Now we have the old fiber mounted to the new fiber. But we don't have anything to trigger the re-rendering- for now, we manually trigger it by adding a button. Since we don't yet have state yet, we use props for mutating the vDOM.

import { render } from "./runtime";
import { createElement, fragment, VDomAttributes, VDomNode } from "./v-dom";

type FuncComponent = (props: VDomAttributes, children: VDomNode[]) => JSX.Element

const App: FuncComponent = (props: VDomAttributes, __: VDomNode[]) => {
    return <div>
        <>
            <h1>H1</h1>
            <h2>{props["example"]?.toString()}</h2>
            {
                props["show"] ? <p>show</p> : <></>
            }
            <h1>H1</h1>
        </>
    </div>
}
const app = document.getElementById('app')

const renderButton = document.createElement('button')
renderButton.textContent = 'Render'
let cnt = 0
renderButton.addEventListener('click', () => {
    const vDom: VDomNode = App({
        "example": (new Date()).toString(),
        "show": cnt % 2 === 0
    }, []) as unknown as VDomNode
    cnt++
    render(vDom, app!)
})
document.body.appendChild(renderButton)

Now if you click the renderButton, the rendered result will repeat once, since, well, all our current logic is simply putting the rendered vDOM into the document.

If you add a console.log in the commit function, you can see the alternate fiber being printed out.

Now we need to define how we handle the old fiber and the new fiber, and mutate the DOM based on the information. The rules is as follows.

For each new fiber,

If there was an old fiber, we compare the content of the old fiber with the new fiber, if they are different, we replace the old DOM node with the new DOM node, or else we copy the old DOM node to the new DOM node. Please note that, by two vDOM being equal, we mean their tags and all properties are equal. Their children can be different.
If there has no old fiber, we create a new DOM node and append it to the parent.
If, for the new fiber, it doesn't have a child or a sibling, but its old fiber has a child or a sibling, we recursively remove the old child or sibling.

Kind of confused? Well, I'll just show the code. We first delete the old DOM creation. Then apply the rules above.

The first rule, if there is an old fiber, we compare the content of the old fiber with the new fiber. If they are different, we replace the old DOM node with the new DOM node, or else we copy the old DOM node to the new DOM node.

export function vDOMEquals(a: VDomNode, b: VDomNode): boolean {
    if (isString(a) && isString(b)) {
        return a === b
    } else if (isElement(a) && isElement(b)) {
        let ret = a.tag === b.tag && a.key === b.key
        if (!ret) return false
        if (a.props && b.props) {
            const aProps = a.props
            const bProps = b.props
            const aKeys = Object.keys(aProps)
            const bKeys = Object.keys(bProps)
            if (aKeys.length !== bKeys.length) return false
            for (let i = 0; i < aKeys.length; i++) {
                const key = aKeys[i]
                if (key === 'key') continue
                if (aProps[key] !== bProps[key]) return false
            }
            for (let i = 0; i < bKeys.length; i++) {
                const key = bKeys[i]
                if (key === 'key') continue
                if (aProps[key] !== bProps[key]) return false
            }
            return true
        } else {
            return a.props === b.props
        }
    } else {
        return false
    }
}

Then I made some small refactor,

function buildDom(fiber: Fiber, fiberIsFragment: boolean) {
    if(fiber.dom) return
    if(fiberIsFragment) return
    fiber.dom = createDom(fiber.vDom)
}

function performUnitOfWork(nextUnitOfWork: Fiber | null): Fiber | null {
    if(!nextUnitOfWork) {
        return null
    }
    const fiber = nextUnitOfWork
    const fiberIsFragment = isFragment(fiber.vDom)

    reconcile(fiber)

    buildDom(fiber, fiberIsFragment);

    if (fiber.child) {
        return fiber.child
    }
    let nextFiber: Fiber | null = fiber
    while (nextFiber) {
        if (nextFiber.sibling) {
          return nextFiber.sibling
        }
        nextFiber = nextFiber.parent
    }
    return null
}

Now, when it comes to commit, we have an extra alternative field to compare the old fiber with the new fiber.

This is the original commit function,


function commit() {
    function commitChildren(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
            fiber.committed = true
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }

        commitChildren(fiber.child)
        commitChildren(fiber.sibling)
    }
    commitChildren(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

We will change the name a bit. The old name is just wrong (I'm sorry for that).

function commit() {
    function commitToParent(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
            fiber.committed = true
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }

        commitToParent(fiber.child)
        commitToParent(fiber.sibling)
    }
    commitToParent(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

Appending, Copying and Replacing

So what should we do? Our old logic is only appending, so we extract that,

function commit() {
    function append(fiber: Fiber) {
        if(!fiber.dom) return

        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
            fiber.committed = true
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }
    }

    function commitToParent(fiber: Fiber | null) {
        if(!fiber) {
            return
        }

        append(fiber)

        commitToParent(fiber.child)
        commitToParent(fiber.sibling)
    }
    commitToParent(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

We need to delay the construction of the DOM until the commit phase, to provide more flexibility.

function commit() {
    function append(fiber: Fiber) {
        if(isFragment(fiber.vDom)) return;
        fiber.dom = createDom(fiber.vDom)

        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
            fiber.committed = true
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }
    }

    function commitToParent(fiber: Fiber | null) {
        if(!fiber) {
            return
        }

        append(fiber)

        commitToParent(fiber.child)
        commitToParent(fiber.sibling)
    }
    commitToParent(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

function reconcile(fiber: Fiber) {
    if (isElement(fiber.vDom)) {
        const elements = fiber.vDom.children ?? []
        let index = 0
        let prevSibling = null

        let currentOldFiber = fiber.alternate?.child ?? null
        while (index < elements.length) {
            const element = elements[index]
            const newFiber: Fiber = {
                parent: fiber,
                dom: null,
                sibling: null,
                child: null,
                vDom: element,
                committed: false,
                alternate: currentOldFiber,
            }

            if (index === 0) {
                fiber.child = newFiber
            } else {
                prevSibling!.sibling = newFiber
            }
            prevSibling = newFiber
            currentOldFiber = currentOldFiber?.sibling ?? null
            index++
        }
    }
}

Following the first and second rule, we refactor them into the following code,


function commit() {
    function findNonFragmentParent(fiber: Fiber): Fiber | null {
        let parent = fiber.parent
        while(parent && isFragment(parent.vDom)) {
            parent = parent.parent!
        }
        return parent
    }

    function append(fiber: Fiber) {
        if(isFragment(fiber.vDom)) return;

        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
            fiber.committed = true
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom) && !fiber.committed) {
            const parent = findNonFragmentParent(fiber);
            parent?.dom?.appendChild(fiber.dom!)
            fiber.committed = true
        }
    }

    function commitToParent(fiber: Fiber | null) {
        if(!fiber) {
            return
        }

        if(!fiber.alternate && !isFragment(fiber.vDom)) {
            fiber.dom = createDom(fiber.vDom)
            append(fiber)
        }
        if(fiber.alternate && vDOMEquals(fiber.vDom, fiber.alternate.vDom)) {
            fiber.dom = fiber.alternate.dom
            append(fiber)
        }
        if(fiber.alternate && !vDOMEquals(fiber.vDom, fiber.alternate.vDom) && !isFragment(fiber.vDom)) {
            fiber.dom = createDom(fiber.vDom)
            fiber.alternate.dom?.replaceWith(fiber.dom)
            append(fiber)
        }

        commitToParent(fiber.child)
        commitToParent(fiber.sibling)
    }
    commitToParent(wip)
    wipParent?.appendChild(wip!.dom!)
    wip!.committed = true
    oldFiber = wip
    wip = null
}

Please always keep in mind that in javascript, all values are references. If we have fiber.dom = fiber.alternate.dom, then fiber.dom and fiber.alternate.dom will point to the same object. If we change fiber.dom, fiber.alternate.dom will also change, and vice versa. That's why when replacing, we simply used fiber.alternate.dom?.replaceWith(fiber.dom). This will replace the old DOM with the new DOM. While previous parents, if copied, have the fiber.alternate.dom for their DOM, their DOM will also be replaced.

However, we hadn't handled deletion yet.

Some Mishaps

Okay, previous code contains some bugs that I spotted as I am writing more complex jsx, so, before implementing the deletion, let's fix them.

Previously there was a bug- we can not pass list to props, let's use this chance to fix it.

export type VDomAttributes = { 
    key?: string | number
    [_: string]: unknown | undefined
}

Then you just fix the type things- only one error for me, so, do it yourself please.

However, if we have the following code,

import { render } from "./runtime";
import { createElement, fragment, VDomAttributes, VDomNode } from "./v-dom";

type FuncComponent = (props: VDomAttributes, children: VDomNode[]) => JSX.Element

const App: FuncComponent = (props: VDomAttributes, __: VDomNode[]) => {
    return <div>
            {
                props["show"] ? <h2>show</h2> : <></>
            }
            <ul>
                {
                    (props["list"] as Array<number>).map((i) => {
                        return (<ul>{i}</ul>)
                    })
                }
            </ul>
    </div>
}
const app = document.getElementById('app')

const renderButton = document.createElement('button')
renderButton.textContent = 'Render'
let cnt = 0
renderButton.addEventListener('click', () => {
    const vDom: VDomNode = App({
        "example": (new Date()).toString(),
        "show": cnt % 2 === 0,
        "list": Array.from({ length: 10 - cnt }, (_, i) => i)
    }, []) as unknown as VDomNode
    cnt++
    console.log(vDom)
    render(vDom, app!)
})
document.body.appendChild(renderButton)

Our thing broke again...

Okay, this is because children can be nested arrays in the above case, we need to flat them.

export function createElement(tag: string | Function, props: VDomAttributes, ...children: Array<VDomNode | Array<VDomNode>>): VDomElement {
    if (tag instanceof Function) {
        return tag(props, children)
    }
    return {
        kind: tag === '' ? 'fragment' : 'element',
        tag,
        children: children.flatMap((child) => Array.isArray(child) ? child : [child]),
        props: props ?? {},
        key: props?.key ?? undefined
    }
}

But that's not enough, ugh, our createDom only recognize either string or element, not integer, so, we need toString the numbers.

(props["list"] as Array<number>).map((i) => {
    return (<ul>{i.toString()}</ul>)
})

Okay, things work now- kind of.

If you hit the render button, the list is updated- but the old element still remains. We need to delete the old element.

Remove

We restate the rule here- for any new fiber, if it does not have a child or a sibling, but its old fiber has a child or a sibling, we recursively remove the old child or sibling.

function recursiveRemoveRelated(fiber: Fiber | null) {
    if(!fiber) {
        return
    }
    if(fiber.sibling) {
        recursiveRemoveRelated(fiber.sibling)
        fiber.sibling.dom?.remove()
    }
    if(fiber.child) {
        recursiveRemoveRelated(fiber.child)
        fiber.child.dom?.remove()
    }
}
if(fiber.alternate?.child && !fiber.child) {
    recursiveRemoveRelated(fiber.alternate.child)
    fiber.alternate.child.dom?.remove()
}
if(fiber.alternate?.sibling && !fiber.sibling) {
    recursiveRemoveRelated(fiber.alternate.sibling)
    fiber.alternate.sibling.dom?.remove()
}

If you don't do recursive remove, some old elements will dangle when you have multiple things requires deletion. You can change to,

renderButton.addEventListener('click', () => {
    const vDom: VDomNode = App({
        "example": (new Date()).toString(),
        "show": cnt % 2 === 0,
        "list": Array.from({ length: 10 - cnt * 2 }, (_, i) => i)
    }, []) as unknown as VDomNode
    cnt++
    console.log(vDom)
    render(vDom, app!)
})

For reference.

Summary

This is a hard chapter- but pretty traditional coding, to be honest. However, up to now, you have understand how React works from bottom to top.

Actually, things can already work now- we can manually trigger a re-render whenever we change the props. However, such frustrating manual work is not what we want. We want the reactivity to be automatic. So, we will talk about hooks in the next chapter.

Build a Tiny React Ch2 Rendering vDOM

Zend — Sun, 20 Oct 2024 09:29:48 +0000

Prev Part Next Part

This tutorial is based on this tutorial, but with JSX, typescript and an easier approach to implement. You can checkout the notes and code on my GitHub repo.

This part we will render the vDOM to the actual DOM. In addition, we will also introduce fiber tree, which a core structure in React.

Rendering vDOM

Rendering vDOM is, simple- too simple. You need to know the following web native APIs.

document.createElement(tagName: string): HTMLElement Creates an actual DOM element.
document.createTextNode(text: string): Text Creates a text node.
.appendChild(child: Node): void Appends a child node to the parent node. A method on HTMLElement
.removeChild(child: Node): void Removes a child node from the parent node. A method on HTMLElement
.replaceChild(newChild: Node, oldChild: Node): void Replaces a child node with a new child node. A method on HTMLElement
.replaceWith(...nodes: Node[]): void Replaces a node with new nodes. A method on Node
.remove(): void Removes a node from the document. A method on Node
.insertBefore(newChild: Node, refChild: Node): void Inserts a new child node before a reference child node. A method on HTMLElement
.setAttribute(name: string, value: string): void Sets an attribute on an element. A method on HTMLElement.
.removeAttribute(name: string): void Removes an attribute from an element. A method on HTMLElement.
.addEventListener(type: string, listener: Function): void Adds an event listener to an element. A method on HTMLElement.
.removeEventListener(type: string, listener: Function): void Removes an event listener from an element. A method on HTMLElement.
.dispatchEvent(event: Event): void Dispatches an event on an element. A method on HTMLElement.

Woa, a bit too much, right? But all you need to do is mirroring the creation of vDOM to the actual DOM. Here is a simple example.

function render(vDom: VDomNode, parent: HTMLElement) {
    if (typeof vDom === 'string') {
        parent.appendChild(document.createTextNode(vDom))
    } else if (vDom.kind === 'element') {
        const element = document.createElement(vDom.tag)
        for (const [key, value] of Object.entries(vDom.props ?? {})) {
            if (key === 'key') continue
            if (key.startsWith('on')) {
                element.addEventListener(key.slice(2).toLowerCase(), value as EventListener)
            } else {
                element.setAttribute(key, value as string)
            }
        }
        for (const child of vDom.children ?? []) {
            render(child, element)
        }
        parent.appendChild(element)
    } else {
        for (const child of vDom.children ?? []) {
            render(child, parent)
        }
    }
}

We registered properties starting with on as event listeners, this is a common practice in React. Also, we ignored the key property, which is used for reconciliation, not for rendering.

Okay, so rendering done and this chapter ends...? No.

Idle Time Rendering

In real react, the rendering process is a bit more complicated. To be more specific, it will use requestIdleCallback, to make more urgent tasks to be done first, lowering its own priority.

Please note that requestIdleCallback is not supported on Safari, on both MacOS and iOS (Apple Engineers, please, why? At least they are working on it, at 2024). If you are on a Mac, use chrome, or replace it with a simple setTimeout. In real react, it uses scheduler to handle this, but the basic idea is the same.

To do so, we need to know the following web native APIs.

requestIdleCallback(callback: Function): void Requests a callback to be called when the browser is idle. The callback will be passed an IdleDeadline object. The callback will have a deadline argument, which is an object with the following properties.
- timeRemaining(): number Returns the time remaining in milliseconds before the browser is no longer idle. So we should finish our work before the time is up.

So we need to split our rendering in chunks, and use requestIdleCallback to handle it. A simple way would be to just render one node at a time. It is easy- but do not be eager to do so- or you'll waste a lot of time, since we also need other work to be done while rendering.

But we can have the following code as a basic framework for what we are going to do.

import { createDom, VDomNode } from "./v-dom"

interface Fiber {
    parent: Fiber | null
    sibling: Fiber | null
    child: Fiber | null
    vDom: VDomNode,
    dom: HTMLElement | Text  | null
}

let nextUnitOfWork: Fiber | null = null

function workLoop(deadline: IdleDeadline) {
  let shouldYield = false
  while (nextUnitOfWork && !shouldYield) {
    nextUnitOfWork = performUnitOfWork(
      nextUnitOfWork
    )
    shouldYield = deadline.timeRemaining() < 1
  }
  requestIdleCallback(workLoop)
}

requestIdleCallback(workLoop)

function performUnitOfWork(nextUnitOfWork: Fiber | null): Fiber | null {
    // TODO
    throw new Error('Not implemented')
}

If you now fill // TODO with rendering vDOM, and return the next vDOM node to be rendered, you can have a simple idle time rendering. But don't be hasty- we need more work.

Fiber Tree

In the next chapter, we will implement reactivity, and the reconciliation is rather complicated- so we move some content into this part, that is the fiber tree.

Fiber tree is just a special data structure. When react handles changes, it does the following process.

Something, may be a user, or initial rendering, triggers a change.
React creates a new vDOM tree.
React calculate the new fiber tree.
React calculates the difference between the old fiber tree and the new fiber tree.
React applies the difference to the actual DOM.

You can see, fiber tree is essential for React.

The fiber tree, a little bit different from traditional tree, has three types of relations between nodes.

child of: A node is a child of another node. Please note that, in fiber tree, every node can have only one child. The traditional tree structure is represented by a child with many siblings.
sibling of: A node is a sibling of another node.
parent of: A node is a parent of another node. Different from child of, many nodes can share the same parent. You can think parent node in fiber tree as a bad parent, who only cares about the first child, but is still, in fact, parent of many children.

For example, for the following DOM,

<div>
    <p></p>
    <div>
        <h1></h1>
        <h2></h2>
    </div>
</div>

We can represent it as a tree.

div
├── p
└── div
    ├── h1
    └── h2

p is a child of the root div, but the secondary div is not a child of the root div, but a sibling of p. h1 and h2 are children of the secondary div.

When it comes to rendering, the order is mainly depth-first, but kind of different- so basically, it follows these rules. For each node, it goes through the following steps.

If this node has a unprocessed child, process the child.
If this node has a sibling, process the sibling. Repeat until all siblings are processed.
Mark this node as processed.
Process its parent.

Now let's implement that. But first, we need to trigger the rendering process. It is simple- just set the nextUnitOfWork to the root of the fiber tree.

export function render(vDom: VDomNode, parent: HTMLElement) {
    nextUnitOfWork = {
        parent: null,
        sibling: null,
        child: null,
        vDom: vDom,
        dom: parent
    }
}

After triggering the rendering, browser will call performUnitOfWork, this is where we, well, perform the work.

The first is that we need to create actual DOM elements. We can do this by creating a new DOM element, and append it to the parent DOM element.

function isString(value: VDomNode): value is string {
    return typeof value === 'string'
}

function isElement(value: VDomNode): value is VDomElement {
    return typeof value === 'object'
}

export function createDom(vDom: VDomNode): HTMLElement | Text | DocumentFragment {
    if (isString(vDom)) {
        return document.createTextNode(vDom)
    } else if (isElement(vDom)) {
        const element = document.createElement(vDom.tag === '' ? 'div' : vDom.tag)
        Object.entries(vDom.props ?? {}).forEach(([name, value]) => {
            if (value === undefined) return
            if (name === 'key') return
            if (name.startsWith('on') && value instanceof Function) {
                element.addEventListener(name.slice(2).toLowerCase(), value as EventListener)
            } else {
                element.setAttribute(name, value.toString())
            }
        })
        return element
    } else {
        throw new Error('Unexpected vDom type')
    }
}

function performUnitOfWork(nextUnitOfWork: Fiber | null): Fiber | null {
    if(!nextUnitOfWork) {
        return null
    }

    if(!nextUnitOfWork.dom) {
        nextUnitOfWork.dom = createDom(nextUnitOfWork.vDom)
    }

    if(nextUnitOfWork.parent && nextUnitOfWork.parent.dom) {
        nextUnitOfWork.parent.dom.appendChild(nextUnitOfWork.dom)
    }

    // TODO
    throw new Error('Not implemented')
}

This is the first part of the work. Now we need to construct the fiber branching out from the current one.

const fiber = nextUnitOfWork
if (isElement(fiber.vDom)) {
    const elements = fiber.vDom.children ?? []
    let index = 0
    let prevSibling = null

    while (index < elements.length) {
        const element = elements[index]
        const newFiber: Fiber = {
            parent: fiber,
            dom: null,
            sibling: null,
            child: null,
            vDom: element,
        }

        if (index === 0) {
            fiber.child = newFiber
        } else {
            prevSibling!.sibling = newFiber
        }
        prevSibling = newFiber
        index++
    }
}

Now we have a fiber tree built for the current node. Now let's follow our rules to process the fiber tree.

if (fiber.child) {
    return fiber.child
}
let nextFiber: Fiber | null = fiber
while (nextFiber) {
    if (nextFiber.sibling) {
        return nextFiber.sibling
    }
    nextFiber = nextFiber.parent
}
return null

Now we can render the vDOM, here it is. Please note that typescript is being stupid here since it can not tell the type of our virtual DOM, we need an ugly bypass here.

import { render } from "./runtime";
import { createElement, fragment, VDomNode } from "./v-dom";

function App() {
    return <>
        <h1>a</h1>
        <h2>b</h2>
    </>
}
const app = document.getElementById('app')
const vDom: VDomNode = App() as unknown as VDomNode
render(vDom, app!)

Now your vDOM is rendered to the actual DOM. Congratulations! You have done a great job. But we are not done yet.

Cumulative Commit

There will be a problem with the current implementation- if we have too many nodes that slows the whole process down, the user will see how the rendering is done. Of course, it won't leak commercial secrets or something, but it is not a good experience. We'd rather hide the dom creation behind the curtain, the submit it all at once.

The solution is simple- instead of directly committing to the document, we create an element without adding it to the document, and when we are done, we add it to the document. This is called cumulative commit.

let wip: Fiber | null = null
let wipParent: HTMLElement | null = null

export function render(vDom: VDomNode, parent: HTMLElement) {
    wip = {
        parent: null,
        sibling: null,
        child: null,
        vDom: vDom,
        dom: null,
    }
    wipParent = parent
    nextUnitOfWork = wip
}

Now, we remove the appendChild from performUnitOfWork, that is, the following part,

if(nextUnitOfWork.parent && nextUnitOfWork.parent.dom) {
    nextUnitOfWork.parent.dom.appendChild(nextUnitOfWork.dom)
}

Now if we finish all the work, we have all the fiber correctly built up with their DOM, but they are not added to the document. When such event dispatches, we call a commit function, which will add the DOM to the document.

function commit() {

}

function workLoop(deadline: IdleDeadline) {
  let shouldYield = false
  while (nextUnitOfWork && !shouldYield) {
    nextUnitOfWork = performUnitOfWork(
      nextUnitOfWork
    )
    if(!nextUnitOfWork && wip) {
        commit()
    }
    shouldYield = deadline.timeRemaining() < 1
  }
  requestIdleCallback(workLoop)
}

Now, the commit function is simple- just add all the children DOM recursively to the wip, then commit wip to the DOM.

function commit() {
    function commitChildren(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent.dom.appendChild(fiber.dom)
        }
        commitChildren(fiber.child)
        commitChildren(fiber.sibling)
    }
    commitChildren(wip)
    wipParent?.appendChild(wip!.dom!)
    wip = null
}

You can test this out by adding a timeout to commitChildren function. previously, the rendering was done step by step, but now it is done all at once.

Nested Components

You may try nested functions- like the following,

import { render } from "./runtime";
import { createElement, fragment, VDomNode } from "./v-dom";

function App() {
    return <>
        <h1>a</h1>
        <h2>b</h2>
        <div>
            <h3>c</h3>
        </div>
    </>
}

function Wrapper() {
    return <div>
        <h1>a</h1>
        <h2>b</h2>
        <div>
            <h3><App/></h3>
        </div>
    </div>
}
const app = document.getElementById('app')
const vDom: VDomNode = Wrapper() as unknown as VDomNode
render(vDom, app!)

But it won't work, since when parsing the JSX, tag is just the label name. Sure, for native elements, it is just a string, but for components, it is a function. So in the process of converting JSX to vDOM, we need to check if the tag is a function, and if so, call it.

export function createElement(tag: string | Function, props: VDomAttributes, ...children: VDomNode[]): VDomElement {
    if (tag instanceof Function) {
        return tag(props, children)
    }
    return {
        kind: tag === '' ? 'fragment' : 'element',
        tag,
        children,
        props: props ?? {},
        key: props?.key ?? undefined
    }
}

Now, props and children are required for each component. In real React, they added extra field to check- you can imagine, just by replacing functions with classes, so you have extra fields- then you provide new function to create objects, a typical factory pattern- but we take a lazy we here.

import { render } from "./runtime";
import { createElement, fragment, VDomAttributes, VDomNode } from "./v-dom";

type FuncComponent = (props: VDomAttributes, children: VDomNode[]) => JSX.Element

const Wrapper: FuncComponent = (_: VDomAttributes, __: VDomNode[]) => {
    return <div>
        <>
            <p>aa</p>
        </>
    </div>
}
const app = document.getElementById('app')
const vDom: VDomNode = Wrapper({}, []) as unknown as VDomNode
console.log(vDom)
render(vDom, app!)

Please note that in the real React, the function component call is delayed to the fiber building stage. Nonetheless, we did so for convenience, and it doesn't really harm the purpose of this series.

Fragment

However, it's still not enough. Previously, we just treated fragment as div, which is not correct. But if you just replace that with a document fragment, it won't work. The reason for this is because fragments is a one-time container- which leads to a strange behaviour- like you cannot take real things out of it, and you can not nest them, and many strange things (really, why it just won't work simpler...). So, fuck, we need to dig this shit up.

So the solution is that, we do not create DOM for fragment- we find the correct parent to add the DOM.

We need,

export function isFragment(value: VDomNode): value is VDomElement {
    return isElement(value) && value.kind === 'fragment'
}

And change the rendering,

function commit() {
    function commitChildren(fiber: Fiber | null) {
        if(!fiber) {
            return
        }
        if(fiber.dom && fiber.parent?.dom) {
            fiber.parent?.dom?.appendChild(fiber.dom)
        }

        if(fiber.dom && fiber.parent && isFragment(fiber.parent.vDom)) {
            let parent = fiber.parent
            // find the first parent that is not a fragment
            while(parent && isFragment(parent.vDom)) {
                // the root element is guaranteed to not be a fragment has has a non-fragment parent
                parent = parent.parent!
            }
            parent.dom?.appendChild(fiber.dom!)
        }

        commitChildren(fiber.child)
        commitChildren(fiber.sibling)
    }
    commitChildren(wip)
    wipParent?.appendChild(wip!.dom!)
    wip = null
}

Now, the fragment is correctly handled.

Build a Tiny React Ch1 JSX

Zend — Sun, 20 Oct 2024 09:26:28 +0000

Next Part

This series will build a tiny frontend framework, functionally similar to React, to illustrate how React works under the hood. This chapter covers JSX.

I will use Bun as runtime. Node may need extra configuration for typescript and compiling.

This tutorial is based on this tutorial, but with JSX, typescript and an easier approach to implement. You can checkout the notes and code on my GitHub repo.

JSX

Now, before we go any deeper, let's look at several important element of react- jsx first.

If you ever took a look at the transpiled code of a React component, you will see that it is just a bunch of function calls. JSX is just a syntactic sugar for React.createElement. That is, for exmaple, the following JSX code:

const element = <h1 className="greeting">Hello, world!</h1>;

Will be transpiled to:

const element = React.createElement(
  'h1',
  {className: 'greeting'},
  'Hello, world!'
);

React.createElement will create a virtual element, which is another core mechanism. To put it simply, virtual element is the element in virtual DOM. A virtual DOM is something that represents the actual DOM. Since operating on virtual DOM is simply operating js objects, it is much faster than operating on actual DOM. We will talk about virtual DOM in the next chapter. But for now, knowing that JSX is just a syntax sugar for React.createElement is enough.

The React.createElement function takes the following arguments in order,

The tag name of the element. Some tags are special, like div, span, h1, etc. These are called intrinsic elements. Some tags are user-defined components.
The props of the element. This is an object that contains the properties of the element. For example, the className of the h1 element above is greeting.
The children of the element. This can be string or element. Please note that this parameter is represented as ...children in the function signature, which means that it can take any number of arguments.

It sounds an easy job, right? So let's do it.

Implement JSX

When it comes to compilation, we can specify the function to use- by default, the function is React.createElement. But we can use our own function.

So we create a v-dom.ts, so as to define the virtual element first.

export type VDomAttributes = { 
    key?: string | number
    [_: string]: string | number | boolean | Function | undefined
}

export interface VDomElement {
  kind: 'element'
  tag: string
  children?: VDomNode[]
  props?: VDomAttributes
  key: string | number | undefined
}

export type VDomNode = 
| string
| VDomElement

Please note that we have a key field in each node (node is just a name for either text or element). This is for reconciliation, which we will talk about in the next chapter. You can safely ignore it for now.

Now we can implement the createElement function. We put it in the same file.


export function createElement(tag: string, props: VDomAttributes, ...children: VDomNode[]): VDomElement {
  console.log('createElement', tag, props, children)
    return {
        kind: 'element',
        tag,
        children,
        props,
        key: props?.key ?? undefined
    }
}

Now we instruct our compiler to use this function. We can do this by adding the following line to the top of our file.

import { createElement as h } from './v-dom'

Please note that since we are adopting react standard, we need to introduce the React type definition. We can do this by adding the following line to the top of our file.

bun i @types/react

Then in the tsconfig.json, we add the following line to the compilerOptions field.

"compilerOptions": {
    "jsx": "react",
    "jsxFactory": "createElement",
}

Now, we can take a look at the virtual element we created.

import { createElement } from "./v-dom";

function App() {
    return <div>
        <h1>a</h1>
        <h2>b</h2>
    </div>
}

console.log(App());

You will see just our defined virtual dom element based on the JSX code.

Furthermore, we can also define a fragment element- the <></>, if you don't know its official name. We can do this by adding the following line to the top of our file.

When dealing with fragment, the compiler will take the configured fragment factory to the tag in the element creation function. This is the same as how functional components work- functional component will take the function to the tag, which we will demonstrate in the next chapter.

Nonetheless, in our implementation, there is no need for complex handling- we just need to set a special tag for fragment.

export type VDomAttributes = { 
    key?: string | number
    [_: string]: string | number | boolean | Function | undefined
}

export interface VDomElement {
  kind: 'element' | 'fragment'
  tag: string
  children?: VDomNode[]
  props?: VDomAttributes
  key: string | number | undefined
}

export type VDomNode = 
| string
| VDomElement

export function createElement(tag: string, props: VDomAttributes, ...children: VDomNode[]): VDomElement {
    console.log('createElement', tag, props, children);
    return {
        kind: tag === '' ? 'fragment' : 'element',
        tag,
        children,
        props: props ?? {},
        key: props?.key ?? undefined
    }
}

export const fragment = ''

Extra compiler options,

"jsxFragmentFactory": "fragment"

Basically, fragment is just a special element with an empty tag. When it comes to the creation of fragment, the compiler will fetch the jsxFragmentFactory and put it into the tag parameter in the first parameter of createElement. So we can easily distinguish fragment from other elements.

import { createElement, fragment } from "./v-dom";

function App() {
    return <>
        <h1>a</h1>
        <h2>b</h2>
    </>
}

console.log(App());

This code will correct yield the virtual DOM. So far, we have implemented the JSX part of our tiny react.

Footnote

Ugh, this is the author from chapter three. Actually, the current implementation of JSX is not perfect. We will fix it in the third chapter. Now it does not support syntax like,

<div>
    {[1, 2, 3].map(i => <span>{i}</span>)}
</div>

This is because every {} is treated as one child, where as the map returns an array. So it will have nested children.

Also we hadn't supported functional components yet, which is in the next chapter.

You can just follow the current one and fix it later. Sorry for the inconvenience.