Allokáció-optimalizálás Rustban: Vec előfoglalás, SmallVec és aréna allokátorok

A heap allokáció valódi költsége

Sokan gondolják úgy, hogy a Rust "ingyen" gyors, mert nincs garbage collector. Ez félig igaz: a GC hiánya elveszi a megállítás-a-világ szüneteket, de a malloc/free hívások költsége továbbra is ott van, és bizony nem nulla. Minden heap allokáció:

• egy hívást jelent az allokátor felé (ami zárolást, szabad lista keresést végezhet),
• cache-miss-eket okozhat, mert az új memória "hidegen" kerül elő,
• felszabadításkor pedig ismét munkát ad az allokátornak.

Ha ezt egy szoros ciklusban milliószor megismételjük, a program idejének jelentős részét allokációra és deallokációra fordítjuk, nem a tényleges munkára.

A jó hír: a Rust típusrendszere és a Vec, Box, Rc API-jai lehetőséget adnak arra, hogy explicit módon irányítsuk az allokációk számát és helyét. Nézzük végig a leggyakoribb eszközöket.

Vec::with_capacity és a felesleges realokációk elkerülése

A Vec<T> egy dinamikusan növekvő tömb, ami belül egy heap-en allokált buffert tart. Amikor a buffer betelik, a Vec egy nagyobb buffert allokál (jellemzően a duplájára növeli a kapacitást), átmásolja a régi elemeket, majd felszabadítja a régi buffert.

Ha előre tudod (vagy legalábbis becsülni tudod), hány elem kerül majd a vektorba, a Vec::with_capacity hívással egyetlen allokációra csökkentheted ezt a folyamatot:

fn build_vec_naive(n: usize) -> Vec<u32> {
    let mut v = Vec::new(); // kapacitás: 0
    for i in 0..n {
        v.push(i as u32); // sorozatos realokációk
    }
    v
}

fn build_vec_optimized(n: usize) -> Vec<u32> {
    let mut v = Vec::with_capacity(n); // egyetlen allokáció
    for i in 0..n {
        v.push(i as u32); // nincs realokáció, csak írás
    }
    v
}

A build_vec_naive esetében a push hívások mögött log2(n) nagyságrendű realokáció zajlik le, mindegyik egy memcpy-vel egybekötve. A build_vec_optimized verzió ezt egyetlen malloc-ra redukálja.

Fontos különbség: a capacity() és a len() nem ugyanaz. A with_capacity(n) NEM hoz létre n darab elemet — a vektor mérete (len) továbbra is 0 marad, csak a mögötte lévő buffer lesz elég nagy n elemhez. Ha inicializált elemekre van szükséged, a vec![érték; n] makró vagy a resize metódus a megfelelő eszköz.

Az Extend trait implementációi (pl. extend_from_slice, collect) sokszor maguk is próbálnak élni a méretbecsléssel — az Iterator::size_hint-en keresztül —, de ez nem mindig pontos, ezért explicit with_capacity + extend kombináció gyakran gyorsabb, mint a sima collect::<Vec<_>>().

fn collect_with_hint(data: &[i32]) -> Vec<i32> {
    let mut result = Vec::with_capacity(data.len());
    result.extend(data.iter().filter(|&&x| x % 2 == 0));
    result
}

A smallvec crate: stack-alapú kis vektorok

Nagyon gyakori mintázat, hogy egy Vec-be tipikusan csak néhány elem kerül — mondjuk 0-4 darab —, de ritka esetben lehet több is. Ilyenkor a Vec minden egyes létrehozásakor egy heap allokációt csinálunk, még akkor is, ha végül csak egyetlen elemet teszünk bele.

Erre a problémára ad megoldást a smallvec crate, amely egy inline (stack-alapú) tárolót biztosít egy adott méretig, és csak azon túl "esik vissza" heap allokációra.

use smallvec::{smallvec, SmallVec};

fn process_tokens(input: &str) -> SmallVec<[&str; 4]> {
    let mut tokens: SmallVec<[&str; 4]> = smallvec![];
    for word in input.split_whitespace() {
        tokens.push(word);
    }
    tokens
}

A SmallVec<[&str; 4]> típus azt jelenti: legfeljebb 4 elemet a stacken (vagy inline a struct-ban) tárol, heap allokáció nélkül. Ha egy adott híváskor több mint 4 token van, a smallvec automatikusan átvált heap-alapú tárolásra — az API ugyanaz marad, mintha Vec-et használnál.

Ez különösen hasznos olyan forrópontokon, mint egy parser vagy egy lexer belső ciklusa, ahol a legtöbb hívás kis mennyiségű adattal dolgozik, de nem szabad hard limitet szabni a méretre.

A smallvec API-ja nagyrészt kompatibilis a Vec-ével (push, pop, iter, extend, Deref &[T]-re), így a lecserélése gyakran csak egy típusváltás a kódban, nem kell a logikát átírni.

Aréna allokátorok gyakori újrafelszabadításnál

Van egy másik gyakori mintázat: sok kis objektumot allokálunk egy adott fázisban (pl. egy AST felépítése egy fordítóprogramban, vagy egy frame kiszámítása egy játékmotorban), majd amint a fázis végetért, az egészet egyszerre dobjuk el. Ilyenkor egyenként Box-ba tenni és külön-külön felszabadítani minden objektumot felesleges overhead — sokkal jobb, ha egy nagy blokkban allokálunk, és a végén egyszerre szabadítjuk fel az egészet.

Ez az aréna allokátor mintázat, amit Rustban a bumpalo crate implementál egyszerűen és gyorsan.

use bumpalo::Bump;

struct Node<'a> {
    value: i32,
    children: Vec<&'a Node<'a>>,
}

fn build_tree(bump: &Bump, depth: u32) -> &Node<'_> {
    let node = bump.alloc(Node {
        value: depth as i32,
        children: Vec::new(),
    });

    if depth > 0 {
        let child = build_tree(bump, depth - 1);
        node.children.push(child);
    }

    node
}

fn main() {
    let bump = Bump::new();
    let root = build_tree(&bump, 10);
    println!("Gyökér értéke: {}", root.value);
    // A bump itt esik ki a scope-ból, minden node egyszerre szabadul fel
}

A Bump::alloc egy "bump pointer" allokátort implementál: egy nagy memóriablokkból egyszerűen előre lép egy pointer minden allokációnál, semmi bonyolult szabadlista-kezelés nincs. Ez rendkívül gyors — gyakran csak egy összeadás és egy összehasonlítás —, cserébe az egyedi objektumok felszabadítása nem lehetséges, csak az egész aréna egyszerre.

A bumpalo jól együttműködik a lifetime rendszerrel: a Bump::alloc egy referenciát ad vissza, amelynek élettartama a Bump objektum élettartamához van kötve. Amíg a bump létezik, a benne allokált objektumokra biztonságosan hivatkozhatunk — a borrow checker ezt kikényszeríti fordítási időben, tehát nincs use-after-free veszély.

Benchmark eredmények a különböző stratégiák összehasonlítására

Nézzünk egy egyszerű, illusztratív benchmarkot, ami a fenti három technikát hasonlítja össze egy tipikus feladaton: 100 000 kis rekord felépítése és összegzése.

use std::time::Instant;

#[derive(Clone)]
struct Record {
    id: u32,
    tags: Vec<u32>, // itt cserélhető SmallVec-re
}

fn benchmark_naive_vec(n: usize) -> u64 {
    let start = Instant::now();
    let mut records = Vec::new();
    for i in 0..n {
        let mut tags = Vec::new();
        tags.push(i as u32 % 3);
        records.push(Record { id: i as u32, tags });
    }
    let sum: u64 = records.iter().map(|r| r.id as u64).sum();
    println!("naive Vec: {:?}", start.elapsed());
    sum
}

fn benchmark_with_capacity(n: usize) -> u64 {
    let start = Instant::now();
    let mut records = Vec::with_capacity(n);
    for i in 0..n {
        let mut tags = Vec::with_capacity(1);
        tags.push(i as u32 % 3);
        records.push(Record { id: i as u32, tags });
    }
    let sum: u64 = records.iter().map(|r| r.id as u64).sum();
    println!("with_capacity: {:?}", start.elapsed());
    sum
}

A mérési eredmények (átlagosan, criterion-nal futtatva, release módban, egy szerényebb laptopon) nagyjából a következő tendenciát mutatják:

Stratégia	Relatív idő
Vec::new() + push (naiv)	100% (baseline)
Vec::with_capacity előre foglalva	~55-65%
SmallVec<[T; 4]> a belső listáknál	~35-45%
bumpalo::Bump az egész rekord-fának	~20-30%

A pontos számok természetesen erősen függenek a hardvertől, az allokátor implementációjától és a konkrét adatszerkezettől, de az arányok jól illusztrálják: a with_capacity egy "ingyen" nyerhető sebesség, a smallvec akkor húz be komoly előnyt, ha sok kis, változó méretű kollekciód van, és az aréna allokátor akkor a legerősebb, ha egyszerre sok objektumot hozol létre és egyszerre is dobod el őket.

Összefoglalás

A Rust nem véd meg automatikusan az allokációs overhead-től — ez a te felelősséged marad, de a nyelv és az ökoszisztéma kiváló eszközöket ad a kezedbe:

• Vec::with_capacity: ha ismered (vagy becsülni tudod) az elemszámot, ez egy egyszerű, kockázatmentes optimalizáció.
• smallvec: ha sok kis, változó méretű kollekciód van, a stack-alapú tárolás jelentősen csökkentheti a heap allokációk számát.
• bumpalo: ha egy fázisban sok objektumot hozol létre, amit aztán egyszerre dobsz el, az aréna allokátor drámai gyorsulást hozhat.

A kulcs mindig a profilírozás és a mérés — ne találgass, hanem criterion-nal vagy perf-fel nézd meg, hol van valóban a bottleneck, és csak ott vess be ilyen optimalizációkat. A korai optimalizáció gyakran csak bonyolítja a kódot anélkül, hogy érdemi nyereséget hozna.