Build a Wasm Compiler in Roc - Part 13
In earlier articles, we have implemented a tokenizer, parser, and transformer to convert WAT syntax to a Wasm AST, and started building a code generator to create bytes from that AST. The structure of the code generator is in place already, so our remaining task is to generate all the other sections of our input module.
Data section
The next section I want to implement is the data section, which is described in the Wasm spec here.
Let’s start by inspecting a basic data section with the following WAT syntax:
(module
(data (offset (i32.const 8)) "hello world")
)
Running wasm-tools dump main.wat
on this file spits out the following
analysis:
0x0 | 00 61 73 6d | version 1 (Module)
| 01 00 00 00
0x8 | 0b 11 | data section
0xa | 01 | 1 count
0xb | 00 | data memory[0]
0xc | 41 08 | i32_const value:8
0xe | 0b | end
0xf |-------------| ... 11 bytes of data
We can confirm that the data section contains the “hello world” string using
xxd. I use the command xxd -c 4 -g 1 main.wasm
. This tells xxd to print four
columns of bytes (-c 4) and to not group bytes together, which has the effect
of placing a space between each byte. This makes comparing the output of xxd
and wasm-tools dump
a little easier, though xxd
doesn’t have the context to
wrap lines at meaningful boundaries:
00000000: 00 61 73 6d .asm
00000004: 01 00 00 00 ....
00000008: 0b 11 01 00 ....
0000000c: 41 08 0b 0b A...
00000010: 68 65 6c 6c hell
00000014: 6f 20 77 6f o wo
00000018: 72 6c 64 rld
The first eight bytes are the header, as usual. Then we see the byte 0b
,
which is the data section identifier (decimal 11
). The next byte is 11
, and
represents the length of all the bytes in the data section. If you are
thinking, “But ‘hello world’ has 11 bytes plus there’s the wrapping vector
logic,” you may need a reminder that 0x11
in hex is 17
in decimal… I did!
The 01
is the start of a vector definition with one element, so we’ll
be able to use generateVector
for that. The remaining bytes encode
the section itself. It’s clear where the actual data bytes are stored,
so we really only need to understand 00 41 08 0b
.
There are three different layouts that a data segment can have, though we’ll
only be supporting one of them. The 00
byte selects the first of those
layouts; 01
and 02
would be used to indicate the other layouts, if we later
want to support them.
Then we need to encode an expression. An expression is always encoded as an
arbitrary list of instructions followed by an explicit 0b
byte as a sentinel.
I’m not clear on which instructions are legal in a data segment, but I’m going
to pretend that it can be any arbitrary collection of instructions. That allows
me to reuse a generateInstruction
function. The 41
byte is the op-code for
i32.const
, and the 08
is the LEB128 encoding of the value passed to
i32.const
(8
, in this case).
Finally, there is a second 0b
before the actual bytes of “hello world”. This
one is unrelated to the “end expression op-code”. In fact, it’s just the count
for another vector: a vector with 11 (0x0b
) elements. The elements in this
vector are bytes, so it is the length of the bytes in the string “hello world”.
So we’ll need these new generator functions:
generateInstruction
to encode aTransformer.Instruction
to the appropriate format.generateExpression
to encode multiple instructions and append theend
byte,0x0b
.generateData
to hard code the00
byte followed by the expressions and data.generateDatas
to hard code the0x0b
section identifier and vector of data segments.
As I hinted in part 12, I want to try to make the last one generic so that we
can reuse the same structure for generateMems
and generateDatas
instead of
copy-pasting the whole thing.
Let’s start from the inside and work our way out.
Generating instructions
Instructions are generally encoded with a byte identifying the op-code, followed
by any op-code specific data that the instruction needs. Our toy compiler only
supports a four-instruction subset of the many
instructions
available in Wasm. For now, I’m just going to handle the i32.const
version.
First, add Instruction
to the export sections of the Transformer
module,
since we’ll need to access it in our generator logic. Then add a simple
expectation:
expect
result = generateInstruction
(
I32Const {
position: { row: 1, column: 1 },
element: 8,
}
)
result == [0x41, 0x08]
This implementation will need some improvement when we implement functions, but it should suffice for now:
generateInstruction : Transformer.Instruction -> List U8
generateInstruction = \instruction ->
when instruction is
I32Const { element } -> List.concat [0x41] (element |> generateU32)
_ -> crash "unexpected instruction encountered"
Generating Expressions
Expressions just need to concatenate the bytecode for each input instruction
and then append an 0x0b
op-code to indicate the “last instruction”.
generateExpression : List Transformer.Instruction -> List U8
generateExpression = \instructions ->
instructions
|> List.walk
[]
(\state, next ->
List.concat state (generateInstruction next)
)
|> List.append 0x0b
If you’re worried that one of the existing instructions might contain an 0x0b
that prematurely ends the expression, note that 0x0b
is only interpreted as
an “end” op-code if it comes at the position of a “next instruction”. If an
existing instruction has an 0x0b
in it (e.g. i32.const 11
) it will be
processed by the inverse of generateInstruction
, rather than as the “end
instruction”.
Generating one data
Since we’re only generating the one data format, we can hard code a single
0x00
, followed by an encoding of the instructions in the
Transformer.Data.offset
, as an expression, and finally, a vector of bytes.
First add the Data
export to the Transformer
module. Next, here’s a
test:
expect
result = generateData {
offset: [
I32Const {
position: { row: 1, column: 1 },
element: 8,
},
],
bytes: "hello world",
}
result == [0x00, 0x41, 0x08, 0x0b, 0x0b, 'h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd']
And here’s the implementation:
generateData : Transformer.Data -> List U8
generateData = \data ->
[0x00]
|> List.concat (generateExpression data.offset)
|> List.concat (generateVector (Str.toUtf8 data.bytes) \b -> [b])
The main thing of note here is that the “transformer” for each byte in a vector
of bytes is just a list containing that byte. The generateVector
function takes
care of combining them.
Generating Sections
We already have a generateMems
function that looks like this:
generateMems : List Transformer.Mem -> List U8
generateMems = \mems ->
when mems is
[] -> []
_ ->
memorySection = generateVector mems generateMem
[0x05]
|> List.concat (memorySection |> List.len |> Num.toU32 |> generateU32)
|> List.concat memorySection
We could copy-paste this to create a generateDatas
function, but copy-paste
is a code smell. Let’s see if we can create a generic generateSection
function instead.
The only two things that need to change between different sections is the section identifier and the function that is called to generate a single element. Can we pass those both in as arguments instead? Let’s try:
generateSection : List a, U8, (a -> List U8) -> List U8
generateSection = \items, sectionId, generateOne ->
when items is
[] -> []
_ ->
section = generateVector items generateOne
[]
|> List.append sectionId
|> List.concat (section |> List.len |> Num.toU32 |> generateU32)
|> List.concat section
This works, but there’s still some duplication in the generate
function:
generate : Common.Positionable Transformer.Module -> Result (List U8) (List Common.Error)
generate = \modulePositioned ->
module = modulePositioned |> Common.extractElement
Ok ([0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00])
|> Result.map \bytes ->
List.concat
bytes
(
module.mems
|> List.map Common.extractElement
|> generateSection 0x05 generateMem
)
|> Result.map \bytes ->
List.concat
bytes
(
module.datas
|> List.map Common.extractElement
|> generateSection 0x0b generateData
)
We can certainly move the Common.extractElement
into the helper function, but
moving the Result.map
and concatenation feels like it deserves its own
generic function. However, I want to take a step back, first. Do we even need a
Result.map
? More specifically, does generate
need to return a Result
? Can
generate
fail?
If the AST is valid, then generation cannot fail. But how valid is the AST? An even more interesting question would be “how valid is the AST if we had a static analysis step?”
So far, none of our encoders have returned errors, and it feels like that will
continue to be the case. If the AST structurally follows the input
specifications, the generator will always return bytes, and never return an
error. It’s not clear to me if that will always be the case, but I’m going to
simplify the generate
function to always return a List U8
:
generate : Common.Positionable Transformer.Module -> List U8
generate = \modulePositioned ->
module = modulePositioned |> Common.extractElement
[0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
|> List.concat (
module.mems
|> List.map Common.extractElement
|> generateSection 0x05 generateMem
)
|> List.concat (
module.datas
|> List.map Common.extractElement
|> generateSection 0x0b generateData
)
I also had to remove some surrounding Ok()
s from the relevant tests an change
the Result.try Generator.generate
to a Result.map Generator.generate
in
main.roc
.
There’s still some duplication, but it’s now clear what we would want to
extract to its own function. My rule of thumb with functional languages is “if
a pipeline contains duplicate structures that don’t fit on one line, make it a
separate function.” Here’s a new concatSection
function:
concatSection : List U8, U8, List (Common.Positionable x), (x -> List U8) -> List U8
concatSection = \existingBytes, sectionId, items, generateOne ->
sectionBytes =
items
|> List.map Common.extractElement
|> generateSection sectionId generateOne
List.concat existingBytes sectionBytes
Now we can make generate
very pretty:
generate : Common.Positionable Transformer.Module -> List U8
generate = \modulePositioned ->
module = modulePositioned |> Common.extractElement
[0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
|> concatSection 0x05 module.mems generateMem
|> concatSection 0x0b module.datas generateData
Having all this infrastructure in place should make writing the other sections a little more straightforward. Let’s test that theory by implementing another section.
Modeling the types section
Our hello world example doesn’t have an explicit types section, but we do generate types implicitly in our import section. The following WAT code generates the same types section as our import section, which makes it a little bit easier to analyze in isolation:
(module
(type (
func (param i32 i32 i32 i32) (result i32))
)
)
wasm-tools dump
on this WAT module looks like this:
0x0 | 00 61 73 6d | version 1 (Module)
| 01 00 00 00
0x8 | 01 09 | type section
0xa | 01 | 1 count
--- rec group 0 (implicit) ---
0xb | 60 04 7f 7f | [type 0] SubType { is_final: true, supertype_idx: None, composite_type: Func(FuncType { params: [I32, I32, I32, I32], results: [I32] }) }
| 7f 7f 01 7f
The type section id is the byte 01
, and it is followed by 09
to reflect the
number of bytes in the section.
Then we have another 01
this time reflecting the number of types in the
section (as a vector encoding). The contents of the vector will be
functypes.
Functypes are encoded as the byte 0x60
to represent the func type, followed
by two vectors. The first vector has 04
elements, and all of those elements
are the byte 0x7f
, which means i32
. The second vector has 01
elements,
and that one element is also 0x7f
. Indeed, i32 is the only primitive type
in our app!
So we’ll need a new generateFuncType
function, but we can make it pretty
simple because we only need parameters and results of type 0x7f
. Other than
that, we can reuse all our existing machinery.
Let’s start by exporting the Type
and FuncType
from Transformer
. It is a
tag union to reflect the fact that there are other types we are not currently
modelling, but the only thing it can contain right now is I32
.
The test is probably the easiest test in the whole app:
expect
result = generateType I32
result == [0x7f]
And the generateType
function is equally simple:
generateType : Transformer.Type -> List U8
generateType = \type ->
when type is
I32 -> [0x7f]
Note that there’s no wildcard handler here, so if we do add new types later, we won’t forget to update this section.
The test for generateFuncType
is interesting:
expect
result = generateFuncType {
identifier: "helloFunc",
param: [
{ element: I32, position: { row: 1, column: 1 } },
{ element: I32, position: { row: 1, column: 1 } },
],
result: [{ element: I32, position: { row: 1, column: 1 } }],
}
result == [0x60, 0x02, 0x7f, 0x7f, 0x01, 0x7f]
Our funcType
supplies an identifier, but the identifier is not used anywhere
in the output bytes.
For now, let’s just get generateFuncType
passing the test:
generateFuncType : Transformer.FuncType -> List U8
generateFuncType = \funcType ->
[]
|> List.append 0x60
|> List.concat (generateVector (funcType.param |> List.map Common.extractElement) generateType)
|> List.concat (generateVector (funcType.result |> List.map Common.extractElement) generateType)
Finally, we can create a overall test for the functypes section:
expect
# (module
# (type (
# func (param i32 i32 i32 i32) (result i32))
# )
# )
result = generate {
position: { row: 1, column: 1 },
element: {
types: [
{
position: { row: 2, column: 3 },
element: {
identifier: "helloFunc",
param: [
{ element: I32, position: { row: 1, column: 1 } },
{ element: I32, position: { row: 1, column: 1 } },
],
result: [{ element: I32, position: { row: 1, column: 1 } }],
},
},
],
funcs: [],
mems: [],
datas: [],
imports: [],
exports: [],
},
}
result
== [
0x00, 0x61, 0x73, 0x6d,
0x01, 0x00, 0x00, 0x00,
0x01, 0x07,
0x01,
0x60,
0x02, 0x7f, 0x7f,
0x01, 0x7f,
]
The only caveat with the updated generate
function is that the sections have
to be in order of their identifier. This means that the module.types
section
(with byte id 0x01
) has to be placed bfore the mems section (0x05
):
generate : Common.Positionable Transformer.Module -> List U8
generate = \modulePositioned ->
module = modulePositioned |> Common.extractElement
[0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
|> concatSection 0x01 module.types generateFuncType
|> concatSection 0x05 module.mems generateMem
|> concatSection 0x0b module.datas generateData
Now that we have types, which in our sample code are only generated by the import section, we can move on to the import section. However, the import section has a few interesting complications, so I think I’ll leave it for the next part!