Why and how we wrote a compiler in Rust (blog post series 2/X): the stack
This post is the second post of the series about how we wrote a compiler in Rust at Cosmian for the CipherCompute product. If you missed the first one here is the link, I recommend you to read it to really understand what we’re trying to achieve and know the context.
As a reminder, the previous article explains our objective is to be able to execute a sMPC (secure multi party computation) program in a widespread and known language, Rust in our case. To achieve this, we developed a compiler named WaSca which takes Wasm (WebAssembly) and compiles it to bytecode understandable by the Scale virtual machine to run sMPC computations.
In this post I’m going to explain a concept developed in the WaSca compiler called the WaSca stack. Indeed there is a fundamental difference between the 2 assembly codes we are dealing with. They have a distinct way of dealing with function calls, memory management and variable types. Wasm is a stack based assembly language. As such, its memory does not contain registers. Instead, all instruction arguments are pushed and popped on a shared stack. Scale asm on the other hand, uses a register based assembly language. Therefore, all arguments and return values are passed with registers. Let’s explain in a more detailed way how each one works.
Wasm instructions
If we take a very simple program, for example a simple addition of 2 numbers (5 + 2).
Representing this program in Wasm is pretty simple, it’s only 3 lines with 2 different kinds of instruction. The first one is i32.const which takes 1 constant as parameter in order to push an i32 constant on the Wasm stack. The other one isi32.add which takes no parameter but instead reads the last 2 elements on the Wasm stack, makes the addition between them and pushes the result on the stack. Here’s an animation to see exactly what happens when running this small program.
Scale asm instructions
Now let’s see the same program but written in Scale asm. It’s also only 3 lines with 2 different kinds of instructions.
The first one is anldint instruction which takes 2 parameters: the first one is the target register to store the constant value located as the second parameter. The second one is theaddint instruction which takes 3 parameters: the first one is the target register to save the result of the addition and the 2 last parameters are registers that will be added together.
So the main difference is that in Wasm almost all instructions use the Wasm stack which will be filled and emptied depending on the function arguments, the instructions, etc. At the end of the execution the Wasm stack is supposed to be empty as all values should be consumed by the instructions, while in Scale asm the arguments or instruction parameters are passed on registers. Registers are memory allocated in some memory zones which keep the value assigned until another value is set for this specific register. In Scale asm when an instruction produces value it goes in a register and not in a global stack.
The WaSca stack
We can easily understand that these two behaviors are very different. To generate our bytecode in Scale we generate Scale asm and then the way we transform Wasm to Scale asm is very different from just doing 1:1 for each instruction. We also have an available stack in Scale by using push and pop instructions but it would be very inefficient to try to use it as it is done in Wasm. By doing this we would have performance issues because a lot more instructions than needed would be used but also we would have less optimizations than we have today.
Those are the reasons why we implemented what we call the “WaSca stack” (to differentiate from the other ones) which only lives in the compiler, this stack doesn’t exist anymore at runtime. This gives us the ability to easily switch from the stack paradigm to the register paradigm. This WaSca stack is basically a stack used in the compiler to reproduce the same behavior as the Wasm stack. Unlike the Wasm stack which lives at runtime, the WaSca stack is only useful for us to compile to the Scale asm. Instead of pushing and popping values as in Wasm stack we push either constant values or registers. Thanks to this behavior we are also able to add some optimizations into our generated code.
But sometimes a stack based approach makes some codes very complex or nearly impossible to write. For this reason Wasm has function-local variables. Every function declares at the start which local variables it has and of which type they are. These local variables are allocated in registers on Scale. The attribution of registers is simply made by incrementing a global counter in the compiler. In addition we have a global hashmap to save which register corresponds to which local variable in Wasm, in order to re-use them. The key of this hashmap is the unique local variable identifier and the value is the register itself. As the registers are suffixed by a number it’s pretty easy to use this global counter to assign new ones.
Here is an illustration to help you understand how WaSca works, following by a written explanation:
So if we take back our previous example program and try to emulate the behavior of the WaSca stack, here is how it works (so basically when we mention a stack in this explanation we talk about the WaSca stack) :
- We allocate the constant value 5 on the stack
- We pop this constant value from the stack and create the asm instruction to load the constant into a register. In the meantime a new register is assigned for this instruction.
- The new assigned register is now pushed on the stack
- We repeat the same behavior for another constant value 2
- Then when we have the
i32.addinstruction in Wasm. We pop the two last elements from the WaSca stack. They are not constant 5 and 2, but the registers where the value 5 and 2 are stored. We can now generate the correct Scale instructionaddint r3, r0, r1,r0andr1are registers popped from the stack and created when we assigned constant values 5 and 2. Andr3is a new register assigned to receive the response of the add operation. - Finally we push the target register
r3into the stack to be able to use it in the next instruction.
Types
When introducing the local variables we talked about types. Unfortunately Wasm only has i32, i64, f32, f64 and v128 types available to us. Technically there is also an externref meta-type that allows the code to create its own types, but LLVM (and thus Rust) does not support this yet. This means that we have five Wasm types that could map to the five Scale register types. The Wasm types have no relation with the types used in Rust, beyond that i64 in Rust is also i64 in Wasm. Unfortunately u64 in Rust is also i64 in Wasm, and the only way to notice this in Wasm is through unsigned operations on i64. To make secure multi-party computation (sMPC) we need some specific types. We therefore encode all operations on types as ClearModp, SecretModp, i64 and SecretI64 in Rust as operations on v128, f32, i64, f64 respectively in Wasm, without ever using them as such.
The reason why it works is that we do not translate operations on these types to Wasm operations, but instead translate them to external function calls which we will lower directly in WaSca to the appropriate Scale asm instructions.
To be continued…
Here is the second blog post in this series just to show you a first look at our compiler internals. The next post will be more focused about how we manage memory allocations in the WaSca compiler. We’ll see that we have encrypted memory zones in addition to the classical heap and stack, and explore together what we have already done and what we plan for the future. If you want to be notified when it’s out, feel free to follow Cosmian’s twitter account.
