Demo of ArrayBuilder, Array and Datatype in Rust

gat.rs

#![feature(generic_associated_types)]

trait ArrayBuilder {
    type ArrayType: Array<Builder = Self>;

    fn new(capacity: usize) -> Self;
    fn append<'a>(
        &'a mut self,
        value: Option<<<Self as ArrayBuilder>::ArrayType as Array>::RefItem<'a>>,
    );
    fn finish(self) -> Self::ArrayType;
}

trait Array {
    type RefItem<'a>
    where
        Self: 'a;
    type Builder: ArrayBuilder<ArrayType = Self>;

    fn value_at<'a>(&'a self, idx: usize) -> Option<Self::RefItem<'a>>;
    fn len(&self) -> usize;
}

struct IntArrayBuilder {
    bitmap: Vec<bool>,
    data: Vec<u64>,
}

impl ArrayBuilder for IntArrayBuilder {
    type ArrayType = IntArray;

    fn new(_capacity: usize) -> Self {
        Self {
            bitmap: vec![],
            data: vec![],
        }
    }

    fn append<'a>(&'a mut self, value: Option<u64>) {
        match value {
            Some(x) => {
                self.bitmap.push(true);
                self.data.push(x);
            }
            None => {
                self.bitmap.push(false);
                self.data.push(0);
            }
        }
    }

    fn finish(self) -> IntArray {
        IntArray {
            bitmap: self.bitmap,
            data: self.data,
        }
    }
}

struct IntArray {
    bitmap: Vec<bool>,
    data: Vec<u64>,
}

impl Array for IntArray {
    type Builder = IntArrayBuilder;
    type RefItem<'a> = u64;

    fn value_at(&self, idx: usize) -> Option<u64> {
        if self.bitmap[idx] {
            Some(self.data[idx])
        } else {
            None
        }
    }

    fn len(&self) -> usize {
        self.bitmap.len()
    }
}

struct BytesArray {
    offset: Vec<usize>,
    bitmap: Vec<bool>,
    data: Vec<u8>,
}

impl Array for BytesArray {
    type Builder = BytesArrayBuilder;
    type RefItem<'a> = &'a [u8];

    fn value_at(&self, idx: usize) -> Option<&[u8]> {
        if self.bitmap[idx] {
            Some(&self.data[self.offset[idx]..self.offset[idx + 1]])
        } else {
            None
        }
    }

    fn len(&self) -> usize {
        self.bitmap.len()
    }
}

struct BytesArrayBuilder {
    offset: Vec<usize>,
    bitmap: Vec<bool>,
    data: Vec<u8>,
}

impl ArrayBuilder for BytesArrayBuilder {
    type ArrayType = BytesArray;

    fn new(_capacity: usize) -> Self {
        Self {
            offset: vec![0],
            data: vec![],
            bitmap: vec![],
        }
    }

    fn append<'a>(&'a mut self, value: Option<&'a [u8]>) {
        match value {
            Some(x) => {
                self.bitmap.push(true);
                self.data.extend_from_slice(x);
                self.offset.push(self.data.len())
            }
            None => {
                self.bitmap.push(false);
                self.offset.push(self.data.len())
            }
        }
    }

    fn finish(self) -> BytesArray {
        BytesArray {
            bitmap: self.bitmap,
            data: self.data,
            offset: self.offset,
        }
    }
}

fn filter<'a, A, F>(data: &'a A, pred: F) -> A
where
    A: Array + 'a,
    F: Fn(Option<A::RefItem<'a>>) -> bool,
{
    let mut builder = A::Builder::new(data.len());
    for i in 0..data.len() {
        if pred(data.value_at(i)) {
            builder.append(data.value_at(i));
        }
    }
    builder.finish()
}

fn main() {
    let mut builder = IntArrayBuilder::new(0);
    for i in 0..60 {
        builder.append(Some(i));
    }
    let array = builder.finish();
    let filtered_array = filter(&array, |x| if let Some(x) = x { x > 50 } else { false });
    for i in 0..filtered_array.len() {
        println!("{:?}", filtered_array.value_at(i));
    }

    let mut builder = BytesArrayBuilder::new(0);
    for i in 0..110 {
        builder.append(Some(format!("{}", i).as_bytes()));
    }
    let array = builder.finish();
    let filtered_array = filter(
        &array,
        |x| if let Some(x) = x { x.len() > 2 } else { false },
    );
    for i in 0..filtered_array.len() {
        println!(
            "{:?}",
            std::str::from_utf8(filtered_array.value_at(i).unwrap()).unwrap()
        );
    }
}