feature image

2026年6月17日 | ブログ記事

Rustのunion型について | またはenumに危険な最適化を掛ける方法

Rust 低レイヤ
helgev icon helgev

Rust unsafeやFFIの日本語文献があんまりない！！！

Rustのunion型について、日本語文献で詳しいものがあまり見つからなかったので書きます。

まずは重要な注意点を:

enum型のことではありません。‌
ここのRustはunsafeです。‌
Rust初心者さん、あるいは、もし貴方が Rust とともに長く幸せな人生を歩みたいと望むなら、今すぐに背を向けて、このブログを見てしまったのを忘れるようお勧めします。
[Rust裏本より一部引用・改変]‌

しかし、それでもenumからサイズを1バイトでも削りたいと思うなら、union型の出番です。

少しでもヒープアクセスを減らし、パフォーマンスを上げたいのなら、union型が便利になる場面も来るかもしれません。

ここではsafe Rustの文法を知っていることが前提です。

unsafeに慣れていない方は、とりあえず未定義動作のことは漆黒の闇で、混沌としていて、危険で、何が起こるか分からない禁じられた領域であると思ってもらえれば大丈夫です。
unsafeは、無償の愛を以って未定義動作から守ってくれるコンパイラの盾を取り除き、あなたの責任でプログラムを守ることを意味します。

traP Blogは数ヶ月ぶりです‌‌
Helgevです

Rustをメイン言語として2年+α書いています

unsafe RustへはFFIプラグイン制作で必要になり参入しました

では本題↓

Rust union型

Rustのenum型の構造について

Rustのenum型は「タグ付き共用体」と呼ばれる構造でメモリに載っています。‌‌中身の各バリアントはメモリ領域を共有しています。enum全体のサイズは、タグ領域+バリアントの最大サイズで決まります。

[           e n u m           ]
-------------------------------
[tag] | [        union        ]
-------------------------------
[001] | [--- variant A ---]
[002] | [ variant B ]
[003] | [----- variant C -----]

// ※タグ領域が実際に配置される場所はコンパイラのメモリ配置最適化による
// ※タグやバリアントのメモリアライメントによってはパディング(何にも使わない領域)が入ることもある
// ※標準ライブラリ内の型はコンパイラ特権でさらに変な最適化が入ることも

enumの各バリアントは共通のメモリ領域を共有しているため、未定義動作を避けるには先頭にタグ領域を用意し、今どのバリアントのデータが有効なのか確定させて処理を分岐させる必要があります。

タグ領域はバリアントの数によってu8だったり、実用上滅多にあり得ませんがu16, u32, u64などと増減します。これはコンパイル時確定です。

しかし、どのバリアントが有効か外から判別できたり、アルゴリズム上確定/保証できる場合はタグをわざわざ付けなくても上手くプログラムできるのではないでしょうか？

こういう最適化したいときが、unionの出番です。

union/タグ無し共用体 (1)

union型はenumからタグ領域を飛ばしたものです。

[        union        ]
-----------------------
[--- variant A ---]
[ variant B ]
[----- variant C -----]

各フィールドが同じメモリ領域を共有していることに注意すれば、構造体宣言はstructに似た感じです

union MyUnion {
    a: u8,
    b: u16,
    c: u32,
}

ジェネリクスも使えます、が、

unionのフィールドに置ける型は以下の4種類に制限されています

見慣れない型が出てきました

std::mem::ManuallyDrop

union/タグ無し共用体の話の続きをする前に、 ManuallyDrop の話をせねばなりません。

問題はdrop処理です。

union型ではコンパイラが有効なバリアントを判断できないため、どのバリアントについてdrop処理を走らせれば良いのかわかりません。もちろん適当にdropしてしまえば未定義動作まっしぐらです。

そのため、すべてのバリアントは ManuallyDrop に包み、プログラマが責任を持って適切なバリアントについてdrop処理を行うと宣言する必要があります。

ManuallyDrop に包まれた値はRAII/スコープによる自動dropが走りません。プログラマの責任でdropを書かない限りはdrop処理は走りません。

ManuallyDrop<Box<T>> などに Box や Vec 、 String など入れた際には手動dropを行わないと即メモリリークです。

また、ManuallyDropのdropメソッドは所有権を消費しません。なので2重dropの危険があり、unsafeです。

この構造体はコンパイラへの一種のメッセージです。実行時コストはゼロです。

union/タグ無し共用体 (2)

unionに戻ります。

unionのdrop時にはコンパイラが自動でdropできないため、プログラマが明示的にdropを書かないといけません。なので、コンパイラがdrop処理不要だと知っている型以外は、ManuallyDropに包まないとunionのフィールドに持てません。‌‌逆にManuallyDropに包んでしまえばジェネリクスも使えます。

// これはコンパイルできない
union MyUnion {
    s: String
}

//　ManuallyDropに包む
union MyUnion {
    s: ManuallyDrop<String>
}

// ジェネリクスも使える
union MyGenericUnion<T> {
    data: ManuallyDrop<T>
}

unionのバリアントへは、作成、代入、書き込み時はunsafe不要、読み出しはunsafeブロック内で行う必要があります。

書き込みの場合、safe部分の構造体、即ちもともと安全な保証がある(とされている)構造体や型のバイト列をそのままコピーする操作にあたるため、未定義動作の余地がないからです。

しかし読み出しはunsafeです。バリアントを間違えると未定義動作だからです。

実際の操作は、 . (ドット)を使って構造体のフィールドのようにアクセスします。

union MyU {
    a: u8,
    b: u64,
    c: ManuallyDrop<Box<u8>>,
}

// 作成はsafe
let mut u = MyU { a: 7u8, };

// 書き込みもsafe
u.b = 23094857u64;

// 代入もsafe
let u2 = u;

// しかし読み出しはunsafe
let b_val = unsafe { u2.b };

// ！！未定義動作！！
// 今のu2には64bit整数のビットパターンが入っているため、Boxとして読み出した時点で未定義動作
// 高確率でDerefした瞬間メモリアクセス違反でOSにKILLされる
let maybe_c_val = unsafe { u2.c };

unionの用途

union型は、共用部分のバリアントが外から判別できるときに使います。あるいは本当に repr(C) した上でメモリをこねくり回して変なことがしたい時でしょうか。

僕の使い方は、SmallBox構造を一部持つ構造体の実装でした。

RustのBoxはヒープ領域にデータを置きます。しかし、u8のような小さな値をわざわざ64ビットのポインタを介して読むのは無駄の極みです。自分で管理している値ならまだしも、関数や構造体などでジェネリクスを使わないといけないときはBoxだけでは最適化できません。

小さな型も大きな型も問答無用でヒープに置くからです。

値をヒープに置くということは単にポインタアクセスが挟まるだけではありません。ヒープアロケーション、ヒープアクセス、デアロケーション全てがコストのかかる操作です。特にアクセスしようとしている領域がキャッシュに載っていなかったとき、運が悪ければ構造体インスタンス1個あたり数百〜千数百クロックの待ち時間が発生します。

ここで、Boxのポインタを表していた64ビット分の領域を使ってデータを直接持てばよいのではないでしょうか？
幸い、Rustには std::mem::size_of::<T>() という便利な関数があります。しかもこれはconst関数なので評価値はコンパイル時に埋め込まれ、この値を使った比較とif分岐は静的評価&分岐削除されるので実行時コストゼロです。

つまり、

if size_of::<T>() <= size_of::<usize>() {
    // ここにポインタを直接データとして読み出す処理
} else {
    // ここにポインタを辿る処理
}

という実装が実行時コストゼロで可能です。

さて、構造体に二種類の解釈を与えるためenumを使いたくなってきました。しかし、今回の場合はenumのうち、使うバリアントが静的に決まっています。unionにしてタグ領域の1バイト分削りましょう。

こう定義します。

union SmallBox<T> {
    inline: usize,
    heap: ManuallyDrop<Box<T>>
}

inline側を直接Tで置かなかったのは、unionはバリアントの最大サイズを確保するため、Tが大きいとたとえinlineを使わなくてもTのサイズ分unionが占有してしまうためです。usizeにして明示的にポインタサイズに抑えます。

これで、enumのタグを消し去ることが出来ました！
これでタグ部分の1バイトが節約できると思いきや、8バイト分の節約になります。
これは、SmallBoxのポインタやusizeが8バイトアラインメントを要求するためです。実は、元のenumのままでは、 タグ+データ で 1+8=9byte と思いきやアラインメントの関係でパディングが入り、全体で タグ+パディング+データ で 1+7+8=16byte になっていました。
そのため、タグを消し去っただけでサイズ半減です。

usize領域の未初期化部分の扱い、データのusize領域への書き込み読み出しと、unionの読み出しがunsafeなことに気をつけて実装すると、一旦はこうなります。

use std::{
    fmt::{Debug, Display},
    mem::{ManuallyDrop, MaybeUninit},
    ops::{Deref, DerefMut},
};

pub struct SmallBox<T> {
    data: SmallBoxData<T>,
}

union SmallBoxData<T> {
    // Tのサイズによってはusize内に未初期化部分が発生するのでMaybeUninitでマークしないと未定義動作
    inline: MaybeUninit<usize>,
    // ManuallyDropでマーク
    heap: ManuallyDrop<Box<T>>,
}

impl<T> SmallBox<T> {
    // 使い回すので関数に
    pub const fn is_inline() -> bool {
        size_of::<T>() <= size_of::<usize>()
    }

    pub fn new(value: T) -> Self {
        if Self::is_inline() {
            // 実体を作る
            let mut data = SmallBoxData {
                inline: MaybeUninit::<usize>::uninit(),
            };
            // unsafeでインライン領域へのポインタを取得&書き込み
            unsafe {
                let dst = data.inline.as_mut_ptr() as *mut T;
                std::ptr::write(dst, value);
            }

            Self { data }
        } else {
            Self {
                data: SmallBoxData {
                    heap: ManuallyDrop::new(Box::new(value)),
                },
            }
        }
    }

    pub fn addr(&self) -> *const T {
        if Self::is_inline() {
            (unsafe { self.data.inline.as_ptr() }) as *const T
        } else {
            unsafe { self.data.heap.as_ref() as *const T }
        }
    }

    pub fn into_inner(self) -> T {
        // このスコープでselfをforget、もしくはManuallyDropに包まないと、
        // 関数を抜けるときに、この後実装する SmallBox::drop() が自動的に呼ばれてしまう
        // データの実体はポインタを経由して既に吸い出されているので二重dropになってしまい未定義動作
        let this = ManuallyDrop::new(self);
        if Self::is_inline() {
            // usizeポインタをTポインタにキャストして読み出し
            unsafe { std::ptr::read(this.data.inline.as_ptr() as *const T) }
        } else {
            // ポインタ操作やキャストはないが、unionの読み出しなのでunsafeは必要
            unsafe { *ManuallyDrop::<Box<T>>::into_inner(std::ptr::read(&this.data.heap)) }
        }
    }
}

impl<T> Deref for SmallBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        if Self::is_inline() {
            // unionの読み出し&ポインタキャスト
            unsafe { &*(self.data.inline.as_ptr() as *const T) }
        } else {
            // unionの読み出し
            unsafe { &*self.data.heap }
        }
    }
}

// Derefとほぼ同じ
impl<T> DerefMut for SmallBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        if Self::is_inline() {
            unsafe { &mut *(self.data.inline.as_ptr() as *mut T) }
        } else {
            unsafe { &mut *self.data.heap }
        }
    }
}

// 明示的にDropを実装しなければならない
impl<T> Drop for SmallBox<T> {
    fn drop(&mut self) {
        if Self::is_inline() {
            // 指定したポインタに対してdrop処理
            unsafe { std::ptr::drop_in_place(self.data.inline.as_ptr() as *mut T) }
        } else {
            // Boxなのでここを忘れるとメモリリーク
            unsafe { ManuallyDrop::drop(&mut self.data.heap) }
        }
    }
}

impl<T: Display> Display for SmallBox<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", **self)
    }
}

impl<T: Debug> Debug for SmallBox<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", **self)
    }
}

ちなみに、これでもClaudeに怒られます。

メモリアライメントについて考慮できていないためです。

普通の変数ならusize以下のサイズの型のメモリアライメントがusizeを超えることはありません。usize以上のサイズの型であれば Box と std::alloc::alloc がいい感じにやってくれます。

しかしRustには手動でメモリアライメントを指定する方法があり、小さい型に途方もない過剰なメモリアライメントを指定することができます。
(これはfalse sharing対策など、ごく稀に実用されることがあります)

実体サイズが1以上の型でメモリアライメントを大きく指定するとき、基本的に型サイズはメモリアライメントに合わせられます。天井関数的な感じです。過剰アライメントされていても実体のある型をSmallBoxに入れれば勝手にヒープに置かれ、あとは Box がよしなにやってくれます。

しかし、ゼロサイズ型では勝手が変わってしまいます。サイズがゼロなのでアライメント整合された型サイズも0です。そのため、SmallBoxではインラインに置きたくなります。しかし、さっき書いたSmallBoxのインラインはusizeとして確保されているため、Tがusize以上のアライメントを持っていた場合にZSTの参照としてのポインタがアライメント整合している保証はありません。

ZSTが実体を持たないことを利用します。実体を持たず、メモリを触らないためどんなポインタであってもアライメントさえ合っていれば問題ないはずです。

ZSTだけ NonNull::<T>::dangling().as_ptr() を使うように分岐します。
NonNull::<T>::dangling() は自動でアライメント整合されたアドレスを与えてくれます。

// 前略

impl<T> SmallBox<T> {
    // 過剰アライメントなZSTのための分岐
    fn inline_ptr(&self) -> *const T {
        if size_of::<T>() == 0 {
            NonNull::<T>::dangling().as_ptr()
        } else {
            unsafe { self.data.inline.as_ptr() as *const T }
        }
    }

    // inline_ptrのmutバージョン
    // コンパイラの借用チェックが機能するように `&mut self` を受ける
    fn inline_ptr_mut(&mut self) -> *mut T {
        if size_of::<T>() == 0 {
            NonNull::<T>::dangling().as_ptr()
        } else {
            unsafe { self.data.inline.as_mut_ptr() as *mut T }
        }
    }

    // ---------------------------------------------------------------------
    // このあとのポインタアクセスはすべて inline_ptr/inline_ptr_mut を経由させる
    // ---------------------------------------------------------------------

    pub fn new(value: T) -> Self {
        if Self::is_inline() {
            let mut data = SmallBoxData {
                inline: MaybeUninit::<usize>::uninit(),
            };
            // ZSTアライメント整合のための分岐
            let dst = if size_of::<T>() == 0 {
                NonNull::<T>::dangling().as_ptr()
            } else {
                unsafe { data.inline.as_mut_ptr() as *mut T }
            };
            unsafe { std::ptr::write(dst, value) };

            Self { data }
        } else {
            Self {
                data: SmallBoxData {
                    heap: ManuallyDrop::new(Box::new(value)),
                },
            }
        }
    }

    pub fn addr(&self) -> *const T {
        if Self::is_inline() {
            self.inline_ptr() as *const T
        } else {
            unsafe { self.data.heap.as_ref() as *const T }
        }
    }

    pub fn into_inner(self) -> T {
        let this = ManuallyDrop::new(self);
        if Self::is_inline() {
            unsafe { std::ptr::read(this.inline_ptr() as *const T) }
        } else {
            unsafe { *ManuallyDrop::<Box<T>>::into_inner(std::ptr::read(&this.data.heap)) }
        }
    }
}

impl<T> Deref for SmallBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        if Self::is_inline() {
            unsafe { &*(self.inline_ptr() as *const T) }
        } else {
            unsafe { &*self.data.heap }
        }
    }
}

impl<T> DerefMut for SmallBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        if Self::is_inline() {
            unsafe { &mut *self.inline_ptr_mut() }
        } else {
            unsafe { &mut *self.data.heap }
        }
    }
}

impl<T> Drop for SmallBox<T> {
    fn drop(&mut self) {
        if Self::is_inline() {
            unsafe { std::ptr::drop_in_place(self.inline_ptr_mut()) }
        } else {
            unsafe { ManuallyDrop::drop(&mut self.data.heap) }
        }
    }
}

これで安全に、小サイズ最適化(Small Object Optimization)を行うスマートポインタの実装ができました。

現在の安定版Rustではマクロを使ったコード生成で型サイズによって実装を変えるなど以外で、これ以上の軽量化はできません。

なので、ライブラリ内のジェネリクス付き構造体や関数などでマクロを使いにくい時は以上の形が事実上の最適解です。

(なお、トレイトオブジェクトをSOOしたい時は、この実装ではトレイトオブジェクトをヒープに置いてしまいます。fat-pointerが16バイト占めるためです。fat-pointerを含めた最適化をやりたいときはまた別の解が必要です。)

ちなみに、ここでは解説のためSmallBoxの最低限の実装を書きました。

Rustではsmallboxクレートがあり、 !Sized な型の小サイズ最適化の実装を提供しています。こちらはfat-pointerなどの対応もしていて、もし使いどころがあればSmallBoxの再発明をする前にAPIと内部実装を覗いてみるといいかもしれません。ちなみにsmallboxクレートはunionを使わずにSOOを実装しているみたいです。

また、大多数の用途ではsmallvecやcompact_strなどのクレートが有用です。
再発明に意気込む前にまずそちらに目を通しておくことを強くお勧めします。

僕の用途では、動画編集ソフトのFFIプラグイン規格と、ホスト側の実装、プラグイン側のSDK実装の際にFFI適正のある小サイズ最適化構造体が欲しくなったためにunsafe/unionに手を出しました。実際のコードはもうちょっとメタデータがあったり込み入った事情&操作が必要になる予定で、そのためにunionアプローチのSOOを自作しています。
(もしかしたらunionから離れるかもしれませんが、、、)

helgev icon
この記事を書いた人
helgev

へるげふです 原義クリエイターです　いろんなものを作っています

この記事をシェア

このエントリーをはてなブックマークに追加
共有

関連する記事

2022年9月26日
競プロしかシラン人間が web アプリ QK Judge を作った話
tqk icon tqk
2024年8月29日
クロスコンパイルRust
H1rono_K icon H1rono_K
2023年12月25日
オレオレRustプロジェクトテンプレート
H1rono_K icon H1rono_K
2024年9月10日
Tauriでtodolist作ってみた
mehm8128 icon mehm8128
2023年12月22日
rust-toolchain.tomlを書こう
H1rono_K icon H1rono_K
2017年12月26日
RustでMCMC(Metropolis-Hasting)
David icon David
記事一覧 タグ一覧 Google アナリティクスについて 特定商取引法に基づく表記