Rust Enum Layout 的优化

今天学到了一点关于 Rust Enum 的冷知识,在开始阅读之前,大家可以猜一下下面的 Rust 代码在常见的 64 bit 机器上的输出是什么?

1
2
3
4
5
6
7
8
struct A (i64, i8);
struct B (i64, i8, bool);
fn main() {
dbg!(std::mem::size_of::<A>());
dbg!(std::mem::size_of::<Option<A>>());
dbg!(std::mem::size_of::<B>());
dbg!(std::mem::size_of::<Option<B>>());
}

在这个 Rust Playground 里可以看到结果。


Rust enum 本质是一种 tagged union,对应代数数据类型中的 sum type,这里不过多展开。在 Rust enum 的实现中,通常用一个 byte 来存储 type tag(大部分 enum 不会超过 256 种类型,更多地会相应扩展),也就是说,理想情况下,以下两个结构体是等价的:

1
2
3
4
enum Attr {
Color(u8, u8, u8),
Shape(u16, u16),
}
1
2
3
4
5
6
7
8
9
struct Color { uint8_t r, g, b };
struct Shape { uint16_t w, h };
struct Attr {
uint8_t tag;
union {
Color color;
Shape shape;
}
}

在这个实现下,enum 在很多场景下并不是 zero overhead 的。幸运的是,Rust 从未定义过 ABI,而带有数据的 enum 甚至是无法被 repr(C) 表示的,这给了 Rust 充分的空间对 enum memory layout 进行细粒度的优化。这篇文章会涉及一些在 rustc 1.60.0-nightly 下相关优化的介绍。

在开始具体的探索之前,我们需要准备一个辅助函数,来帮我我们查看变量的内存结构:

1
2
3
4
5
fn print_memory<T>(v: &T) {
println!("{:?}", unsafe {
core::slice::from_raw_parts(v as *const _ as *const u8, std::mem::size_of_val(v))
})
}

以上面的 Attr 为例:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
print_memory(&Attr::Color(1, 2, 3));
// [0, 1, 2, 3, 0, 0]
// ^ ^ ^ ^ ^^^^
// | | | | |
// | | | | --- padding
// | | | -------- b
// | | |---------- g
// | |------------- r
// ----------------- tag

print_memory(&Attr::Shape(257, 258));
// [1, 0, 1, 1, 2, 1]
// ^ ^ ^^^^ ^^^^
// | | | |
// | | | --- h, 256 + 2
// | | --------- w, 256 + 1
// | |------------- padding
// ----------------- tag

Option<P<T>>

P 是常见的智能指针类型,包括 &/&mut/Box。这应该是关于 enum layout 优化里最著名的一个例子了。Rust 推荐使用 Option<P<T>> 来处理可空指针,这实现了 null safety.

Option<T> 在 rust 中被表示为一种 enum:

1
2
3
4
pub enum Option<T> {
None,
Some(T),
}

如果不作任何优化的话,显然是存在不必要的 overhead 的,空指针可以完整地表示 None 的语义。由于这种情况太过常见,rustc 不仅针对性地做了优化,而且将其标准化了。

If T is an FFI-safe non-nullable pointer type, Option is guaranteed to have the same layout and ABI as T and is therefore also FFI-safe. As of this writing, this covers &, &mut, and function pointers, all of which can never be null.

1
2
3
4
5
let p = Box::new(0u64);
print_memory(&p);
// [208, 185, 38, 40, 162, 85, 0, 0]
print_memory(&Some(p));
// [208, 185, 38, 40, 162, 85, 0, 0]

一个不算太冷的冷知识是,这种 hack 并不是针对 Option 的,而是针对指针类型的。任何自定义的 enum 满足条件也可以达到相同的效果。

1
2
3
4
5
6
7
8
9
10
11
12
enum MyOption<T> {
MySome(T),
MyNone,
}

let p = Box::new(0u64);
print_memory(&MyOption::MySome(p));
// [208, 185, 38, 40, 162, 85, 0, 0]
// The address of `p`.
print_memory(&MyOption::<Box<u64>>::MyNone);
// [0, 0, 0, 0, 0, 0, 0, 0]
// Use nullptr to represent `MyNone`.

Option<P<T>> 可以优化的根本原因是,P 的内存表示下有一个永远非法的值,而相应的 enum 仅需要表示一个额外的值来表达多余的类型。超出这个约束就会导致这个优化失效。

1
2
3
4
5
6
7
8
9
10
11
enum MyOption2<T> {
MySome(T),
MyNone,
MyNone2,
}

let p = Box::new(0u64);
print_memory(&MyOption2::MySome(p));
// [0, 0, 0, 0, 0, 0, 0, 0, 208, 185, 38, 40, 162, 85, 0, 0]
print_memory(&MyOption2::<Box<u64>>::MyNone2);
// [2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

bool, Ordering

rust 中的 bool 占用一个 byte,且仅有两个合法的值,TrueFalse,对应的内存表示为 1u080u08。我们可以理解为 bool 有 254 个非法值可以供 type tag 挥霍。

1
2
3
4
5
6
print_memory(&Some(false));
// [0]
print_memory(&Some(true));
// [1]
print_memory(&(None as Option<bool>));
// [2]

更进一步地,我们可以更给力一点:

1
2
3
4
5
6
7
8
print_memory(&Some(Some(false)));
// [0]
print_memory(&Some(Some(true)));
// [1]
print_memory(&(Some(None) as Option<Option<bool>>));
// [2]
print_memory(&(None as Option<Option<bool>>));
// [3]

对应的,Ordering 有三种合法值,同样适用于这个优化。

1
2
3
4
5
6
7
8
print_memory(&Some(std::cmp::Ordering::Less));
// [255]
print_memory(&Some(std::cmp::Ordering::Greater));
// [1]
print_memory(&Some(std::cmp::Ordering::Equal));
// [0]
print_memory(&(None as Option<std::cmp::Ordering>));
// [2]

Enum

事实上,编译器并没有对 bool、Ordering 进行特判,任何种类少于 256 的 enum 本身都满足被优化的条件,即 type tag 里会有 (256 - kinds) 个空位。

1
2
3
4
5
6
7
8
9
enum ShapeKind { Square, Circle }
print_memory(&Some(ShapeKind::Square));
// [0]

print_memory(&MyOption::MySome(MyOption::MySome(1u8)));
// [0, 1]
print_memory(&(MyOption::MyNone as MyOption<MyOption<u8>>));
// [2, 0]
// 尽管 u8 本身没有空位,但是 MyOption 的 type tag 有 254 个空位,因此外层的 MyOption 的 typetag 被优化掉了。

Struct、Tuple

Struct、Tuple 等都属于 Product type。在实现中,往往是将所有字段依次存下来,并做额外的 padding。那么一个理所当然的优化是,如果 struct 的其中一个字段有空位,那么就可以将 enum tag 塞进去。

1
2
3
4
5
6
7
struct A1 (i8, bool);
print_memory(&Some(A1(1, false)));
// [1, 0]
print_memory(&Some(A1(1, true)));
// [1, 1]
print_memory(&(None as Option<A1>));
// [0, 2]

我们再回到文章开头的例子。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
struct A (i64, i8);
struct B (i64, i8, bool);

dbg!(std::mem::size_of::<A>());
// 16
dbg!(std::mem::size_of::<Option<A>>());
// 24
dbg!(std::mem::size_of::<B>());
// 16
dbg!(std::mem::size_of::<Option<B>>());
// 16

print_memory(&Some(A(1, 1)));
// [1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
// ^ ^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^^^ ^ ^^^^^^^^^^^^^^^^^^^
// | | | | |
// | | | | ------------ padding
// | | | ------------------------ .1
// | | ------------------------------------- .0
// | ------------------------------------------------------------ padding
// ------------------------------------------------------------------------ tag

A 和 B 由于 padding,都需要占用 16 个 byte,而 Option<B> 由于存在一个 bool 字段 .2,tag 被优化进 bool 了,因此也只需要 16 个 byte。反而 Option<A> 实打实地用了 24 个 byte。

一个很容易想到的优化是,使用 padding 中未定义的内存来存储 type tag。比较遗憾的是,A 的 layout 是在编译 A 自身时确定的,而 Option<A>A 的 padding 中存储的数据是未定义行为。这也导致了一个比较滑稽的结果,多存了一个字段,Option 占用的空间反而减少了。

当然,使用 padding 存储数据是完全可能的,但前提是不能影响子数据结构的 memory layout。如果我们将 AOption手动展开

1
2
3
4
5
6
use std::mem::size_of;
struct A(i64, i8);
enum OptionA { Some(i64, i8), None }
dbg!(size_of::<A>()); // 16
dbg!(size_of::<Option<A>>()); // 24
dbg!(size_of::<OptionA>()); // 16

这种情况下,由于 OptionA 的数据直接保存在 Some 内,实例化的时候完全可以使用 padding 存储 type tag,而不会引起潜在的未定义行为。

优化自定义结构的可能性

目前,所有 enum layout 相关的优化都适合由编译器针对特定的类型进行 hack 来实现的,我们无法自己控制我们自定义的 struct 在 enum 中的 layout。为了优化一些常用场景,rust 又提供了 NonZero* 等辅助结构体,用来表示非 0 的整数,与此同时编译器会让 size_of::<Option<NonZeroU8>> == size_of::<u8>。但这只能由标准库 case by case 处理,而真实的需求是非常复杂的,比如有时候我们可能需要 NonMaxU64,或者例如使用了第三方库的 Option<ordered_float::NotNan<f64>> 就无法被优化。

针对自定义接口的优化需要引入非常复杂的机制,在编译期告诉编译器一个类型非法的内存结构有哪些。我目前感觉一个可能的实现是 const trait + const iterator,给编译器提供潜在的非法值的迭代器。不过目前没有看到相关的 RFC。

这篇文章所有的 example 可以在 Rust Playground 找到。