Last active
October 30, 2022 00:33
-
-
Save lovasoa/fdf6a558be15bce33703a3b67d9d99ce to your computer and use it in GitHub Desktop.
autovec : alternative to rust's standard vector type which uses column-oriented storage for better code auto-vectorization. In short: makes rust code working with vectors of structs magically faster.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#[derive(Clone, Debug, PartialEq)] | |
struct B { | |
x: bool, | |
y: f64, | |
} | |
pub trait AutoVec { | |
type Vec; | |
} | |
pub trait LikeB { | |
fn x(&self) -> bool; | |
fn y(&self) -> f64; | |
} | |
#[derive(Default)] | |
pub struct SizedAutoVecB { | |
x: [bool; SIZE], | |
y: [f64; SIZE], | |
} | |
#[derive(Clone, Copy)] | |
pub struct SizedAutoVecBIter<'a> { | |
item: &'a SizedAutoVecB, | |
i: usize, | |
} | |
impl<'a> LikeB for SizedAutoVecBIter<'a> { | |
fn x(&self) -> bool { | |
self.item.x[self.i] | |
} | |
fn y(&self) -> f64 { | |
self.item.y[self.i] | |
} | |
} | |
impl<T: LikeB> From<T> for B { | |
fn from(a: T) -> Self { | |
Self { | |
x: a.x(), | |
y: a.y(), | |
} | |
} | |
} | |
impl AutoVecB { | |
#[inline(always)] | |
fn len(&self) -> usize { self.size } | |
#[inline(always)] | |
fn get(&self, index: usize) -> Option<SizedAutoVecBIter<'_>> { | |
self.bins.get(index / SIZE) | |
.map(|bin| SizedAutoVecBIter { i: index % SIZE, item: bin }) | |
.filter(|_| index < self.len()) | |
} | |
} | |
const SIZE: usize = 32; | |
#[derive(Default)] | |
pub struct AutoVecB { | |
bins: Vec<SizedAutoVecB>, | |
size: usize, | |
} | |
impl AutoVecB { | |
pub fn iter(&self) -> AutoVecBIter<'_> { | |
self.into_iter() | |
} | |
} | |
impl From<Vec<B>> for AutoVecB { | |
fn from(v: Vec<B>) -> Self { | |
let bins = v.chunks(SIZE) | |
.map(|fixed| { | |
let mut d = SizedAutoVecB::default(); | |
for (i, b) in fixed.into_iter().enumerate() { | |
d.x[i] = b.x; | |
d.y[i] = b.y; | |
} | |
d | |
}) | |
.collect(); | |
Self { bins, size: v.len() } | |
} | |
} | |
pub struct AutoVecBIter<'a> { | |
base: &'a AutoVecB, | |
j: usize, | |
} | |
impl<'a> Iterator for AutoVecBIter<'a> { | |
type Item = SizedAutoVecBIter<'a>; | |
#[inline(always)] | |
fn next(&mut self) -> Option<Self::Item> { | |
let res = self.base.get(self.j); | |
self.j += 1; | |
res | |
} | |
} | |
impl<'a> IntoIterator for &'a AutoVecB { | |
type Item = SizedAutoVecBIter<'a>; | |
type IntoIter = AutoVecBIter<'a>; | |
#[inline(always)] | |
fn into_iter(self) -> Self::IntoIter { | |
AutoVecBIter { base: self, j: 0 } | |
} | |
} | |
impl AutoVec for B { | |
type Vec = AutoVecB; | |
} | |
#[test] | |
fn test_iter() { | |
let original_vec: Vec<B> = (0..10000).map(|i| B { x: i % 2 == 0, y: f64::from(i) }).collect(); | |
let a = AutoVecB::from(original_vec.clone()); | |
let back_to_vec: Vec<B> = a.into_iter().map(B::from).collect(); | |
assert_eq!(original_vec, back_to_vec) | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment