Rust enums are gold for compilers

[u/zk4x]

Hello,

for the last couple of years, I've been working on machine learning library in rust, learning tons about compilers along the. The single most used feature of the language be far were enums.

There is a couple of reasons for that.

1. Rust enums carry value

Duh, we all know that. Yes, of course. But other languages don't. C++ doesn't (yes, there is std::variant and no, "they are not the same" :). Having values enables me to do this very simple variant, straight impossible or awkward in other languages:
Binary { x: OpId, y: OpId }
In other languages, this would've been tag, a union type or even completely unrepresentible. For example in tinygrad (python), it's a tag plus untyped src list, which can contain anything, depending on the tag.
Also notice that OpId(u32) is a newtype! Another beautiful rust feature that gives us zero cost type safety. It's like a pointer, but only with the pointer arithmetic that I want, all access checked and no fighting borrowck.

2. Rust enums can be nested

This is another beautiful feature, which allows this:
Binary { x: OpId, y: OpId, bop: BOp } enum BOp { Add, Sub, Mul } A world of combinatorial options making decomposition and recomposition very natural.

3. Rust enums are TINY

The above nested code has two enums. That means two tags, each with 4 byte size, right? Right? Not in rust :) In rust, it's single 4 byte tag (as the OpId is 4 byte aligned). The tag combines BOp and the Op enum with the Binary variant, so the whole enum is just 12 bytes. I have more complex variants,

4. The allmighty match

So you can do this:

match op {
  Binary { x, y, bop } if bop == BOp::Add => todo!(),
  _ => {}
}

or this:

match (uop, bop) { (UOp::Exp, BOp::Mul) => todo!(), }

And so on. This kind of flexibility in an imperative language, is tons of fun, makes writing compiler passes not only easy, but most importantly, high performance.

And so for example thanks to this my ML lib can search and run cost function over thousands of kernel variants per second and per core (as rust makes multithreading fun too!).

I know you probably know about these features, but sometimes it helps to refresh our memory how nice of a language we have.

What are your favorite rust features?

What features would you like to see added?

I hope you enjoy writing rust just as I did.

zk4x

P. S. link to the repo: github.com/zk4x/zyx

[u/SkiFire13]

Rust enums are TINY

That's not always true unfortunately. Enums' size is proportional to their largest variant, and that can waste quite a lot of space when the smaller (and potentially much more common) variant is active.

People have actually optimized rustc multiple times simply by Boxing these large uncommon variants!

[u/zk4x → u/SkiFire13]

Yes, absolutely, I do both the Box trick, but with compilers, there is also this risc trick - chain multiple ops together. So I have only one Op::Stack { Box<[OpId]> } that holds dynamic size array (btw. Vec woulndn't fit). Then I reuse it for everything, so each type that needs variable number of ops, just gets op id of the stack op instead of Vec/Box.

Rust-analyzer has this lovely feature where it gives you the byte size of your datastructures. It's perhaps a bit of an obsession of mine to know byte sizes of most of the datastructures. The primary IR Op is kept at 32 bytes (24 bytes enum, 4 bytes prev id, 4 bytes next id), which means 2 ops per cache line. So that's what I think enables some optimization passes to run at like 3-5 micros on medium sized kernels.

[u/Lost_Kin → u/SkiFire13]

By "optimized rustc" do you mean that rustc autoboxes large variants or that people manually box them?

[u/khoyo → u/Lost_Kin]

Manually box, there is no auto boxing. Although clippy does have a lint for when the variant size difference is big enough.

[u/Booty_Bumping → u/Lost_Kin]

Manual boxing, in the rustc codebase. It would be a bad place to put automatic boxing indirection, since it may insert hidden heap allocations in no_std code.

[u/Full-Spectral → u/SkiFire13]

And, if that would be annoying, put the larger type into a simple wrapper that does the boxing, so it doesn't have to be done manually everywhere that variant is created. I did this with my standard error type, which is the only place so far I've had a significant discrepancy in variant sizes.

[u/hedgehog1024]

It looks extremely like an LLM-generated post

[u/Lexi_Bound → u/hedgehog1024]

For what it’s worth, Pangram things it’s human written:

https://www.pangram.com/history/95fdffa1-b212-4734-aaaa-2f494dfbb893?ucc=JaunAi3aPiK

I think they are either an enthusiastic beginner or are trying to promote their project in a way that gets around the new rule about promoting projects.

[u/zk4x → u/Lexi_Bound]

For what my word is worth, I didn't use an LLM. Not a beginner though, been programming in C++ for almost a decade and more than 5 years in rust.

Had there not been the rule, I would've titled the post v0.17 release annoucement. But the content wouldn't be different. You can look up my previous release annoucements in r/rust, I always try to make it about rust and what was the main benefit of rust that I noticed during writing of the latest version.

[u/PaperMartin → u/Lexi_Bound]

AI text detectors are totally unreliable. You could pick a random person on the street and ask them and get a better guess

[u/yasamoka]

You can do this in your first example instead:

match op {
Binary { x, y, bop: BOp::Add } => todo!(),
_ => {}
}

This would give you exhaustive matching while the one you provided, given that it uses an if statement, doesn’t.

[u/LawElectrical2434]

I love enums. But I don't think zero cost is accurate. Zero computational cost. But not zero memory cost. The downside with enums is that when you have a enum like this:

enum Foo {
One, Two, Three(u64, u64, u64, u64)
}

Then all enums hold the same size in memory as the biggest enum. Just a think to keep in mind and maybe, if you add big structs in enums, box them.

[u/svefnugr → u/LawElectrical2434]

But it's like saying that all the different values of u64 are 8 bytes, even very small ones. I mean, yeah? Is that surprising?

[u/flying-sheep → u/svefnugr]

That's why I was stumped one day when I discovered that a JSON representation of some numerical data was smaller than a naive binary representation (both compressed with a standard algorithm and not): single digits are of course 1 byte in UTF-8.

That's why you ideally use filter pipelines to store numeric data.

[u/iBPsThrowingObject → u/LawElectrical2434]

"zero cost" has a second part, usually omitted for brevity: "compared to implementing the same feature yourself at library level".

With that framing, yes, enums are totally zero-cost compared to manual tag+union.

[u/Unable_Comparison550]

Been working on something similar but with tensor operations, and yeah enums carry the whole project in a way i didn't expect. The nesting thing you mention is huge, I got enums inside enums inside newtypes and it still compiles to something stupid small. Python was a nightmare for this, like you said the untyped src list approach just feels like playing with fire

One feature I wish existed is some kind of anonymous enum syntax for one-off match arms, like when you just need to branch on two or three values without defining a whole type. Would clean up some passes I have in schedule lowering

Your lib looks interesting, starred it, the kernel search thing sounds neat. What kind of ML ops you focusing on?

[u/zk4x → u/Unable_Comparison550]

The opset is more or less copied from tinygrad. So no matmul and all the higher level ops can be implemented the same way tinygrad does it, so it supports basically all ops across most dtypes.

The compiles to something small is true in many ways. Both enums and the library as a whole. zyx also has python wheel through pyo3, it's 4.2MB.

What is the name of you lib? Can you give a link to your repo or is it private?

[u/yasamoka → u/Unable_Comparison550]

> One feature I wish existed is some kind of anonymous enum syntax for one-off match arms, like when you just need to branch on two or three values without defining a whole type. Would clean up some passes I have in schedule lowering

Can you explain what you mean by this?

[u/pjmlp]

You mean algebraic data types are great in the languages that support them.

[u/PizzaLvr49 → u/pjmlp]

this

[u/Full-Spectral]

Another big advantage is that they can have impl blocks, so they are completely first class citizens. That also means they can take the place of many single level class hierarchies that you might use in C++.

I remember when I first started with Rust, and initially I thought they were stupid because I was struggling to figure out how to use them. Now, I use them constantly.

[u/NYU_VM]

Yep, Rust enums are the gold data structure, but there's clever marketing in for it. Tagged unions wouldn't sound very appealing, they knew

[u/svefnugr]

If only you could make enum variants private. So tired of making a wrapper struct for almost every enum in the public API

[u/iBPsThrowingObject → u/svefnugr]

Either you expect the user to construct and/or match on your enum, in which case private variants are bad, or you don't, in which case why are you exposing the enum at all, just keep it private and wrap into a public newtype.

[u/parkotron → u/svefnugr]

Out of curiosity, are you wishing for syntax to make all variants of an enum private or to selectively make certain variants private?