<commit_before>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};
use std::mem::replace;

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub grouped_fields: Vec<usize>,
    pub sorted_fields: Vec<(usize, Direction)>,
    pub chunked: bool,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

impl Source {
    fn prepare(&self, mut rows: Vec<Vec<Value>>) -> Vec<Vec<Value>> {
        // TODO compensate for inputs if primitive
        rows.sort_by(|a, b| {
            for &ix in self.grouped_fields.iter() {
                match a[ix].cmp(&b[ix]) {
                    Ordering::Greater => return Ordering::Greater,
                    Ordering::Less => return Ordering::Less,
                    Ordering::Equal => ()
                }
            }
            for &(ix, direction) in self.sorted_fields.iter() {
                match (a[ix].cmp(&b[ix]), direction) {
                    (Ordering::Greater, Direction::Ascending) => return Ordering::Greater,
                    (Ordering::Greater, Direction::Descending) => return Ordering::Less,
                    (Ordering::Less, Direction::Ascending) => return Ordering::Less,
                    (Ordering::Less, Direction::Descending) => return Ordering::Greater,
                    (Ordering::Equal, _) => ()
                }
            }
            return Ordering::Equal;
        });
        let mut groups = Vec::new();
        let row_len = rows.iter().next().map_or(0, |row| row.len());
        let mut group = vec![Value::Null; row_len];
        let mut ix = 1;
        for mut row in rows.into_iter() {
            row.push(Value::Float(ix as f64));
            let is_same_group = self.grouped_fields.iter().all(|&ix| row[ix] == group[ix]);
            match (self.chunked, is_same_group) {
                (true, true) => {
                    for &(ix, _) in self.sorted_fields.iter() {
                        group[ix].as_column_mut().push(row[ix].clone())
                    }
                    ix += 1;
                }
                (true, false) => {
                    groups.push(group);
                    group = row;
                    for &(ix, _) in self.sorted_fields.iter() {
                        let value = replace(&mut group[ix], Value::Null);
                        group[ix] = Value::Column(vec![value]);
                    }
                    ix = 1;
                }
                (false, true) => {
                    groups.push(group);
                    group = row;
                    ix += 1;
                }
                (false, false) => {
                    groups.push(group);
                    group = row;
                    ix = 1;
                }
            }
        }
        groups.push(group);
        groups.remove(0); // remove the dummy group we started with
        groups
    }
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

// TODO this algorithm is incredibly naive and also clones excessively
fn join_step(join: &Join, ix: usize, inputs: &[Vec<Vec<Value>>], state: &mut Vec<Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.iter().map(|ix| state[*ix].clone()).collect());
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{..} => {
                for values in inputs[ix].iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = values[field_ix].clone();
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                // values returned from primitives don't include inputs, so we will have to offset accesses by input_len
                let input_len = input_bindings.len();
                for mut values in primitive.eval_from_join(&input_bindings[..], &state[..]).into_iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = replace(&mut values[field_ix - input_len], Value::Null);
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix - input_len]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, upstream: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), upstream.len());
                for select in union.selects.iter() {
                    for values in select.select(&upstream[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.clone();
                let inputs = join.sources.iter().map(|source|
                    match source.input {
                        Input::Primitive{..} => {
                            vec![]
                        }
                        Input::View{input_ix, ..} => {
                            source.prepare(upstream[input_ix].index.iter().map(|row| row.clone()).collect())
                        }
                    }).collect::<Vec<_>>();
                join_step(join, 0, &inputs[..], &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&upstream[..]);
                let mut inner = aggregate.inner.select(&upstream[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<commit_msg>fix the prepare function to correctly build groups across all the possible states chunked/unchunked, grouped/ungrouped, sorted/unsorted<commit_after>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};
use std::mem::replace;

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub grouped_fields: Vec<usize>,
    pub sorted_fields: Vec<(usize, Direction)>,
    pub chunked: bool,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

impl Source {
    fn prepare(&self, mut rows: Vec<Vec<Value>>) -> Vec<Vec<Value>> {
        // TODO compensate for inputs if primitive
        rows.sort_by(|a, b| {
            for &ix in self.grouped_fields.iter() {
                match a[ix].cmp(&b[ix]) {
                    Ordering::Greater => return Ordering::Greater,
                    Ordering::Less => return Ordering::Less,
                    Ordering::Equal => ()
                }
            }
            for &(ix, direction) in self.sorted_fields.iter() {
                match (a[ix].cmp(&b[ix]), direction) {
                    (Ordering::Greater, Direction::Ascending) => return Ordering::Greater,
                    (Ordering::Greater, Direction::Descending) => return Ordering::Less,
                    (Ordering::Less, Direction::Ascending) => return Ordering::Less,
                    (Ordering::Less, Direction::Descending) => return Ordering::Greater,
                    (Ordering::Equal, _) => ()
                }
            }
            return Ordering::Equal;
        });

        if self.chunked {
            // if the source is chunked, we need to turn all of the fields into columns for each grouping
            // we want the grouped fields to be scalars so that they can be joined against and every other
            // field should be a column of the values for that group
            let mut groups = Vec::new();
            let mut ordinal = 0;
            let grouped_fields = &self.grouped_fields;
            let row_len = rows.iter().next().map_or(0, |row| row.len());
            let ungrouped_fields = &(0..row_len).filter(|field_ix| !self.grouped_fields.contains(field_ix)).collect::<Vec<usize>>();
            let mut group: Vec<Value> = vec![Value::Column(vec![]); row_len]; // this creates a vector of the given length with all the values set to Value::Null
            for mut row in rows.into_iter() {
                // check if this is a new group
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == group[ix])
                };
                if is_new_group {
                    // add the current ordinal to the group
                    group.push(Value::Float(ordinal as f64));
                    // push the old group
                    groups.push(group);
                    // create a new group with every field being a column
                    ordinal += 1;
                    group = vec![Value::Column(vec![]); row_len];
                    // any grouped field should have a scalar value for its index so that it can be joined on
                    // e.g. ["foo", [..], [..]] if you've grouped on index 0
                    for &grouped_field_ix in grouped_fields {
                        group[grouped_field_ix] = row[grouped_field_ix].clone();
                    }
                }
                // add the fields for this row
                for &field_ix in ungrouped_fields {
                    group[field_ix].as_column_mut().push(row[field_ix].clone());
                }
            }
            groups.push(group);
            // if there are grouped fields then the first group is the initial placeholder and
            // should be removed
            if self.grouped_fields.len() > 0 {
                groups.remove(0);
            }
            groups
        } else {
            // in the case where the source is not chunked all we need to worry about is setting the ordinal correctly
            // if there are no groups, then this is just normal rows and we number them monotonically
            // if there are groups, each new group resets the ordinal to 1
            let otherwise = vec![];
            let mut groups = Vec::new();
            let mut ordinal = 1;
            let mut prev_row = rows.iter().next().map_or(otherwise, |row| row.clone());
            for mut row in rows.into_iter() {
                let is_new_group = match self.grouped_fields.len() {
                    0 => false,
                    _ => !self.grouped_fields.iter().all(|&ix| row[ix] == prev_row[ix])
                };
                if is_new_group {
                    ordinal = 1;
                    prev_row = row.clone(); // @TODO: is there a way to do this without having to clone it?
                }
                row.push(Value::Float(ordinal as f64));
                groups.push(row);
                ordinal += 1;
            }
            groups
        }
    }
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

// TODO this algorithm is incredibly naive and also clones excessively
fn join_step(join: &Join, ix: usize, inputs: &[Vec<Vec<Value>>], state: &mut Vec<Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.iter().map(|ix| state[*ix].clone()).collect());
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{..} => {
                for values in inputs[ix].iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = values[field_ix].clone();
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                // values returned from primitives don't include inputs, so we will have to offset accesses by input_len
                let input_len = input_bindings.len();
                for mut values in primitive.eval_from_join(&input_bindings[..], &state[..]).into_iter() {
                    for &(field_ix, variable_ix) in source.output_bindings.iter() {
                        state[variable_ix] = replace(&mut values[field_ix - input_len], Value::Null);
                    }
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        state[variable_ix] == values[field_ix - input_len]
                    ) {
                        join_step(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, upstream: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), upstream.len());
                for select in union.selects.iter() {
                    for values in select.select(&upstream[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.clone();
                let inputs = join.sources.iter().map(|source|
                    match source.input {
                        Input::Primitive{..} => {
                            vec![]
                        }
                        Input::View{input_ix, ..} => {
                            source.prepare(upstream[input_ix].index.iter().map(|row| row.clone()).collect())
                        }
                    }).collect::<Vec<_>>();
                join_step(join, 0, &inputs[..], &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&upstream[..]);
                let mut inner = aggregate.inner.select(&upstream[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<|endoftext|>


<commit_before>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ConstraintOp {
    EQ,
    NEQ,
    LT,
    GT,
    LTE,
    GTE,
}

#[derive(Clone, Debug)]
pub struct Constraint {
    pub left: usize,
    pub op: ConstraintOp,
    pub right: usize,
}

impl Constraint {
    fn is_satisfied_by(&self, state: &[&Value]) -> bool {
        match self.op {
            ConstraintOp::EQ => state[self.left] == state[self.right],
            ConstraintOp::NEQ => state[self.left] != state[self.right],
            ConstraintOp::LT => state[self.left] < state[self.right],
            ConstraintOp::GT => state[self.left] > state[self.right],
            ConstraintOp::LTE => state[self.left] <= state[self.right],
            ConstraintOp::GTE => state[self.left] >= state[self.right],
        }
    }
}

#[derive(Clone, Debug)]
pub enum JoinSource {
    Relation{
        input: usize
    },
    Primitive{
        primitive: Primitive,
        arguments: Vec<usize>,
    },
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<JoinSource>,
    pub constraints: Vec<Vec<Constraint>>,
    pub select: ViewSelect,
    pub join2: Join2,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

#[derive(Clone, Debug)]
pub struct Join2 {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
    Join2(Join2),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

fn join_step<'a>(join: &'a Join, ix: usize, inputs: &[&'a Relation], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.select(&state[..]));
    } else {
        match join.sources[ix] {
            JoinSource::Relation{input} => {
                for values in inputs[input].index.iter() {
                    push_all(state, values);
                    if join.constraints[ix].iter().all(|constraint| constraint.is_satisfied_by(&state[..])) {
                        join_step(join, ix+1, inputs, state, index)
                    }
                    pop_all(state, values);
                }
            }
            JoinSource::Primitive{ref primitive, ref arguments, ..} => {
                for values in primitive.eval_from_join(&arguments[..], &state[..]).into_iter() {
                    // promise the borrow checker that we will pop values before we exit this scope
                    // TODO this is not panic-safe - we should use CowString in Value instead
                    let values = unsafe { ::std::mem::transmute::<&Vec<Value>, &'a Vec<Value>>(&values) };
                    push_all(state, values);
                    if join.constraints[ix].iter().all(|constraint| constraint.is_satisfied_by(&state[..])) {
                        join_step(join, ix+1, inputs, state, index)
                    }
                    pop_all(state, values);
                }
            }
        }
    }
}

fn join_step2<'a>(join: &'a Join2, ix: usize, inputs: &[&'a Relation], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.select(&state[..]));
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{input_ix} => {
                for values in inputs[input_ix].index.iter() {
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        *state[variable_ix] == values[field_ix]
                    ) {
                        for &(field_ix, variable_ix) in source.output_bindings.iter() {
                            state[variable_ix] = &values[field_ix];
                        }
                        join_step2(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                for values in primitive.eval_from_join2(&input_bindings[..], &state[..]).into_iter() {
                    // promise the borrow checker that we wont read these values after exiting this scope
                    // TODO this is not panic-safe - we should use CowString in Value instead
                    let values = unsafe { ::std::mem::transmute::<&Vec<Value>, &'a Vec<Value>>(&values) };
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        *state[variable_ix] == values[field_ix]
                    ) {
                        for &(field_ix, variable_ix) in source.output_bindings.iter() {
                            state[variable_ix] = &values[field_ix];
                        }
                        join_step2(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, inputs: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), inputs.len());
                for select in union.selects.iter() {
                    for values in select.select(&inputs[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.iter().collect();
                join_step(join, 0, inputs, &mut state, &mut output.index);
                let output2 = {
                    let mut output = Relation::new(
                        old_output.view.clone(),
                        old_output.fields.clone(),
                        old_output.names.clone()
                        );
                    let mut state = join.join2.constants.iter().collect();
                    join_step2(&join.join2, 0, inputs, &mut state, &mut output.index);
                    output
                };
                if output.index != output2.index {
                    println!("Failed for:");
                    println!("{:#?}", join);
                    println!("{:#?}", output);
                    println!("{:#?}", output2);
                }
                Some(output)
            }
            View::Join2(ref join) => {
                let mut state = join.constants.iter().collect();
                join_step2(join, 0, inputs, &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&inputs[..]);
                let mut inner = aggregate.inner.select(&inputs[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<commit_msg>Correctly offset output field indexes when running primitives<commit_after>use std::collections::BTreeSet;

use value::Value;
use relation::{Relation, IndexSelect, ViewSelect};
use primitive::{Primitive, resolve_as_scalar};
use std::cmp::{min, max, Ordering};

#[derive(Clone, Debug)]
pub struct Table {
    pub insert: Option<IndexSelect>,
    pub remove: Option<IndexSelect>,
}

#[derive(Clone, Debug)]
pub struct Union {
    pub selects: Vec<IndexSelect>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ConstraintOp {
    EQ,
    NEQ,
    LT,
    GT,
    LTE,
    GTE,
}

#[derive(Clone, Debug)]
pub struct Constraint {
    pub left: usize,
    pub op: ConstraintOp,
    pub right: usize,
}

impl Constraint {
    fn is_satisfied_by(&self, state: &[&Value]) -> bool {
        match self.op {
            ConstraintOp::EQ => state[self.left] == state[self.right],
            ConstraintOp::NEQ => state[self.left] != state[self.right],
            ConstraintOp::LT => state[self.left] < state[self.right],
            ConstraintOp::GT => state[self.left] > state[self.right],
            ConstraintOp::LTE => state[self.left] <= state[self.right],
            ConstraintOp::GTE => state[self.left] >= state[self.right],
        }
    }
}

#[derive(Clone, Debug)]
pub enum JoinSource {
    Relation{
        input: usize
    },
    Primitive{
        primitive: Primitive,
        arguments: Vec<usize>,
    },
}

#[derive(Clone, Debug)]
pub struct Join {
    pub constants: Vec<Value>,
    pub sources: Vec<JoinSource>,
    pub constraints: Vec<Vec<Constraint>>,
    pub select: ViewSelect,
    pub join2: Join2,
}

#[derive(Clone, Debug)]
pub enum Input {
    Primitive{
        primitive: Primitive,
        input_bindings: Vec<(usize, usize)>,
    },
    View{
        input_ix: usize,
    },
}

#[derive(Clone, Debug)]
pub struct Source {
    pub input: Input,
    pub constraint_bindings: Vec<(usize, usize)>,
    pub output_bindings: Vec<(usize, usize)>,
}

#[derive(Clone, Debug)]
pub struct Join2 {
    pub constants: Vec<Value>,
    pub sources: Vec<Source>,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub struct Reducer {
    pub primitive: Primitive,
    pub arguments: Vec<usize>,
}

#[derive(Clone, Debug, Copy)]
pub enum Direction {
    Ascending,
    Descending,
}

#[derive(Clone, Debug)]
pub struct Aggregate {
    pub constants: Vec<Value>,
    pub outer: IndexSelect,
    pub inner: IndexSelect,
    pub directions: Vec<Direction>,
    pub limit_from: Option<usize>,
    pub limit_to: Option<usize>,
    pub reducers: Vec<Reducer>,
    pub selects_inner: bool,
    pub select: ViewSelect,
}

#[derive(Clone, Debug)]
pub enum View {
    Table(Table),
    Union(Union),
    Join(Join),
    Aggregate(Aggregate),
    Join2(Join2),
}

fn push_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for value in input.iter() {
        state.push(value);
    }
}

fn pop_all<'a>(state: &mut Vec<&'a Value>, input: &'a Vec<Value>) {
    for _ in input.iter() {
        state.pop();
    }
}

fn join_step<'a>(join: &'a Join, ix: usize, inputs: &[&'a Relation], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.select(&state[..]));
    } else {
        match join.sources[ix] {
            JoinSource::Relation{input} => {
                for values in inputs[input].index.iter() {
                    push_all(state, values);
                    if join.constraints[ix].iter().all(|constraint| constraint.is_satisfied_by(&state[..])) {
                        join_step(join, ix+1, inputs, state, index)
                    }
                    pop_all(state, values);
                }
            }
            JoinSource::Primitive{ref primitive, ref arguments, ..} => {
                for values in primitive.eval_from_join(&arguments[..], &state[..]).into_iter() {
                    // promise the borrow checker that we will pop values before we exit this scope
                    // TODO this is not panic-safe - we should use CowString in Value instead
                    let values = unsafe { ::std::mem::transmute::<&Vec<Value>, &'a Vec<Value>>(&values) };
                    push_all(state, values);
                    if join.constraints[ix].iter().all(|constraint| constraint.is_satisfied_by(&state[..])) {
                        join_step(join, ix+1, inputs, state, index)
                    }
                    pop_all(state, values);
                }
            }
        }
    }
}

fn join_step2<'a>(join: &'a Join2, ix: usize, inputs: &[&'a Relation], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if ix == join.sources.len() {
        index.insert(join.select.select(&state[..]));
    } else {
        let source = &join.sources[ix];
        match source.input {
            Input::View{input_ix} => {
                for values in inputs[input_ix].index.iter() {
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        *state[variable_ix] == values[field_ix]
                    ) {
                        for &(field_ix, variable_ix) in source.output_bindings.iter() {
                            state[variable_ix] = &values[field_ix];
                        }
                        join_step2(join, ix+1, inputs, state, index);
                    }
                }
            }
            Input::Primitive{primitive, ref input_bindings} => {
                // values returned from primitives don't include inputs, so we will have to offset accesses by input_len
                let input_len = input_bindings.len();
                for values in primitive.eval_from_join2(&input_bindings[..], &state[..]).into_iter() {
                    // promise the borrow checker that we wont read these values after exiting this scope
                    // TODO this is not panic-safe - we should use CowString in Value instead
                    let values = unsafe { ::std::mem::transmute::<&Vec<Value>, &'a Vec<Value>>(&values) };
                    if source.constraint_bindings.iter().all(|&(field_ix, variable_ix)|
                        *state[variable_ix] == values[field_ix - input_len]
                    ) {
                        for &(field_ix, variable_ix) in source.output_bindings.iter() {
                            state[variable_ix] = &values[field_ix - input_len];
                        }
                        join_step2(join, ix+1, inputs, state, index);
                    }
                }
            }
        }
    }
}

fn aggregate_step<'a>(aggregate: &Aggregate, input_sets: &'a [&[Vec<Value>]], state: &mut Vec<&'a Value>, index: &mut BTreeSet<Vec<Value>>) {
    if input_sets.len() == 0 {
        index.insert(aggregate.select.select(&state[..]));
    } else {
        for values in input_sets[0].iter() {
            push_all(state, values);
            aggregate_step(aggregate, &input_sets[1..], state, index);
            pop_all(state, values);
        }
    }
}

fn compare_in_direction(xs: &[Value], ys: &[Value], directions: &[Direction]) -> Ordering {
    for ((x,y), direction) in xs.iter().zip(ys.iter()).zip(directions.iter()) {
        let cmp = match *direction {
            Direction::Ascending => x.cmp(y),
            Direction::Descending => y.cmp(x),
        };
        match cmp {
            Ordering::Greater => return Ordering::Greater,
            Ordering::Equal => (),
            Ordering::Less => return Ordering::Less,
        };
    }
    return Ordering::Equal;
}

impl View {
    pub fn run(&self, old_output: &Relation, inputs: &[&Relation]) -> Option<Relation> {
        let mut output = Relation::new(
            old_output.view.clone(),
            old_output.fields.clone(),
            old_output.names.clone()
            );
        match *self {
            View::Table(_) => None,
            View::Union(ref union) => {
                assert_eq!(union.selects.len(), inputs.len());
                for select in union.selects.iter() {
                    for values in select.select(&inputs[..]) {
                        output.index.insert(values);
                    }
                }
                Some(output)
            }
            View::Join(ref join) => {
                let mut state = join.constants.iter().collect();
                join_step(join, 0, inputs, &mut state, &mut output.index);
                let output2 = {
                    let mut output = Relation::new(
                        old_output.view.clone(),
                        old_output.fields.clone(),
                        old_output.names.clone()
                        );
                    let mut state = join.join2.constants.iter().collect();
                    join_step2(&join.join2, 0, inputs, &mut state, &mut output.index);
                    output
                };
                if output.index != output2.index {
                    println!("Failed for:");
                    println!("{:#?}", join);
                    println!("{:#?}", output);
                    println!("{:#?}", output2);
                }
                Some(output)
            }
            View::Join2(ref join) => {
                let mut state = join.constants.iter().collect();
                join_step2(join, 0, inputs, &mut state, &mut output.index);
                Some(output)
            }
            View::Aggregate(ref aggregate) => {
                let mut outer = aggregate.outer.select(&inputs[..]);
                let mut inner = aggregate.inner.select(&inputs[..]);
                outer.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                outer.dedup();
                inner.sort_by(|a,b| compare_in_direction(&a[..], &b[..], &aggregate.directions[..]));
                let constants = &aggregate.constants[..];
                let mut group_start = 0;
                for outer_values in outer.into_iter() {
                    let mut group_end = group_start;
                    while (group_end < inner.len())
                    && (inner[group_end][0..outer_values.len()] == outer_values[..]) {
                        group_end += 1;
                    }
                    let (group, output_values) = {
                        let limit_from = match aggregate.limit_from {
                            None => group_start,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_to = match aggregate.limit_to {
                            None => group_end,
                            Some(ix) => group_start
                                + resolve_as_scalar(ix, constants, &outer_values[..]).as_usize(),
                        };
                        let limit_from = min(max(limit_from, group_start), group_end);
                        let limit_to = min(max(limit_to, limit_from), group_end);
                        let group = &inner[limit_from..limit_to];
                        let output_values = aggregate.reducers.iter().map(|reducer| {
                            reducer.primitive.eval_from_aggregate(&reducer.arguments[..], constants, &outer_values[..], group)
                        }).collect::<Vec<_>>();
                        (group, output_values)
                    };
                    let mut output_sets = vec![];
                    let null = Value::Null;
                    let mut state = constants.iter().chain(outer_values.iter()).collect::<Vec<_>>();
                    if aggregate.selects_inner {
                        output_sets.push(group);
                    } else {
                        // nasty hack - fill null values for inner so that the ixes work out right
                        for _ in aggregate.inner.mapping.iter() {
                            state.push(&null);
                        }
                    }
                    for output in output_values.iter() {
                        output_sets.push(output);
                    }
                    aggregate_step(aggregate, &output_sets[..], &mut state, &mut output.index);
                    group_start = group_end;
                }
                Some(output)
            }
        }
    }
}<|endoftext|>


<commit_before>// Copyright (c) 2017 Colin Finck, RWTH Aachen University
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Implementation of the Allocator for dynamically allocating heap memory
//! in the kernel.
//!
//! The data structures used to manage heap memory require dynamic memory allocations
//! themselves. To solve this chicken-egg problem, eduOS-rs first uses a
//! "Bootstrap Allocator". This is a simple single-threaded implementation using some
//! preallocated space, along with an index variable.
//! Freed memory is never reused, but this can be neglected for bootstrapping.

use crate::arch::{BasePageSize, PageSize};
use crate::consts::*;
use crate::logging::*;
use crate::mm;
use alloc::alloc::Layout;
use core::alloc::GlobalAlloc;

/// Size of the preallocated space for the Bootstrap Allocator.
const BOOTSTRAP_HEAP_SIZE: usize = 0x1000;

/// Alignment of pointers returned by the Bootstrap Allocator.
/// Note that you also have to align the HermitAllocatorInfo structure!
const BOOTSTRAP_HEAP_ALIGNMENT: usize = CACHE_LINE;

/// The Allocator structure is immutable, so we need this helper structure
/// for our allocator information.
#[repr(align(64))]
#[repr(C)]
struct AllocatorInfo {
	heap: [u8; BOOTSTRAP_HEAP_SIZE],
	index: usize,
	is_bootstrapping: bool,
}

impl AllocatorInfo {
	const fn new() -> AllocatorInfo {
		AllocatorInfo {
			heap: [0xCC; BOOTSTRAP_HEAP_SIZE],
			index: 0,
			is_bootstrapping: true,
		}
	}

	fn switch_to_system_allocator(&mut self) {
		info!("Switching to the System Allocator");
		self.is_bootstrapping = false;
	}
}

static mut ALLOCATOR_INFO: AllocatorInfo = AllocatorInfo::new();

pub struct Allocator;

unsafe impl<'a> GlobalAlloc for &'a Allocator {
	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
		if ALLOCATOR_INFO.is_bootstrapping {
			alloc_bootstrap(layout)
		} else {
			alloc_system(layout)
		}
	}

	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
		let address = ptr as usize;

		// We never deallocate memory of the Bootstrap Allocator.
		// It would only increase the management burden and we wouldn't save
		// any significant amounts of memory.
		// So check if this is a pointer allocated by the System Allocator.
		debug!("Deallocate {} bytes at {:#X}", layout.size(), ptr as usize);
		if address >= mm::kernel_end_address() {
			dealloc_system(address, layout);
		}
	}
}

/// An allocation using the always available Bootstrap Allocator.
unsafe fn alloc_bootstrap(layout: Layout) -> *mut u8 {
	let ptr = &mut ALLOCATOR_INFO.heap[ALLOCATOR_INFO.index] as *mut u8;
	debug!(
		"Allocating {:#X} bytes at {:#X}, index {}",
		layout.size(),
		ptr as usize,
		ALLOCATOR_INFO.index
	);

	if ALLOCATOR_INFO.index + layout.size() >= BOOTSTRAP_HEAP_SIZE {
		panic!("Bootstrap Allocator Overflow! Increase BOOTSTRAP_HEAP_SIZE.");
	}

	// Bump the heap index and align it up to the next BOOTSTRAP_HEAP_ALIGNMENT boundary.
	ALLOCATOR_INFO.index = align_up!(
		ALLOCATOR_INFO.index + layout.size(),
		BOOTSTRAP_HEAP_ALIGNMENT
	);

	ptr
}

/// An allocation using the initialized System Allocator.
fn alloc_system(layout: Layout) -> *mut u8 {
	debug!(
		"Allocating {} bytes using the System Allocator",
		layout.size()
	);

	let size = align_up!(layout.size(), BasePageSize::SIZE);
	mm::allocate(size, true) as *mut u8
}

/// A deallocation using the initialized System Allocator.
fn dealloc_system(virtual_address: usize, layout: Layout) {
	debug!(
		"Deallocating {} bytes at {:#X} using the System Allocator",
		layout.size(),
		virtual_address
	);

	let size = align_up!(layout.size(), BasePageSize::SIZE);
	mm::deallocate(virtual_address, size);
}

pub fn init() {
	unsafe {
		ALLOCATOR_INFO.switch_to_system_allocator();
	}
}
<commit_msg>disable temporary system allocator<commit_after>// Copyright (c) 2017 Colin Finck, RWTH Aachen University
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Implementation of the Allocator for dynamically allocating heap memory
//! in the kernel.
//!
//! The data structures used to manage heap memory require dynamic memory allocations
//! themselves. To solve this chicken-egg problem, eduOS-rs first uses a
//! "Bootstrap Allocator". This is a simple single-threaded implementation using some
//! preallocated space, along with an index variable.
//! Freed memory is never reused, but this can be neglected for bootstrapping.

use crate::arch::{BasePageSize, PageSize};
use crate::consts::*;
use crate::logging::*;
use crate::mm;
use alloc::alloc::Layout;
use core::alloc::GlobalAlloc;

/// Size of the preallocated space for the Bootstrap Allocator.
const BOOTSTRAP_HEAP_SIZE: usize = 0x140000;

/// Alignment of pointers returned by the Bootstrap Allocator.
/// Note that you also have to align the HermitAllocatorInfo structure!
const BOOTSTRAP_HEAP_ALIGNMENT: usize = CACHE_LINE;

/// The Allocator structure is immutable, so we need this helper structure
/// for our allocator information.
#[repr(align(64))]
#[repr(C)]
struct AllocatorInfo {
	heap: [u8; BOOTSTRAP_HEAP_SIZE],
	index: usize,
	is_bootstrapping: bool,
}

impl AllocatorInfo {
	const fn new() -> AllocatorInfo {
		AllocatorInfo {
			heap: [0x00; BOOTSTRAP_HEAP_SIZE],
			index: 0,
			is_bootstrapping: true,
		}
	}

	fn switch_to_system_allocator(&mut self) {
		info!("Switching to the System Allocator");
		self.is_bootstrapping = false;
	}
}

static mut ALLOCATOR_INFO: AllocatorInfo = AllocatorInfo::new();

pub struct Allocator;

unsafe impl<'a> GlobalAlloc for &'a Allocator {
	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
		if ALLOCATOR_INFO.is_bootstrapping {
			alloc_bootstrap(layout)
		} else {
			alloc_system(layout)
		}
	}

	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
		let address = ptr as usize;

		// We never deallocate memory of the Bootstrap Allocator.
		// It would only increase the management burden and we wouldn't save
		// any significant amounts of memory.
		// So check if this is a pointer allocated by the System Allocator.
		debug!("Deallocate {} bytes at {:#X}", layout.size(), ptr as usize);
		if address >= mm::kernel_end_address() {
			dealloc_system(address, layout);
		}
	}
}

/// An allocation using the always available Bootstrap Allocator.
unsafe fn alloc_bootstrap(layout: Layout) -> *mut u8 {
	let ptr = &mut ALLOCATOR_INFO.heap[ALLOCATOR_INFO.index] as *mut u8;
	debug!(
		"Allocating {:#X} bytes at {:#X}, index {}",
		layout.size(),
		ptr as usize,
		ALLOCATOR_INFO.index
	);

	if ALLOCATOR_INFO.index + layout.size() >= BOOTSTRAP_HEAP_SIZE {
		panic!("Bootstrap Allocator Overflow! Increase BOOTSTRAP_HEAP_SIZE.");
	}

	// Bump the heap index and align it up to the next BOOTSTRAP_HEAP_ALIGNMENT boundary.
	ALLOCATOR_INFO.index = align_up!(
		ALLOCATOR_INFO.index + layout.size(),
		BOOTSTRAP_HEAP_ALIGNMENT
	);

	ptr
}

/// An allocation using the initialized System Allocator.
fn alloc_system(layout: Layout) -> *mut u8 {
	debug!(
		"Allocating {} bytes using the System Allocator",
		layout.size()
	);

	let size = align_up!(layout.size(), BasePageSize::SIZE);
	mm::allocate(size, true) as *mut u8
}

/// A deallocation using the initialized System Allocator.
fn dealloc_system(virtual_address: usize, layout: Layout) {
	debug!(
		"Deallocating {} bytes at {:#X} using the System Allocator",
		layout.size(),
		virtual_address
	);

	let size = align_up!(layout.size(), BasePageSize::SIZE);
	mm::deallocate(virtual_address, size);
}

pub fn init() {
	/*unsafe {
		ALLOCATOR_INFO.switch_to_system_allocator();
	}*/
}
<|endoftext|>


<commit_before>use super::*;

pub trait Notification {
    type Params;
    const METHOD: &'static str;
}

#[macro_export]
macro_rules! lsp_notification {
    ("$/cancelRequest") => { $crate::notification::Cancel };
    ("initialized") => { $crate::notification::Initialized };
    ("exit") => { $crate::notification::Exit };

    ("window/showMessage") => { $crate::notification::ShowMessage };
    ("window/logMessage") => { $crate::notification::LogMessage };

    ("telemetry/event") => { $crate::notification::Event };

    ("client/registerCapability") => { $crate::notification::RegistrationParams };
    ("client/unregisterCapability") => { $crate::notification::UnregistrationParams };

    ("textDocument/didOpen") => { $crate::notification::DidOpenTextDocument };
    ("textDocument/didChange") => { $crate::notification::DidChangeTextDocument };
    ("textDocument/willSave") => { $crate::notification::WillSaveTextDocument };
    ("textDocument/didSave") => { $crate::notification::DidSaveTextDocument };
    ("textDocument/didClose") => { $crate::notification::DidCloseTextDocument };
    ("textDocument/publishDiagnostics") => { $crate::notification::PublishDiagnostics };

    ("workspace/didChangeConfiguration") => { $crate::notification::DidChangeConfiguration };
    ("workspace/didChangeWatchedFiles") => { $crate::notification::DidChangeWatchedFiles };
}

/// The base protocol now offers support for request cancellation. To cancel a request,
/// a notification message with the following properties is sent:
///
/// A request that got canceled still needs to return from the server and send a response back.
/// It can not be left open / hanging. This is in line with the JSON RPC protocol that requires
/// that every request sends a response back. In addition it allows for returning partial results on cancel.
#[derive(Debug)]
pub enum Cancel {}

impl Notification for Cancel {
    type Params = CancelParams;
    const METHOD: &'static str = "$/cancelRequest";
}

/// The initialized notification is sent from the client to the server after the client received
/// the result of the initialize request but before the client is sending any other request or
/// notification to the server. The server can use the initialized notification for example to
/// dynamically register capabilities.
#[derive(Debug)]
pub enum Initialized {}

impl Notification for Initialized {
    type Params = ();
    const METHOD: &'static str = "initialized";
}

/**
 * A notification to ask the server to exit its process.
 * The server should exit with success code 0 if the shutdown request has been received before;
 * otherwise with error code 1.
 */
#[derive(Debug)]
pub enum Exit {}

impl Notification for Exit {
    type Params = ();
    const METHOD: &'static str = "exit";
}

/**
 * The show message notification is sent from a server to a client to ask the client to display a particular message
 * in the user interface.
 */
#[derive(Debug)]
pub enum ShowMessage {}

impl Notification for ShowMessage {
    type Params = ShowMessageParams;
    const METHOD: &'static str = "window/showMessage";
}

/**
 * The log message notification is sent from the server to the client to ask the client to log a particular message.
 */
#[derive(Debug)]
pub enum LogMessage {}

impl Notification for LogMessage {
    type Params = LogMessageParams;
    const METHOD: &'static str = "window/logMessage";
}

/**
 * The telemetry notification is sent from the server to the client to ask the client to log a telemetry event.
 */
#[derive(Debug)]
pub enum TelemetryEvent {}

impl Notification for TelemetryEvent {
    type Params = ();
    const METHOD: &'static str = "telemetry/event";
}

/**
 * The client/registerCapability request is sent from the server to the client to register for a new capability on the client side. Not all clients need to support dynamic capability registration. A client opts in via the ClientCapabilities.GenericCapability property.
 */
#[derive(Debug)]
pub enum RegisterCapability {}

impl Notification for RegisterCapability {
    type Params = RegistrationParams;
    const METHOD: &'static str = "client/registerCapability";
}

/// The client/unregisterCapability request is sent from the server to the client to unregister a
/// previously register capability.
#[derive(Debug)]
pub enum UnregisterCapability {}

impl Notification for UnregisterCapability {
    type Params = UnregistrationParams;
    const METHOD: &'static str = "client/unregisterCapability";
}

/**
 * A notification sent from the client to the server to signal the change of configuration settings.
 */
#[derive(Debug)]
pub enum DidChangeConfiguration {}

impl Notification for DidChangeConfiguration {
    type Params = DidChangeConfigurationParams;
    const METHOD: &'static str = "workspace/didChangeConfiguration";
}

/**
 * The document open notification is sent from the client to the server to signal newly opened text documents.
 * The document's truth is now managed by the client and the server must not try to read the document's truth
 * using the document's uri.
 */
#[derive(Debug)]
pub enum DidOpenTextDocument {}

impl Notification for DidOpenTextDocument {
    type Params = DidOpenTextDocumentParams;
    const METHOD: &'static str = "textDocument/didOpen";
}

/**
 * The document change notification is sent from the client to the server to signal changes to a text document.
 * In 2.0 the shape of the params has changed to include proper version numbers and language ids.
 */
#[derive(Debug)]
pub enum DidChangeTextDocument {}

impl Notification for DidChangeTextDocument {
    type Params = DidChangeTextDocumentParams;
    const METHOD: &'static str = "textDocument/didChange";
}

/// The document will save notification is sent from the client to the server before the document
/// is actually saved.
#[derive(Debug)]
pub enum WillSave {}

impl Notification for WillSave {
    type Params = ();
    const METHOD: &'static str = "textDocument/willSave";
}

/// The document will save request is sent from the client to the server before the document is
/// actually saved. The request can return an array of TextEdits which will be applied to the text
/// document before it is saved. Please note that clients might drop results if computing the text
/// edits took too long or if a server constantly fails on this request. This is done to keep the
/// save fast and reliable.
#[derive(Debug)]
pub enum WillSaveWaitUntil {}

impl Notification for WillSaveWaitUntil {
    type Params = ();
    const METHOD: &'static str = "textDocument/willSaveWaitUntil";
}

/**
 * The document close notification is sent from the client to the server when the document got closed in the client.
 * The document's truth now exists where the document's uri points to (e.g. if the document's uri is a file uri
 * the truth now exists on disk).
 */
#[derive(Debug)]
pub enum DidCloseTextDocument {}

impl Notification for DidCloseTextDocument {
    type Params = DidCloseTextDocumentParams;
    const METHOD: &'static str = "textDocument/didClose";
}

/**
 * The document save notification is sent from the client to the server when the document was saved in the client.
 */
#[derive(Debug)]
pub enum DidSaveTextDocument {}

impl Notification for DidSaveTextDocument {
    type Params = DidSaveTextDocumentParams;
    const METHOD: &'static str = "textDocument/didSave";
}

/**
 * The watched files notification is sent from the client to the server when the client detects changes to files
 * watched by the language client.
 */
#[derive(Debug)]
pub enum DidChangeWatchedFiles {}

impl Notification for DidChangeWatchedFiles {
    type Params = DidChangeWatchedFilesParams;
    const METHOD: &'static str = "workspace/didChangeWatchedFiles";
}

/**
 * Diagnostics notification are sent from the server to the client to signal results of validation runs.
 */
#[derive(Debug)]
pub enum PublishDiagnostics {}

impl Notification for PublishDiagnostics {
    type Params = PublishDiagnosticsParams;
    const METHOD: &'static str = "textDocument/publishDiagnostics";
}
<commit_msg>Fix [un]registerCapabilit<commit_after>use super::*;

pub trait Notification {
    type Params;
    const METHOD: &'static str;
}

#[macro_export]
macro_rules! lsp_notification {
    ("$/cancelRequest") => { $crate::notification::Cancel };
    ("initialized") => { $crate::notification::Initialized };
    ("exit") => { $crate::notification::Exit };

    ("window/showMessage") => { $crate::notification::ShowMessage };
    ("window/logMessage") => { $crate::notification::LogMessage };

    ("telemetry/event") => { $crate::notification::Event };

    ("client/registerCapability") => { $crate::notification::RegisterCapability };
    ("client/unregisterCapability") => { $crate::notification::UnregisterCapability };

    ("textDocument/didOpen") => { $crate::notification::DidOpenTextDocument };
    ("textDocument/didChange") => { $crate::notification::DidChangeTextDocument };
    ("textDocument/willSave") => { $crate::notification::WillSaveTextDocument };
    ("textDocument/didSave") => { $crate::notification::DidSaveTextDocument };
    ("textDocument/didClose") => { $crate::notification::DidCloseTextDocument };
    ("textDocument/publishDiagnostics") => { $crate::notification::PublishDiagnostics };

    ("workspace/didChangeConfiguration") => { $crate::notification::DidChangeConfiguration };
    ("workspace/didChangeWatchedFiles") => { $crate::notification::DidChangeWatchedFiles };
}

/// The base protocol now offers support for request cancellation. To cancel a request,
/// a notification message with the following properties is sent:
///
/// A request that got canceled still needs to return from the server and send a response back.
/// It can not be left open / hanging. This is in line with the JSON RPC protocol that requires
/// that every request sends a response back. In addition it allows for returning partial results on cancel.
#[derive(Debug)]
pub enum Cancel {}

impl Notification for Cancel {
    type Params = CancelParams;
    const METHOD: &'static str = "$/cancelRequest";
}

/// The initialized notification is sent from the client to the server after the client received
/// the result of the initialize request but before the client is sending any other request or
/// notification to the server. The server can use the initialized notification for example to
/// dynamically register capabilities.
#[derive(Debug)]
pub enum Initialized {}

impl Notification for Initialized {
    type Params = ();
    const METHOD: &'static str = "initialized";
}

/**
 * A notification to ask the server to exit its process.
 * The server should exit with success code 0 if the shutdown request has been received before;
 * otherwise with error code 1.
 */
#[derive(Debug)]
pub enum Exit {}

impl Notification for Exit {
    type Params = ();
    const METHOD: &'static str = "exit";
}

/**
 * The show message notification is sent from a server to a client to ask the client to display a particular message
 * in the user interface.
 */
#[derive(Debug)]
pub enum ShowMessage {}

impl Notification for ShowMessage {
    type Params = ShowMessageParams;
    const METHOD: &'static str = "window/showMessage";
}

/**
 * The log message notification is sent from the server to the client to ask the client to log a particular message.
 */
#[derive(Debug)]
pub enum LogMessage {}

impl Notification for LogMessage {
    type Params = LogMessageParams;
    const METHOD: &'static str = "window/logMessage";
}

/**
 * The telemetry notification is sent from the server to the client to ask the client to log a telemetry event.
 */
#[derive(Debug)]
pub enum TelemetryEvent {}

impl Notification for TelemetryEvent {
    type Params = ();
    const METHOD: &'static str = "telemetry/event";
}

/**
 * The client/registerCapability request is sent from the server to the client to register for a new capability on the client side. Not all clients need to support dynamic capability registration. A client opts in via the ClientCapabilities.GenericCapability property.
 */
#[derive(Debug)]
pub enum RegisterCapability {}

impl Notification for RegisterCapability {
    type Params = RegistrationParams;
    const METHOD: &'static str = "client/registerCapability";
}

/// The client/unregisterCapability request is sent from the server to the client to unregister a
/// previously register capability.
#[derive(Debug)]
pub enum UnregisterCapability {}

impl Notification for UnregisterCapability {
    type Params = UnregistrationParams;
    const METHOD: &'static str = "client/unregisterCapability";
}

/**
 * A notification sent from the client to the server to signal the change of configuration settings.
 */
#[derive(Debug)]
pub enum DidChangeConfiguration {}

impl Notification for DidChangeConfiguration {
    type Params = DidChangeConfigurationParams;
    const METHOD: &'static str = "workspace/didChangeConfiguration";
}

/**
 * The document open notification is sent from the client to the server to signal newly opened text documents.
 * The document's truth is now managed by the client and the server must not try to read the document's truth
 * using the document's uri.
 */
#[derive(Debug)]
pub enum DidOpenTextDocument {}

impl Notification for DidOpenTextDocument {
    type Params = DidOpenTextDocumentParams;
    const METHOD: &'static str = "textDocument/didOpen";
}

/**
 * The document change notification is sent from the client to the server to signal changes to a text document.
 * In 2.0 the shape of the params has changed to include proper version numbers and language ids.
 */
#[derive(Debug)]
pub enum DidChangeTextDocument {}

impl Notification for DidChangeTextDocument {
    type Params = DidChangeTextDocumentParams;
    const METHOD: &'static str = "textDocument/didChange";
}

/// The document will save notification is sent from the client to the server before the document
/// is actually saved.
#[derive(Debug)]
pub enum WillSave {}

impl Notification for WillSave {
    type Params = ();
    const METHOD: &'static str = "textDocument/willSave";
}

/// The document will save request is sent from the client to the server before the document is
/// actually saved. The request can return an array of TextEdits which will be applied to the text
/// document before it is saved. Please note that clients might drop results if computing the text
/// edits took too long or if a server constantly fails on this request. This is done to keep the
/// save fast and reliable.
#[derive(Debug)]
pub enum WillSaveWaitUntil {}

impl Notification for WillSaveWaitUntil {
    type Params = ();
    const METHOD: &'static str = "textDocument/willSaveWaitUntil";
}

/**
 * The document close notification is sent from the client to the server when the document got closed in the client.
 * The document's truth now exists where the document's uri points to (e.g. if the document's uri is a file uri
 * the truth now exists on disk).
 */
#[derive(Debug)]
pub enum DidCloseTextDocument {}

impl Notification for DidCloseTextDocument {
    type Params = DidCloseTextDocumentParams;
    const METHOD: &'static str = "textDocument/didClose";
}

/**
 * The document save notification is sent from the client to the server when the document was saved in the client.
 */
#[derive(Debug)]
pub enum DidSaveTextDocument {}

impl Notification for DidSaveTextDocument {
    type Params = DidSaveTextDocumentParams;
    const METHOD: &'static str = "textDocument/didSave";
}

/**
 * The watched files notification is sent from the client to the server when the client detects changes to files
 * watched by the language client.
 */
#[derive(Debug)]
pub enum DidChangeWatchedFiles {}

impl Notification for DidChangeWatchedFiles {
    type Params = DidChangeWatchedFilesParams;
    const METHOD: &'static str = "workspace/didChangeWatchedFiles";
}

/**
 * Diagnostics notification are sent from the server to the client to signal results of validation runs.
 */
#[derive(Debug)]
pub enum PublishDiagnostics {}

impl Notification for PublishDiagnostics {
    type Params = PublishDiagnosticsParams;
    const METHOD: &'static str = "textDocument/publishDiagnostics";
}
<|endoftext|>


<commit_before>use crate::data_provider::OnDiskDataProvider;
use crate::errors::Result;
use crate::iterator::{ParallelIterator, PointQuery};
use crate::octree::{build_octree, Octree};
use crate::{AttributeData, NumberOfPoints, PointsBatch};
use cgmath::{EuclideanSpace, Point3, Vector3};
use collision::{Aabb, Aabb3};
use tempdir::TempDir;

const NUM_POINTS: usize = 100_001;

impl NumberOfPoints for std::vec::IntoIter<PointsBatch> {
    fn num_points(&self) -> usize {
        self.clone().map(|b| b.position.len()).sum()
    }
}

fn build_big_test_octree() -> Octree {
    let mut batch = PointsBatch {
        //ECEF parking lot porter dr
        position: vec![Vector3::new(-2_699_182.0, -4_294_938.0, 3_853_373.0); 2 * NUM_POINTS],
        attributes: vec![(
            "color".to_string(),
            AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); 2 * NUM_POINTS]),
        )]
        .into_iter()
        .collect(),
    };

    batch.position[3] = Vector3::new(-2_702_846.0, -4_291_151.0, 3_855_012.0); // ECEF STANFORD

    let p = Point3::new(6_400_000.0, 6_400_000.0, 6_400_000.0);
    let bounding_box = Aabb3::new(-1.0 * p, p);

    let tmp_dir = TempDir::new("octree").unwrap();

    build_octree(
        &tmp_dir,
        1.0,
        bounding_box,
        vec![batch].into_iter(),
        &["color"],
    );
    Octree::from_data_provider(Box::new(OnDiskDataProvider {
        directory: tmp_dir.into_path(),
    }))
    .unwrap()
}

fn build_test_octree() -> Octree {
    let mut batch = PointsBatch {
        position: vec![Vector3::new(0.0, 0.0, 0.0); NUM_POINTS],
        attributes: vec![(
            "color".to_string(),
            AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); NUM_POINTS]),
        )]
        .into_iter()
        .collect(),
    };

    batch.position[NUM_POINTS - 1] = Vector3::new(-200., -40., 30.);

    let bounding_box = Aabb3::zero().grow(Point3::from_vec(batch.position[NUM_POINTS - 1]));

    let tmp_dir = TempDir::new("octree").unwrap();

    build_octree(
        &tmp_dir,
        1.0,
        bounding_box,
        vec![batch].into_iter(),
        &["color"],
    );
    Octree::from_data_provider(Box::new(OnDiskDataProvider {
        directory: tmp_dir.into_path(),
    }))
    .unwrap()
}

#[test]
fn test_batch_iterator() {
    let batch_size = 5000;
    // define function
    let mut callback_count: usize = 0;
    let max_num_points = 13_000; // 2*batch size + 3000
    let mut delivered_points = 0;
    println!(
        "batch_size= {} ,  num_points= {}",
        batch_size, max_num_points
    );
    let callback_func = |points_batch: PointsBatch| -> Result<()> {
        callback_count += 1;
        delivered_points += points_batch.position.len();
        println!(
            "Callback_count {:}, delivered points {:}",
            callback_count, delivered_points
        );
        if delivered_points >= max_num_points {
            println!("Callback: Max Points reached!");
            return Err(std::io::Error::new(
                std::io::ErrorKind::Interrupted,
                format!("Maximum number of {} points reached.", max_num_points),
            )
            .into());
        }
        Ok(())
    };

    // octree and iterator
    let octree = build_test_octree();
    let location = PointQuery {
        attributes: vec!["color"],
        ..Default::default()
    };
    let octree_slice: &[Octree] = std::slice::from_ref(&octree);
    let mut parallel_iterator = ParallelIterator::new(
        octree_slice,
        &location,
        batch_size,
        std::cmp::max(1, num_cpus::get() - 1),
        4,
    );

    let _err_stop = parallel_iterator
        .try_for_each_batch(callback_func)
        .expect_err("Test error");

    //  requiring and returning less points than the total
    assert!(3 <= callback_count);
    assert!(3 * batch_size <= delivered_points);
    assert!(callback_count >= 3);
    assert!((3 + 1) * batch_size >= delivered_points);
    assert!(delivered_points as i32 - max_num_points as i32 >= 0);
}

#[test]
fn test_batch_iterator_more_points() {
    let batch_size = NUM_POINTS / 2;
    // define function

    // requiring more points than in the octree
    let mut callback_count = 0;
    let mut delivered_points = 0;
    let max_num_points = NUM_POINTS + 30_000;
    println!(
        "batch_size= {} ,  num_points= {}",
        batch_size, max_num_points
    );
    let callback_func = |points_batch: PointsBatch| -> Result<()> {
        callback_count += 1;
        delivered_points += points_batch.position.len();
        println!(
            "Callback_count {:}, delivered points {:}",
            callback_count, delivered_points
        );
        if delivered_points >= max_num_points {
            println!("Callback: Max Points reached!");
            return Err(std::io::Error::new(
                std::io::ErrorKind::Interrupted,
                format!("Maximum number of {} points reached.", max_num_points),
            )
            .into());
        }
        Ok(())
    };
    // octree and iterator
    let octree = build_test_octree();
    let location = PointQuery {
        attributes: vec!["color"],
        ..Default::default()
    };

    let octree_slice: &[Octree] = std::slice::from_ref(&octree);
    let mut parallel_iterator = ParallelIterator::new(octree_slice, &location, batch_size, 2, 2);

    let _err_stop = parallel_iterator
        .try_for_each_batch(callback_func)
        .expect("Test OK");
    assert!(delivered_points == NUM_POINTS); // only delivers all the points in the octree
}

#[test]
//#[ignore]
fn test_batch_iterator_big_octree() {
    let batch_size = 5000;
    // define function
    let mut callback_count: u64 = 0;
    let max_num_points = 13_001; // 2*batch size + 3000
    let mut delivered_points: u64 = 0;
    println!(
        "batch_size= {} ,  num_points= {}",
        batch_size, max_num_points
    );
    let callback_func = |points_batch: PointsBatch| -> Result<()> {
        callback_count += 1;
        delivered_points += points_batch.position.len() as u64;
        println!(
            "Callback_count {:}, delivered points {:}",
            callback_count, delivered_points
        );
        if delivered_points >= max_num_points {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Interrupted,
                format!("Maximum number of {} points reached.", max_num_points),
            )
            .into());
        }
        Ok(())
    };

    // octree and iterator
    let octree = build_big_test_octree();
    let location = PointQuery {
        attributes: vec!["color"],
        ..Default::default()
    };
    let octree_slice: &[Octree] = std::slice::from_ref(&octree);
    let mut parallel_iterator = ParallelIterator::new(
        octree_slice,
        &location,
        batch_size,
        std::cmp::max(1, num_cpus::get() - 1),
        4,
    );

    let _err_stop = parallel_iterator
        .try_for_each_batch(callback_func)
        .expect_err("Test error");

    assert!(callback_count >= 3);
    assert!((3 + 1) * batch_size as u64 >= delivered_points);
    assert!(delivered_points as i32 - max_num_points as i32 >= 0);
    assert!(delivered_points as i32 - max_num_points as i32 <= batch_size as i32);
}
<commit_msg>Streamline tests (#416)<commit_after>use crate::data_provider::OnDiskDataProvider;
use crate::errors::Result;
use crate::iterator::{ParallelIterator, PointQuery};
use crate::octree::{build_octree, Octree};
use crate::{AttributeData, NumberOfPoints, PointsBatch};
use cgmath::{EuclideanSpace, Point3, Vector3};
use collision::{Aabb, Aabb3};
use tempdir::TempDir;

const NUM_POINTS: usize = 100_001;

impl NumberOfPoints for std::vec::IntoIter<PointsBatch> {
    fn num_points(&self) -> usize {
        self.clone().map(|b| b.position.len()).sum()
    }
}

fn build_test_octree() -> Octree {
    let mut batch = PointsBatch {
        position: vec![Vector3::new(0.0, 0.0, 0.0); NUM_POINTS],
        attributes: vec![(
            "color".to_string(),
            AttributeData::U8Vec3(vec![Vector3::new(255, 0, 0); NUM_POINTS]),
        )]
        .into_iter()
        .collect(),
    };

    batch.position[NUM_POINTS - 1] = Vector3::new(-200., -40., 30.);

    let bounding_box = Aabb3::zero().grow(Point3::from_vec(batch.position[NUM_POINTS - 1]));

    let tmp_dir = TempDir::new("octree").unwrap();

    build_octree(
        &tmp_dir,
        1.0,
        bounding_box,
        vec![batch].into_iter(),
        &["color"],
    );
    Octree::from_data_provider(Box::new(OnDiskDataProvider {
        directory: tmp_dir.into_path(),
    }))
    .unwrap()
}

struct Consumer {
    max_num_points: usize,
    num_received_points: usize,
    num_received_callbacks: usize,
}

impl Consumer {
    fn new(max_num_points: usize) -> Self {
        Self {
            max_num_points,
            num_received_points: 0,
            num_received_callbacks: 0,
        }
    }

    // Consumes points until it reaches max_num_points, then errors out
    fn consume(&mut self, points_batch: PointsBatch) -> Result<()> {
        self.num_received_callbacks += 1;
        self.num_received_points += points_batch.position.len();
        println!(
            "Callback_count {}, delivered points {}",
            self.num_received_callbacks, self.num_received_points
        );
        if self.num_received_points >= self.max_num_points {
            println!("Callback: Max Points reached!");
            return Err(std::io::Error::new(
                std::io::ErrorKind::Interrupted,
                format!("Maximum number of {} points reached.", self.max_num_points),
            )
            .into());
        }
        Ok(())
    }
}

#[test]
fn test_batch_iterator() {
    let batch_size = 5000;
    let max_num_points = 13_000;
    let mut c = Consumer::new(max_num_points);
    // octree and iterator
    let octree = build_test_octree();
    let location = PointQuery {
        attributes: vec!["color"],
        ..Default::default()
    };
    let octree_slice: &[Octree] = std::slice::from_ref(&octree);
    let mut parallel_iterator = ParallelIterator::new(
        octree_slice,
        &location,
        batch_size,
        std::cmp::max(1, num_cpus::get() - 1),
        4,
    );

    parallel_iterator
        .try_for_each_batch(|points_batch| c.consume(points_batch))
        .expect_err("Iterator did not error even though callback errored.");

    assert!(c.num_received_points >= c.max_num_points);
    // The number of points doesn't fit in two batches
    assert_eq!(c.num_received_callbacks, 3);
    // The callback received full batches
    assert_eq!(c.num_received_points, 3 * batch_size);
}

#[test]
fn test_batch_iterator_more_points() {
    let batch_size = NUM_POINTS / 2;
    // This test asks for more points than exist in the octree.
    // Expected result: The number of points returned is equal to the number of points in the octree.
    let max_num_points = NUM_POINTS + 30_000;
    let mut c = Consumer::new(max_num_points);

    // octree and iterator
    let octree = build_test_octree();
    let location = PointQuery {
        attributes: vec!["color"],
        ..Default::default()
    };

    let octree_slice: &[Octree] = std::slice::from_ref(&octree);
    let mut parallel_iterator = ParallelIterator::new(octree_slice, &location, batch_size, 2, 2);

    parallel_iterator
        .try_for_each_batch(|points_batch| c.consume(points_batch))
        .expect("Iterator errored even though callback should not have errored.");
    assert_eq!(c.num_received_points, NUM_POINTS);
}
<|endoftext|>