Iterator của Rust thật sự không tốn gì

Đến từ Swift, nơi map và filter cấp phát các mảng trung gian trừ khi bạn nhớ viết .lazy, tôi đã không tin lời tuyên bố của Rust rằng một chuỗi iterator biên dịch ra đúng đoạn code như một vòng lặp viết tay.

Nên tôi đi kiểm tra.

Thí nghiệm

pub fn sum_of_even_squares(values: &[i32]) -> i32 {
    values.iter()
        .filter(|&&x| x % 2 == 0)
        .map(|&x| x * x)
        .sum()
}

pub fn sum_of_even_squares_loop(values: &[i32]) -> i32 {
    let mut total = 0;
    for &x in values {
        if x % 2 == 0 {
            total += x * x;
        }
    }
    total
}

Biên dịch với --release rồi soi bằng cargo asm, hai hàm này sinh ra cùng một mã máy — cả hai đều được tự động vector hóa thành các lệnh SIMD xử lý vài số nguyên mỗi chu kỳ.

Không phải “hiệu năng tương đương”. Mà là cùng những lệnh y hệt.

Vì sao nó chạy được

Ba thứ kết hợp lại, và không thứ nào là phép màu.

Iterator là struct, không phải đối tượng. values.iter().filter(…).map(…) dựng ra một giá trị thuộc kiểu Map<Filter<Iter<i32>, closure>, closure>. Không cấp phát, không đóng hộp, và không điều phối động — các kiểu closure được nướng thẳng vào cái kiểu.

Mọi thứ đều được nội tuyến. Hàm next() của mỗi adapter là một hàm nhỏ được đánh dấu để nội tuyến. Sau khi nội tuyến, cả chuỗi lời gọi next() sập lại thành một thân vòng lặp duy nhất.

LLVM tối ưu phần còn lại. Một khi nó đã là một vòng lặp duy nhất trên một lát cắt với một nhánh rẽ và một phép nhân, trình tối ưu thông thường lo được — bao gồm cả việc vector hóa.

Điểm cấu trúc then chốt là việc duyệt mang tính kéo. Không gì xảy ra cho tới khi sum() xin một giá trị, và mỗi yêu cầu chạy xuống dọc chuỗi rồi quay lại. Không có tập hợp trung gian nào ở bất kỳ điểm nào, vì không adapter nào giữ nhiều hơn một phần tử.

Chỗ nó thôi miễn phí

Ba trường hợp, tất cả đều nhìn thấy được ngay trong code.

collect() cấp phát. Hiển nhiên — bạn đã xin một tập hợp mà. Thứ ít hiển nhiên hơn là nối collect() vào giữa một pipeline sẽ vô hiệu hóa cả cơ chế:

// hai lần cấp phát, hai lượt duyệt
let result: Vec<_> = values.iter().map(|x| x * 2).collect();
let result: Vec<_> = result.iter().filter(|x| **x > 10).collect();

// không cấp phát cho tới cuối
let result: Vec<_> = values.iter().map(|x| x * 2).filter(|x| *x > 10).collect();

Box<dyn Iterator> phá vỡ việc nội tuyến. Xóa kiểu một iterator nghĩa là trình biên dịch không còn nhìn xuyên qua next() được nữa, nên chẳng có gì sập lại và mỗi phần tử tốn một lời gọi ảo.

Vài adapter buộc phải đệm. sorted không phải một adapter iterator là có lý do — việc sắp xếp cần mọi phần tử. rev() trên một iterator không hai đầu, peekable, và chunks đều giữ trạng thái, dù thường là một lượng hằng số.

collect() khôn hơn vẻ ngoài của nó

collect là generic trên kiểu đích, thứ cho phép nó làm những việc mà ở ngôn ngữ khác trông như các hàm riêng biệt:

let vec: Vec<i32> = iter.collect();
let set: HashSet<i32> = iter.collect();
let map: HashMap<String, i32> = pairs.collect();
let string: String = chars.collect();

// cái đáng biết
let results: Result<Vec<i32>, ParseError> = strings.iter().map(|s| s.parse()).collect();

Cái cuối thật sự hữu ích và không hiển nhiên: gom một iterator của Result thành một Result<Vec<_>, _> sẽ dừng ngay ở lỗi đầu tiên và trả về lỗi đó. Điều tương tự áp dụng cho Option. Nó thay một vòng lặp có return sớm bằng một dòng.

Tôi đã đổi gì trong code của mình

Chủ yếu là tôi thôi lo lắng. Cụ thể:

Tôi thôi viết vòng lặp thủ công vì hiệu năng. Chuỗi iterator rõ ràng hơn và sinh ra cùng những lệnh máy.

Tôi thôi dùng collect() ở giữa chừng. Đó là cái giá thật duy nhất trong code của tôi, và có mấy chỗ như vậy — thường được viết ra vì một chuỗi dài quá nên tôi chẻ nó ra cho dễ đọc. Một phép gán let cho chính cái iterator làm được đúng việc đó mà không cấp phát.

Tôi dùng iter(), iter_mut() và into_iter() một cách có chủ đích. Mượn, mượn khả biến, tiêu thụ. Phần lớn những trận vật lộn ban đầu của tôi với borrow checker là do với nhầm một trong ba cái này.

Lời cảnh báo thành thật

“Chi phí bằng không” nghĩa là không tốn thêm chi phí lúc chạy so với bản viết tay tương đương. Nó không miễn phí theo hai nghĩa khác.

Thời gian biên dịch. Mỗi adapter là một kiểu mới và monomorphisation sinh code cho từng cái. Các chuỗi iterator dài thì biên dịch chậm hơn vòng lặp một cách đo được.

Khả năng gỡ lỗi. Bước qua một chuỗi mười adapter đã bị nội tuyến trong một bản release thì khó chịu, và một cú panic bên trong một closure cho bạn một backtrace đầy nội tạng của iterator.

Không cái nào trong hai thứ đó đủ để khiến tôi quay về viết vòng lặp. Nhưng “chi phí bằng không” là một lời tuyên bố về mã được sinh ra, không phải về toàn bộ trải nghiệm, và tôi thà nói vậy còn hơn nhắc lại cái khẩu hiệu.