alloca in Rust using inline ASM

• go back
• pt.1 <
• pt.2

If you need a debugger, I highly suggest writing your own.

inception

The stack is like a stack of plates. You can't change things in the middle or the whole thing falls down and you need to know the size of plates you want to remove.

Usually the compiler ensures that the sizes are known at compile-time, but we can store variable-size data on the stack by keeping track of the size at runtime so we know how much to pop.

In C there is an alloca function that can allocate variable-size data, but no one really talks about it because it's bad practice!

But hold on, that function doesn't exist in the Rust std library! There is a library that implements alloca in Rust but it just links to C which is no fun. We want alloca in Pure Rust™.

learning asm

Here's what we're going to do. We're going to move the stack pointer down to allocate some space, and then we're going to move the stack pointer back up to deallocate the space. Pretty crazy stuff right. (Remember the stack grows downwards)

unsafe {
    asm!("sub rsp, 8"); // allocate 8 bytes
    do_some_stuff();
    asm!("add rsp, 8"); // deallocate 8 bytes
}
        

Hmm, that's definitely going to be a problem. If you look back up at the code we wrote you can see the stack pointer isn't restored at the end of the asm!("sub rsp, 8") block. That's going to cause weird bugs and undefined behaviour if we aren't careful.

unsafe {
    asm!(
        "sub rsp, 8",
        "add rsp, 8",
    )
}
        

Perfect! Now we've allocated and deallocated stack memory, but how do we actually use it? As soon as the block exits, the stack memory is already deallocated, that's not very useful.

Let's make the assembly call a function! To call a function in assembly we can use the call instruction.

unsafe {
    extern "C" fn my_function() {}

    asm! {
        "sub rsp, 8",
        "call {f}",
        "add rsp, 8",

        f = in(reg) my_function,
    }
}
        

We have to specify the function pointer as an input register with f = in(reg) my_function. You can find more information about setting input/output registers in the Rust reference book. But basically, you tell the compiler what the asm block receives with in(reg) and what it returns with out(reg).

Ok I lied, it isn't *that* easy to call a function in assembly. First of all, on certain CPU architectures the stack pointer has to be a multiple of 16 before we jump to a function (known as stack-alignment). And second of all, we need to tell the compiler what registers our code modifies (known as register clobbering).

Since we are calling an extern "C" function, we are going to be clobbering any registers that the extern "C" function would. We can tell the compiler that we are clobbering registers in a C way with the clobber_abi("C") flag. Let's change the stack pointer to be a multiple of 16 too.

unsafe {
    extern "C" fn my_function() {}

    asm! {
        "sub rsp, 16", // change to 16
        "call {f}",
        "add rsp, 16", // change to 16

        f = in(reg) my_function,
        clobber_abi("C"), // specify C ABI
    }
}
        

Now we are correctly telling the compiler what registers get clobbered and we are ensuring the stack pointer is a multiple of 16. Great!

There is still an issue though... we still can't use our allocated data! The function has no parameters! It needs to take a pointer to the data we allocated so that we can use it.

In the C ABI the first parameter of a function is passed in the rdi register. So let's set the rdi register to a pointer so that the function can use it.

unsafe {
    extern "C" fn my_function(ptr: *mut c_void) {}

    asm! {
        "sub rsp, 16",
        "mov rdi, rsp", // set first parameter to rsp
        "call {f}",
        "add rsp, 16",

        f = in(reg) my_function,
        clobber_abi("C"),
    }
}
        

And that's it! Assembly is wildly unsafe, especially because Rust expects your program to "follow all the rules". There are likely tons of edge-cases I missed that would cause weird obscure heisenbugs. I really wouldn't use this in production code but it is a fun exercise. I hope you learnt something reading this because I learnt a lot making it.

final thoughts

There's one detail I didn't cover and that's making alloca work for any closure and not just extern "C" functions. You can see the full source code here. I made this a while ago but I thought I may aswell turn it into a blog post.

P.S If you find any of those aforementioned edge-cases that I missed, shoot me a message johnnycambodia09@gmail.com