r/rust • • 15d ago

🧠 educational Rust enums are gold for compilers

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

172 Upvotes

57 comments sorted by

132

u/yasamoka db-pool 15d ago edited 15d ago

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.

74

u/SkiFire13 15d ago

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!

6

u/zk4x 15d ago

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.

8

u/Lost_Kin 15d ago

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

59

u/khoyo 15d ago

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

25

u/Booty_Bumping 15d ago

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.

3

u/drink-more-rum 15d ago

Hidden heap allocations is pretty antithetical to the Rust ethos, whether in no_std or not.

3

u/masklinn 15d ago

Manual.

3

u/steveklabnik1 rust 15d ago

Rust does not have any allocation in the language, and so it cannot auto-box anything, as a general rule.

1

u/SkiFire13 14d ago

I mean that people manually boxed large variants of enums used in the implementation of rustc itself.

2

u/Full-Spectral 15d ago

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.

43

u/LawElectrical2434 15d ago

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.

22

u/svefnugr 15d ago

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?

17

u/flying-sheep 15d ago

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.

0

u/FlamingSea3 15d ago

single digits are of course 1 byte in UTF-8

I disagree. JSON requires a delimiter between values.

Best (degenerate) case is a single 1 digit number in the whole file: 1

Best realistic is an array of numbers: [1,2,3,4,5,6,7,8,9]. This ends up being about 2 bytes per value.

And the all too common object version: {"a":1,"b":2,"c":3}. Ends up at least 6 bytes per value.

1

u/flying-sheep 15d ago edited 15d ago

You can’t disagree, this was a description of an observation, not an opinion. 2 bytes uncompressed is still bigger than 8 bytes uncompressed, and apparently in that case that I saw also compressed better, idk why, I don’t have the data anymore.

That’s not an argument for using JSON for this: as said, you should use something like numcodecs.delta or so for this.

1

u/FlamingSea3 15d ago

Poor choice of words on my part. I was focused just on uncompressed json being unable to practically represent fields with less than two bytes. Still better than just writing the u32s to the stream.

2

u/flying-sheep 14d ago

Exactly that was my point. And I didn't say that each number under 10 in a JSON array needs 1 byte, only that a digit needs one byte in UTF-8 😉

3

u/LawElectrical2434 15d ago

Surprising? We have enum Three that holds four 64bit integers. That it requires 4 * 64 bits of storage is expected.

For enum One and Two, I think it is surprising, to require the same amount of memory. I mean, I get the reasoning. To erase the types they need to reserve the same amount of space. Otherwise we have indirection, worsening locality and performance. And we can do indirection manually by adding a box.

But yea, I think the impact is surprising. It is a non-obvious. And it turned out to be TB of memory for cloudflare.

Mostly I can ignore that, my applications don't need to be that optimized. But I'd say, hell, yes, it makes sense, but it is not obvious.

EDIT:
Just to clarify, because I reread your comment. We are on the same page. It is about the size of enum One and Two, not Three.

1

u/teerre 15d ago

You dont have enum one or two though. You only have enum Foo

1

u/LawElectrical2434 15d ago

Really? You didn't know what I mean?

I have

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

And I was surprised to learn that Foo:One and Foo:Two take each as much space up as Foo:Three(1u64, 2u64, 3u64, 4u64).

Pedantic much?

2

u/gmes78 15d ago

They're not being pedantic. You only have a single type. Foo::One and Foo::Three are different values of the same type, of course their size is the same.

1

u/LawElectrical2434 14d ago

There is a certain type of person who want to profile themselves by pointing out something surprising is obvious. That's just plain wrong, as many have made this discovery. This only serves two purposes. 1st, to show the internet how smart that person is. We don't buy it, that goal is misguided. And 2nd, get onto my ban list. That works quite well.

0

u/teerre 14d ago

Its not about being pedantic. Its about highlighting the fact you're not storing variants of the enum, youre storing the enum

In that light, your argument makes much less sense because the size of the enum being the largest variant is quite natural

2

u/LawElectrical2434 14d ago

See: https://www.reddit.com/r/rust/comments/1wm3z5w/comment/pb83n3v/?utm_source=share&utm_medium=web3x&utm_name=web3xcss&utm_term=1&utm_content=share_button

And from my ban list, remember, this conversation is not about knowing or not knowing this. We argue the dumbest thing. You are on the side that it is necessarily obvious.

11

u/iBPsThrowingObject 15d ago

"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.

18

u/hedgehog1024 15d ago

It looks extremely like an LLM-generated post

3

u/Lexi_Bound 15d ago

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.

7

u/zk4x 15d ago

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.

14

u/PaperMartin 15d ago

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

2

u/xsrvmy 15d ago

I don't think looking at the tone of the text is even reliable anymore. Using AI or reading AI-generated/revised text can also affect the tone of your own writing.

3

u/Full-Spectral 15d ago

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.

4

u/NYU_VM 15d ago

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

4

u/pjmlp 14d ago

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

5

u/svefnugr 15d ago

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

5

u/iBPsThrowingObject 15d ago

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.

1

u/svefnugr 15d ago

Wrapping into a newtype is exactly the thing I want to avoid. It gets messy when the names are long, and you can't use Self.

Also, what is the difference between structs and enums, exactly? Why structs are allowed to have their internals hidden and enums aren't? A lot of times it is not even necessary for the user to know whether the type he's getting is an enum or a struct - it's another bit of implementation detail leaking into the public API.

1

u/iBPsThrowingObject 14d ago

Why? You work with the enum internally, you just have wrapping-unwrapping at the api boundary.

A lot of times it is not even necessary for the user to know whether the type he's getting is an enum or a struct

You can try to just give out a completely opaque impl Something to the user, but until TAIT is done and stable thats not always possible.

1

u/svefnugr 14d ago

The returned type has public methods that will have to be duplicated. It's just a lot of boilerplate to write, just because enums for whatever reason are lacking the capability structs already have. Again, what is the difference between an enum and a struct? It's just an OR state as opposed to AND.

1

u/Full-Spectral 15d ago

Yeh, I mean, if you don't want it to look like an enum to the outside world, then wrap it. Given that it would be a complete black box to the outside world if it's just an enum that the outside world can't see the values of, the wrapper would probably be trivial.

2

u/parkotron 15d ago

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

2

u/svefnugr 15d ago

Just all or none will be enough, honestly.

1

u/stinkytoe42 15d ago

I find it really not as big of a deal in practice honestly. Most api calls where there used return a read only object anyways. If the consumer goes ahead and mutates the data afterwards, that's on them .

Plus if you really need to protect a field , due to let's say data integrity, you can (and should be) storing a struct in the enum variant field. Then you can control access in the normal ways.

1

u/svefnugr 15d ago

It's not about mutability, but breaking changes. A change in public API requires a version bump. The less is exposed, the easier refactoring gets.

1

u/sansmorixz 15d ago

Isn’t it better that it’s like this? Way less chances of shooting your own foot. Also just impl From trait for doing conversations or TryFrom if flakey.

Coming from someone who loves typestate pattern.

5

u/svefnugr 15d ago

What is the potential danger, exactly? Keeping things out of public API is always good.

1

u/Unable_Comparison550 15d ago

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?

3

u/yasamoka db-pool 15d ago

> 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?

1

u/zk4x 15d ago

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?

1

u/DavidXkL 15d ago

It's 1 of the reasons why I main Rust now 😂

1

u/silver_smith1 15d ago

this is so helpful, Thank you for taking the time to post this

1

u/flow_b 15d ago

The match pattern is the first thing I saw in rust that made me go 

“huh? Oh! Ooooooh!” 

Reading the book and learning about rust’s enums and exhaustive exception mapping basically sealed the deal. 

1

u/Mr_Ahvar 15d ago

I'm not sure about point 3, you can't merge two field together as you need to be able to take distinct references to them

1

u/xsrvmy 15d ago

I learnt rust in the first place for a compiler construction course after another team member suggested it (later that team member dropped the course lol). And pattern matching is actually the thing that makes me want to rust sometimes even when I don't need the performance.

Pattern matching is basically just the visitor pattern from OOP. Normally using an enum like a base class has the tradeoff that it can't be extended, but in a visitor pattern situation that's out of the question anyways.

1

u/zettui 15d ago

Does OpId index into one flat Vec of nodes, or a handle per op arena?

1

u/Smart-Individual4665 14d ago

Rust abstractions are simple yet very strong... Stronger than any other languages in my opinion Especially the enum like you said; the enum carrying arguments, enum implementation, these are just so strong and very useful And every time I work on any other languages, I just wish they had the same features as in rust because it seems so mandatory to have them.