Using Rust to Build a JVM

A while back, I started my project java_rs, with the eventual goal of making it into a JVM (Java Virtual Machine). This blog post is about a few of the choices I made writing it and the effect that Rust's safety guarantees had on the final product.

I assume you have basic knowledge about Java, and roughly Bachelor's-level knowledge of Computer Science. I'll try to explain the Rust as we go, but unless you're already familiar, it may be difficult to follow.

Benefits of Rust

java_rs is written in Rust, and anyone who's programmed in Rust before will know how difficult it is to represent reference cycles. Even the extemely common doubly-linked list requires quite the explanation. Given that objects in the JVM can have arbitrary references to other objects, including themselves, I was expecting a challenge.

First off, I don't think you could actually represent JVM memory using Rust's references, so objects in java_rs are wrapped by Arcs and RwLocks. If you're unfamiliar, Arc is the Rust stdlib's version of a thread-safe reference-counted pointer, and RwLock is its one-write-many-read synchronization lock (like a Mutex, but more than one thread can read at a time).

While it may seem that these types would negate many of Rust's benefits, they don't; as long as you don't use unsafe code, there are still a lot of guarantees about what you can and can't do. For example, let's try to fail to release our lock on a RwLock:

use std::sync::*;

fn test(val: RwLock<i32>) {
    val.write().unwrap()
}

Predictably, we get an error at val.write().unwrap() telling us the compiler expected a () but got a RwLockWriteGuard. Helpfully, it gives us a suggestion with the exact return type it's getting:

use std::sync::*;

fn test(val: RwLock<i32>) -> std::sync::RwLockWriteGuard<'_, i32> {
    val.write().unwrap()
}

Trying again, we get another error. It's pointing to '_, the lifetime argument on our new return type. It says "this function's return type contains a borrowed value with an elided lifetime, but the lifetime cannot be derived from the arguments". That means the compiler doesn't have enough information to tell what lifetime we mean in our return type, so we have to specify. Again, it has a helpful suggestion that maybe we want to use the 'static lifetime, which is a complicated topic but basically means "this lives forever".

In a typical program, you rarely want to use 'static except when dealing with compile-time constants like string literals. We're trying to leak our hold on a lock though, so let's go for it:

use std::sync::*;

fn test(val: RwLock<i32>) -> std::sync::RwLockWriteGuard<'static, i32> {
    val.write().unwrap()
}

Another error! And this time, we get to the heart of the issue. The compiler says "cannot return value referencing function parameter `val`" and that our function "returns a value referencing data owned by the current function". Finally, it lets us know that we're borrowing the data when we call val.write().

If you've used Rust before, that will probably make good sense, but if you haven't it's probably nonsense. The exact reason this happens requires more space to explain than I care to use here, so in short, the issue is that the write() function references data inside our RwLock, but the RwLock itself isn't returned along with the reference. That means that when the function returns, and val is dropped (read: deallocated/deleted), the reference would no longer be valid. And rust hates invalid references.

That's not to say you can't return a RwLockWriteGuard safely; for example, this code works just fine:

use std::sync::*;

fn test(val: &RwLock<i32>) -> std::sync::RwLockWriteGuard<'_, i32> {
    val.write().unwrap()
}

There's a bit of compiler magic going on here with elided references, but the idea is this time borrow of val occurs outside the function, meaning that borrow can live longer. The compiler sees that the code will work just fine as long as '_ is at most as long as our borrow of val. We can see the relationship between the lifetime of the borrow in val and the '_ is important by defining them explicitly:

use std::sync::*;

fn test<'a, 'b>(val: &'a RwLock<i32>) -> std::sync::RwLockWriteGuard<'b, i32> {
    val.write().unwrap()
}

This code fails to compile, telling us that 'b cannot outlive 'a so that the reference we return can't outlive the value it points to. That was possibly too much time to spend on such a small detail, but you hopefully now have some kind of understanding of the benefits Rust's extensive compile-time restrictions can have on complex and multi-threaded software like a JVM.

Representing Primitives and References

Now, getting to my code, let's take a look at JavaType, which encompasses any value a variable may have in JVM bytecode:

pub enum JavaType {
    Boolean(bool),
    Byte(i8),
    Short(i16),
    Char(u16),
    Int(i32),
    Float(f32),
    Long(i64),
    Double(f64),
    Object { class: ClassRef, fields: HashMap<String, JavaType> },
    Array { class: ClassRef, data: Box<[JavaType]> },
    Reference { class: ClassRef, val: Arc<RwLock<JavaType>> },
    Null
}

There are more design decisions happening here than you might think at first. Most interesting, in my opinion, is the decision to include the Object and Array variants. These variants contain the actual data contained by object and array instances, which is something JVM bytecode can never operate on - it can only use references. My reasoning is that having a separate type for instances would require 1. representing arrays as objects directly, 2. having separate reference types for objects and arrays, or 3. having different array types for primitives vs. objects.

Initially, I did actually represent arrays as objects with a special instance field containing their data, but I found it to be unnecessarily clunky and not type-safe enough for my liking.

As a short aside, programmers who use Rust but not Java may be confused by the fact that the array variant uses a Box containing an array rather than a Vec. As it stands, if you wanted to change the length of the array you'd have to do all the shifting of values around yourself - you'd never have to, though, because arrays in Java are fixed-length.

The last interesting design decision I'll mention here is that references are just Arc<RwLock>s. The use of Arcs instead of a bespoke type managed by the garbage collector means the GC has to keep record of, and deallocate manually at collection times, every object ever allocated. That's not great for performance, since there are a lot of performance benefits to be had with better design, but I'm not particularly interested in performance in this project, so I kept it simple.

Representing Shared Class Information

Next up, let's check out ClassRef, a type that showed up several times in JavaType:

pub type ClassRef = &'static Class;

ClassRef is just a convenience type alias for a reference to a Class with a static lifetime. I'm mainly bringing this up and not skipping directly to the definition of Class because the use of a static reference here has a few important implications. First is that we'll basically never deallocate a Class, but that's not a big deal since there's few circumstances you would actually want to. Second is that the inside of a Class will be immutable pretty much immediately after construction. That's not immediately an issue, but it will be soon:

    pub struct Class {
    pub minor_version: u16,
    pub major_version: u16,
    pub constant_pool: RuntimeConstantPool,
    pub access_flags: u16,
    pub name: String,
    /// may be `None` if and only if this is a primitive class or java/lang/Object
    pub super_class: Option<ClassRef>,
    pub interfaces: Vec<ClassRef>,
    pub fields: HashMap<String, Field>,
    pub instance_fields: Vec<InstanceFieldInfo>,
    pub methods: HashMap<String, &'static Method>,
    pub attributes: Vec<Attribute>,
    pub array_inner: Option<ClassRef>,
    /// whether the <clinit> method has been called for this type
    pub is_initialized: Mutex<ClassInitStatus>,
    }

Most of this is pretty unremarkable, but the eagle-eyed may be seeing the issue with Classes being immutable. Classes, themselves, contain ClassRefs, but circular references are possible. How, then, can we possibly use plain old ClassRefs that don't have interior mutability, like a Mutex<ClassRef> or RwLock<ClassRef>? The answer is unsafe Rust.

Now that you're done screaming in horror, let me explain. Despite its name, "unsafe Rust" is not necessarily actually unsafe. The deal is just that you have to make sure any invariants the compiler expects to be upheld are upheld. To explain, let's take a look at how I mutate ClassRefs in the code:

    unsafe {
    let c = std::mem::transmute::<ClassRef, *mut Class>(cx.0);
    (&mut *c).initialize(&cx.1, cx.0)?;
    }

The first unsafe part of what I'm doing here is telling Rust to reinterpret cx.0 as a *mut Class instead of as a &'static Class. (*mut T just means a mutable raw pointer to T). It's also unsafe to dereference a raw pointer, mutable or otherwise, but since it's coming directly from a reference we at least know it's valid (i.e. not null, not dangling, and actually containing the kind of data we expect).

By turning an (immutable) reference into a mutable raw pointer, you make a promise with the compiler not to do things it doesn't expect, the same way you promise a C++ compiler you won't use a pointer after you've freed the data it points to. If you break those promises, the result is undefined behavior. The most obvious invariant to uphold is that no other thread may attempt to read from the pointed-at data while you mutate it. We take care of that by making sure 1. Classes aren't used until they're initialized, and 2. initialization is single threaded.

My code actually does fail to avoid undefined behavior though; it is technically undefined behavior to mutate data reached through a shared reference. I believe this is undefined behavior in Rust primarily because of possible noalias optimizations, but I'm not sure. Discussion of noalias and the exact results of the undefined behavior here are too much for this post, but the main thing that's important is that the generated code is very probably correct in this case, and even if it's not, it's not a big deal for a toy JVM. As a final note here, if it turned out to be a problem, you could solve it by using an UnsafeCell.

Now, for completeness' sake, I'll address the Field type you probably noticed if you were reading the code. The type isn't anything special - it's just a bunch of information about the field, such as its type and name, along with an Arc<RwLock<JavaType>> holding the actual value of the field.

The VM Stack

The same as just about every modern language, Java uses a separate stack and heap, where the stack holds most function-local data that isn't dynamically allocated. Unlike typical assembly, which typically manipulates a stack pointer directly with instructions like sub sp, 20, the JVM's stack is allocated only declaritavely in the binary by giving a maximum number of local variables and a maximum size of the operand stack (the JVM is a stack machine). The specification leaves it completely up to the implementation to decide how it handles "stack" data like local variables and return addresses.

If you were implementing a JVM and you were ok with using assembly (in a JIT or otherwise), you'd probably do the allocation of the JVM stack frames the exact same way you allocate other space on the stack, e.g. sub sp, words_needed. However, I've never seen a non-assembly language allow you to allocate an array on the stack unless its length is a compile-time constant, and that includes Rust. That means we can't allocate our entire JVM stack on our implementing language's stack without using something like inline assembly.

There is the option to use a fixed-size array on the stack for a portion of the JVM stack, and only heap allocating if you exceed your predefined size limit, but that's a little too much of an optimization for my toy JVM, and I'm not sure if you'd even get a performance boost on average. As for inline assembly, Rust does actually have impressive support for it, but again that's far beyond the (current) scope of this JVM. Instead of either complex option, the stack frame is just dynamically allocated:

pub struct StackFrame {
    pub current_method: &'static Method,
    pub this: Option<JavaType>,
    pub pc: usize,
    pub stack: Vec<JavaType>,
    pub locals: Vec<JavaType>,
    pub is_native: bool
}

Technically Rust is fully within its rights to allocate Vecs on the stack since Rust is actually capable of determining whether the vec will ever leave the scope of the function it's allocated in (or even moving it to another stack frame when ownership is transferred), but that's not really important.

Conclusion

One of the reasons I wrote java_rs was because I was curious what kind of impact Rust's safety guarantees might have on a piece of software as complex and orthogonal to Rust's strengths as a JVM. Unfortunately, I think the answer is "not much". For a small, one-person project, Rust undoubtebly prevented a huge number of bugs that I otherwise would've happily coded in, but I don't think that would be a real benefit to a team making a JVM they actually indend for people to use.

The combination of coding around a garbage collector and coding around bytecode means that you're actually looking at very few locations where memory is (de)allocated. In a simple JVM, allocations only really happen in three main locations:

1. During the loading of new classes 2. During the start of a function (the stack) 3. In a very small number of opcodes that directly call for it

Compared to other programs similar in their amount of complexity but different in their kind of complexity, that leaves very little room for Rust to shine. If I were to do it over again, and write another simple JVM, would I use Rust? Of course! But that's because I've already spent a ton of time in Rust; I'm used to wrangling with the borrow checker and using the trait system. Further, I find Rust's general style to be very nice, especially its enums.

If someone were to ask me whether they should pick up Rust specifically to code up a VM, I'd tell them to probably just stick with what they know. I love Rust, but it turns out this problem domain just doesn't benefit much from its innovations.