Build a Wasm Compiler in Roc - Part 15
In earlier articles, we have implemented a tokenizer, parser, and transformer to convert WAT syntax to a Wasm AST, and got a good portion of the code generation step in place.
Unless there are any unexpected detours, this part will cover more code generation and take us through a messy refactor that Roc helpfully makes not so messy!
The function section
At the end of part 14, I had to make a modification to our transformer to make sure that it was generating a types section for each function we defined. Functions are like imports (which are just external functions, after all) in that they have to update two sections.
Consider this basic function that stores the i32 value 8 at position 0 in
memory (acknowledging that in fact, there is no memory in this module):
(module
(func $main
(i32.const 0)
(i32.const 8)
(i32.store)
)
)
The annotated output looks like this:
0x0 | 00 61 73 6d | version 1 (Module)
| 01 00 00 00
0x8 | 01 04 | type section
0xa | 01 | 1 count
--- rec group 0 (implicit) ---
0xb | 60 00 00 | [type 0] SubType { is_final: true, supertype_idx: None, composite_type: Func(FuncType { params: [], results: [] }) }
0xe | 03 02 | func section
0x10 | 01 | 1 count
0x11 | 00 | [func 0] type 0
0x12 | 0a 0b | code section
0x14 | 01 | 1 count
============== func 0 ====================
0x15 | 09 | size of function
0x16 | 00 | 0 local blocks
0x17 | 41 00 | i32_const value:0
0x19 | 41 08 | i32_const value:8
0x1b | 36 02 00 | i32_store memarg:MemArg { align: 2, max_align: 2, offset: 0, memory: 0 }
0x1e | 0b | end
0x1f | 00 0e | custom section
0x21 | 04 6e 61 6d | name: "name"
| 65
0x26 | 01 07 | function name section
0x28 | 01 | 1 count
0x29 | 00 04 6d 61 | Naming { index: 0, name: "main" }
The first eight bytes are the same as usual, and the next six bytes define a
type section that contains one function with no params and no results. This
section is what I knew was missing from the AST and prompted the fixes to
transformFunc in the last part.
After the type section, which we hopefully already handle correctly, comes
the func section. We’ll go into detail of that very shortly, but first, note
that, much like imports, there is a custom section named “name” that we’ll
be ignoring.
For the example code, we need to focus on this part:
0xe | 03 02 | func section
0x10 | 01 | 1 count
0x11 | 00 | [func 0] type 0
0x12 | 0a 0b | code section
0x14 | 01 | 1 count
============== func 0 ====================
0x15 | 09 | size of function
0x16 | 00 | 0 local blocks
0x17 | 41 00 | i32_const value:0
0x19 | 41 08 | i32_const value:8
0x1b | 36 02 00 | i32_store memarg:MemArg { align: 2, max_align: 2, offset: 0, memory: 0 }
0x1e | 0b | end
The function section has the id 3, so the first byte is 03. The next byte is
02 which means the function section itself is only two bytes long! The first
of those bytes are the number of functions in a vector, and the second one is
00, which is an encoding of the index in the types section for the type of
the function at position 00.
So the funcs section is just a sort of lookup to map functions to their
types. I assume this is so we don’t have to specify the types of functions that
have the same type more than once. If there are two functions that don’t have
any return values or parameters, they can use the same index in the types
section. Ideally, our transformer would have performed this deduplication on
the AST before passing it into the code generator, but I’m going to kick that
can down the road. Our target sample only has two function types and they have
different signatures anyway.
After that tiny little func section, we get the code section, which is
where all the juicy bits live. The code section has an id of 10 in decimal,
which is 0xa in hex. There are 11 bytes in this particular code section, so
the second byte of the code section is 0xb. And there is only one func
in this code section, so the next byte is 01 for the length of the vector.
The function starts with the LEB-encoded size of the function in bytes, which is 9 in this example (this makes it easy for a runtime to skip over the function without reading it).
The rest of the function is collected into a vector followed by an expression.
In this example, the vector is empty, so we only need a single 00 byte to
record the length of that vector. This first vector holds any “locals” in the
function, and would map to the types of local variables. Since our one and only
function doesn’t have locals (and we don’t even support instructions that can
set local values), I’m just going to hardcode it to 0!
The expression encodes a set of instructions followed by an 0b instruction to
denote the end of the expression. We already generate an expression in the
generateData function that we’ll be able to reuse. First, though, we’ll have
to add some instructions to the generateInstruction function that it calls.
In this example, we have two i32_const instructions, which
generateInstruction already knows how to handle; the opcode 0x41 followed by
the numbers to be pushed onto the stack (0 and 8). The i32_store instruction
has an opcode of 36. It is followed by two bytes representing the alignment
and offset of those bytes in memory. The alignment is represented as a power of
two (so in this case 2 ** 2 = 4) and offset is 0. It took me a bit of hunting
to figure out where these values came from. They are the default values (a
32-bit integer takes up four bytes, so an alignment of 2 makes sense). It is
possible to change this default by passing attributes such as alignment=1 to
the WAT syntax. However, our parser doesn’t even know how to handle attributes,
and I think even the tokenizer would choke on the = syntax.
So to make a long story extremely short, I’m going to hardcode this implementation detail.
The last byte is 0b: the usual instruction to end an expression.
Before tackling all the instructions in our “hello world” example, let’s see
if we can get the above three-instruction sample to cooperate. First, we’ll need
to support i32.store in generateInstruction. I’ll start with an expectation:
expect
result = generateInstruction (I32Store { row: 1, column: 1 })
result == [0x36, 0x02, 0x00]
This causes our compiler to hit the crash arm in generateInstruction,
which is currently defined like this:
generateInstruction : Transformer.Instruction -> List U8
generateInstruction = \instruction ->
when instruction is
I32Const { element } -> List.concat [0x41] (element |> generateU32)
_ -> crash "unexpected instruction encountered"
To satisfy our “ignore alignment and offset altogether” strategy, we can add
a dead simple arm to that when statement:
I32Store {} -> [0x36, 0x02, 0x00]
The func section, even though it is really simple, needs to be able to look
up a function in the types array. So we’ll need to use the same trick we used
for generateImportFactory. If we were feeling really zealous, we could even
try to make that function generic! Instead, I’m going to pull the typedDict
pipeline out of generateImportFactory into its own function:
typesToDict : List (Common.Positionable Transformer.FuncType) -> Dict Str U32
typesToDict = \funcTypes ->
funcTypes
|> List.map Common.extractElement
|> List.mapWithIndex \funcType, index -> (funcType.identifier, index |> Num.toU32)
|> Dict.fromList
This shortens the generateImportFactory function a bit, as follows:
generateImportFactory : List (Common.Positionable Transformer.FuncType) -> (Transformer.Import -> List U8)
generateImportFactory = \funcTypes ->
typeDict = typesToDict funcTypes
\importDef ->
when importDef.definition is
Func identifier if Dict.contains typeDict identifier ->
[]
|> List.concat (generateVector (Str.toUtf8 importDef.namespace) \b -> [b])
|> List.concat (generateVector (Str.toUtf8 importDef.name) \b -> [b])
|> List.append 0x00
|> List.concat (generateU32 ((Dict.get typeDict identifier) |> Result.withDefault 999999))
Func _ ->
crash "Unexpected identifier shouldn't get past AST generation"
_ ->
crash "Haven't implemented non func imports"
The majority of that code is custom and specific to imports, so I’m comfortable
duplicating it for the functions section. Roc may even be smart enough not to
call typesToDict on the same value twice; it at least has the option to
cache the response since everything is immutable!
The generateFuncFactory version looks like this:
generateFuncFactory : List (Common.Positionable Transformer.FuncType) -> (Transformer.Func -> List U8)
generateFuncFactory = \funcTypes ->
typeDict = typesToDict funcTypes
\func ->
when func.identifier is
identifier if Dict.contains typeDict identifier ->
typeDict
|> Dict.get func.identifier
|> Result.withDefault 999999
|> generateU32
_ ->
crash "Unexpected identifier shouldn't get past AST generation"
In the funcs section, each individual function is just an index into the
types array. As long as we can find the type (which the AST generator and
static analysis should guarantee, if we had implemented that part), we just
encode the index as a u32 and call it a day.
We also need to add a call to this new function in generate:
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 0x02 module.imports (generateImportFactory module.types)
|> concatSection 0x03 module.funcs (generateFuncFactory module.types)
|> concatSection 0x05 module.mems generateMem
|> concatSection 0x0b module.datas generateData
That should have the func section taken care of, so now we can move onto the
code section. The code section doesn’t need access to the types section, so we
can skip the Factory thing.
The generateCode function, which generates the code section for just one
function, has three pieces: the number of bytes in the function, the locals
vector (which we’ll hardcode to zero because I am feeling IMPATIENT), and
the expression containing all the instructions.
generateCode : Transformer.Func -> List U8
generateCode = \func ->
bytes =
[0x00] # hardcoding the locals to be empty for now
|> List.concat (generateExpression func.instructions)
bytes
|> List.len
|> Num.toU32
|> generateU32
|> List.concat bytes
I collect the bytes first so I can set up a pipeline to count them and convert that to U32, then concatenate the bytes onto the end of that encoded count.
Then we need to add another call to concatSection to our pipeline. It’s at position
10 so it needs to go between the calls to generate Mem and Data:
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 0x02 module.imports (generateImportFactory module.types)
|> concatSection 0x03 module.funcs (generateFuncFactory module.types)
|> concatSection 0x05 module.mems generateMem
|> concatSection 0x0a module.funcs generateCode
|> concatSection 0x0b module.datas generateData
Finally, for completeness, here’s a rather long expectation to test our little function:
expect
# (module
# (func $main
# (i32.const 0)
# (i32.const 8)
# (i32.store)
# )
# )
result = generate {
position: { row: 1, column: 1 },
element: {
types: [
{
position: { row: 2, column: 4 },
element: {
identifier: "main",
param: [],
result: [],
},
},
],
funcs: [
{
position: { row: 2, column: 4 },
element: {
identifier: "main",
instructions: [
I32Const {
position: { row: 3, column: 8 },
element: 0,
},
I32Const {
position: { row: 4, column: 8 },
element: 8,
},
I32Store { row: 5, column: 8 },
],
},
},
],
mems: [],
datas: [],
imports: [],
exports: [],
},
}
result
==
[
0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00,
0x01, # type section
0x04, 0x01, 0x60, 0x00, 0x00,
0x03, # func section
0x02, 0x01, 0x00,
0x0a, # code section
0x0b, # 11 bytes
0x01, # 1 function
0x09, # nine bytes in this function
0x00, # 0 local blocks
0x41, # i32_const
0x00,
0x41, # i32_const
0x08,
0x36, # i32_store with hardcoded alignment/offset
0x02,
0x00,
0x0b, # end
]
Remaining instructions
The function in our hello world example has to call two other instructions, and one of them is going to throw a wrench in the works.
Let’s start with the simple instruction: drop is encoded as the single byte
0x1a and doesn’t have any fancy parameters with it. Just a simple match
arm in our generateInstruction pattern:
generateInstruction : Transformer.Instruction -> List U8
generateInstruction = \instruction ->
when instruction is
I32Const { element } -> List.concat [0x41] (element |> generateU32)
I32Store {} -> [0x36, 0x02, 0x00]
Drop {} -> [0x1a]
_ -> crash "unexpected instruction encountered"
The last one, however, is going to be annoying. The call instruction
in our AST is associated with an identifier, fd_write. fd_write maps
to a specific index in the function index space.
But wait… the fd_write function isn’t even in the funcs table. It’s an
import. How would it know whether to search imports or functions?
I had to scour the damn specification before I found this sneaky paragraph:
The index space for functions, tables, memories and globals includes respective imports declared in the same module. The indices of these imports precede the indices of other definitions in the same index space.
I thought this was going to mean a horrible refactor to pipe imports to our
typesToDict function, but it turns out its (sort of) ok. Imports and
functions both update the funcTypes table already. So long as the WAT code
always happens to place any imports before any functions, that table already
contains all the possible identifiers. Since our sample code DOES happen
to place imports before functions, it’s “good enough.”
One thing this project has taught me is not to complain about unexpected or poorly explained compiler errors because the authors neglected to cover every edge case. As a super young language, Roc has a lot of these, of course. Now I know why! I’m sure they’ll chip away at them until their 1.0 release, but holy crap what a job!
We still have a fairly nasty refactor ahead of us, though. The
generateInstruction function doesn’t have access to the module.types field.
We can’t trivially change it to a generateInstructionFactory-style function
because it’s only called from generateExpression. generateExpression is
called from generateData and generateCode. Making all of those support the
Factory is just going to be messy.
Refactoring with confidence
In fact, I think it’s better to ditch our “clever” factory pattern, and pass the types array down to either all functions or just the ones that need it.
If I was expecting to maintain this code, I would probably write a Context
object that gets passed to all of the functions. There’s a good chance that
types are not the only thing that would need passing around, so I could put
more of them on a context. But since I’m not planning to maintain it, I’m
instead just going to pass the types dictionary.
We are too close to our goal to let good design get in the way!
These words are the reason we have things like make and Javascript. Oh
dear.
I am glad that I feel so confident refactoring Roc code at this point; once it compiles and the existing unit tests pass, we’ll know we’re done.
I started by committing all my current changes to git. Always a good policy
when starting an uncertain refactor. Then I put the typeDict definition in
generate and passed it to all the concatSection calls:
generate : Common.Positionable Transformer.Module -> List U8
generate = \modulePositioned ->
module = modulePositioned |> Common.extractElement
typeDict = typesToDict module.types
[0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
|> concatSection 0x01 typeDict module.types generateFuncType
|> concatSection 0x02 typeDict module.imports (generateImportFactory module.types)
|> concatSection 0x03 typeDict module.funcs (generateFuncFactory module.types)
|> concatSection 0x05 typeDict module.mems generateMem
|> concatSection 0x0a typeDict module.funcs generateCode
|> concatSection 0x0b typeDict module.datas generateData
Now it’s time for another round of “follow the errors”. Six errors right now tell
me that concatSection is getting too many arguments. These can be fixed by
updating the signature of that function:
concatSection : List U8, U8, Dict Str U32, List (Common.Positionable x), (x -> List U8) -> List U8
concatSection = \existingBytes, sectionId, typeDict, items, generateOne ->
That solves all six errors, but introduces an Unused Argument warning. We can solve
that by passing the typeDict into generateSection:
sectionBytes =
items
|> List.map Common.extractElement
|> generateSection sectionId typeDict generateOne
Next error tells me to add the new argument to generateSection, and since I
know it’s going to complain about me not using it, I’m just going to forward it
to generateVector while I’m at it:
generateSection : List a, U8, Dict Str U32, (a -> List U8) -> List U8
generateSection = \items, sectionId, typeDict, generateOne ->
when items is
[] -> []
_ ->
section = generateVector items typeDict generateOne
[]
|> List.append sectionId
|> List.concat (section |> List.len |> Num.toU32 |> generateU32)
|> List.concat section
Now, generateVector is called from a lot of different places, so it seems
clear that we’re going to be passing this Dict everywhere. This is why object
oriented programming got popular. If we used an OOP language, we’d just add
typeDict as an instance variable on an object and it would be available in
every method. Unfortunately, it took the industry decades to realize that is a
really good way to have no idea where anything came from or is stored.
Note: You can do object oriented design in any language. Remember how I suggested creating a context and passing that to every function? That is really a form of OOD. After all, that’s basically what the
selfargument in Python or Rust classes is.
Anyway, as a first pass and to move on from our current error, generateVector
now needs to look like this:
generateVector : List a, Dict Str U32, (a -> List U8) -> List U8
generateVector = \items, typeDict, encodeOne ->
result = items |> List.len |> Num.toU32 |> generateU32
items
|> List.walk result \current, next ->
List.concat current (encodeOne typeDict next)
Notice that I passed the typeDict down to encodeOne. But encodeOne in the
generateVector signature doesn’t accept a typeDict. So we’ll need to update
that as well:
generateVector : List a, Dict Str U32, (Dict Str U32, a -> List U8) -> List U8
generateVector = \items, typeDict, encodeOne ->
Now I have seven errors telling me that I’m passing the wrong types to
generateVector in various places. The first such call site is one we already
changed: generateSection needs to be updated again so that the generateOne
function accepts two arguments:
generateSection : List a, U8, TypeDict, (TypeDict, a -> List U8) -> List U8
I can tell I’m going to get real sick of typing Dict Str U32, so I changed it
to TypeDict and added a type alias at the top of the file:
TypeDict : Dict Str U32
concatSections needs to be updated again as well, so that the generateOne
param also accepts a TypedDict (so it can be passed down to
generateSection).
At this point I have 12 errors. About half are calls to generateVector that
don’t have the TypeDict passed down, and half are the various functions we pass
into concatSection in generate. We can kind of kill two birds with each
stone now, because the functions called from generate are the ones that
(directly or indirectly) call generateVertor. Here’s the list of functions:
|> concatSection 0x01 typeDict module.types generateFuncType
|> concatSection 0x02 typeDict module.imports (generateImportFactory module.types)
|> concatSection 0x03 typeDict module.funcs (generateFuncFactory module.types)
|> concatSection 0x05 typeDict module.mems generateMem
|> concatSection 0x0a typeDict module.funcs generateCode
|> concatSection 0x0b typeDict module.datas generateData
Four of those functions need to be changed to accept a TypeDict, and the two
factory functions can be simplified (thank god, I need some simpler in my
life right now) to accept a TypeDict instead of generating it by itself.
Let’s start with generateFuncType. It currently looks like this:
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)
The function needs to accept two parameters now, and also pass the new
parameter down to generateVector:
generateFuncType : TypeDict, Transformer.FuncType -> List U8
generateFuncType = \typeDict, funcType ->
[]
|> List.append 0x60
|> List.concat (generateVector (funcType.param |> List.map Common.extractElement) typeDict generateType)
|> List.concat (generateVector (funcType.result |> List.map Common.extractElement) typeDict generateType)
Sadly, we still have 12 errors, because generateType doesn’t currently accept
typedDict. Let’s fix that:
generateType : TypeDict, Transformer.Type -> List U8
generateType = \_, type ->
when type is
I32 -> [0x7f]
This function has the luxury of ignoring the TypedDict (the variable _
ignores it), but to cooperate with generateVector, it still needs to accept
it. This knocks our error count down to 11. I feel like that deserves a major
celebration!
We’ll need to do something similar with generateMem, generateCode, and
generateData. generateMem can just ignore the parameter again, since the
only function it calls is generateU32, and there is no way that LEB128
encoding needs to know about the types in our app.
While the generateExpression function doesn’t need TypeDict right now, it’s
going to. It generates instructions and this whole refactor started because of
needing the TypeDict in instructions.
Here’s generateData as an example:
generateData : TypeDict, Transformer.Data -> List U8
generateData = \typeDict, data ->
[0x00]
|> List.concat (generateExpression typeDict data.offset)
|> List.concat (generateVector (Str.toUtf8 data.bytes) typeDict \_, b -> [b])
Note that, in addition to generateExpression, I also passed typeDict down
to generateVector and the encodeOne lambda ignores it. generateCode
is similar, so I want show the changes.
The generateExpression signature needs to be changed as well, but it can
ignore the new parameter for now.
And we’re down to 9 errors! Progress!
Removing the Factories
Let’s change generate to call non-factory functions that accept typeDict as a
parameter instead of passing them down. Currently we have these two lines:
|> concatSection 0x02 typeDict module.imports (generateImportFactory module.types)
|> concatSection 0x03 typeDict module.funcs (generateFuncFactory module.types)
The simpler version will look more like the others:
|> concatSection 0x02 typeDict module.imports generateImport
|> concatSection 0x03 typeDict module.funcs generateFunc
We also have to rename and simplify the called functions, of course. This means
changing generateFuncFactory from the following:
generateFuncFactory : List (Common.Positionable Transformer.FuncType) -> (Transformer.Func -> List U8)
generateFuncFactory = \funcTypes ->
typeDict = typesToDict funcTypes
\func ->
when func.identifier is
identifier if Dict.contains typeDict identifier ->
typeDict
|> Dict.get func.identifier
|> Result.withDefault 999999
|> generateU32
_ ->
crash "Unexpected identifier shouldn't get past AST generation"
to this:
generateFunc : TypeDict, Transformer.Func -> List U8
generateFunc = \typeDict, func ->
when func.identifier is
identifier if Dict.contains typeDict identifier ->
typeDict
|> Dict.get func.identifier
|> Result.withDefault 999999
|> generateU32
_ ->
crash "Unexpected identifier shouldn't get past AST generation"
Similarly generateImportFactory becomes as follows, paying attention to the updates to generateVector:
generateImport : TypeDict, Transformer.Import -> List U8
generateImport = \typeDict, importDef ->
when importDef.definition is
Func identifier if Dict.contains typeDict identifier ->
[]
|> List.concat (generateVector (Str.toUtf8 importDef.namespace) typeDict \_, b -> [b])
|> List.concat (generateVector (Str.toUtf8 importDef.name) typeDict \_, b -> [b])
|> List.append 0x00
|> List.concat (generateU32 ((Dict.get typeDict identifier) |> Result.withDefault 999999))
Func _ ->
crash "Unexpected identifier shouldn't get past AST generation"
We’re down to five errors, and all of them are in unit tests. They can all be
solved by passing typesToDict [] to various function invocations. For
example, the test for generateVector needs to look like this:
expect
result = generateVector [624485, 3, 0, 123456] (typesToDict []) \_, num -> generateU32 num
result == [0x04, 0xE5, 0x8E, 0x26, 0x03, 0x00, 0xc0, 0xc4, 0x7]
You could also use (Dict.empty {}), but (typesToDict []) documents itself better.
And with that, our refactor is complete. The code compiles and the tests pass. Now we just need to remember why the hell we did this!
The call instruction
The call instruction has the opcode 0x10, and accepts a single U32 (LEB128
encoded as usual) as its parameter. The U32 needs to come from the
typeDict.
Right now generateExpression calls generateInstruction without passing it down,
so the first step is to fix that:
generateExpression : TypeDict, List Transformer.Instruction -> List U8
generateExpression = \typeDict, instructions ->
instructions
|> List.walk
[]
(\state, next ->
List.concat state (generateInstruction typeDict next)
)
|> List.append 0x0b
And of course the generateInstruction signature needs to accept two parameters now:
generateInstruction : TypeDict, Transformer.Instruction -> List U8
generateInstruction = \typeDict, instruction ->
A couple unit tests that call generateInstruction will also need to be updated.
Now we just need to add the Call arm to generateInstruction:
generateInstruction : TypeDict, Transformer.Instruction -> List U8
generateInstruction = \typeDict, instruction ->
when instruction is
I32Const { element } -> List.concat [0x41] (element |> generateU32)
I32Store {} -> [0x36, 0x02, 0x00]
Drop {} -> [0x1a]
Call { element: { identifier } } ->
when Dict.get typeDict identifier is
Err _ -> crash "Should not receive call instruction with invalid identifier"
Ok num -> num |> generateU32 |> List.prepend 0x10
I also removed the wildcard arm, since we are now covering all possible
instructions and it will be nice to get a compiler error here if we add a new
one to Transformer.Instruction.
This code matches the Call structure and looks up the id in typeDict It
converts the result to a U32 and prepends the 0x10 opcode byte. I’m pretty
sure that is sufficient. This post is getting long, so I won’t include yet
another crazy long unit test, but know that I did write one! It was fiddly.
This part ended up going a little long, but I think we’ll be able to wrap it up in one more post! See you there.
