use csv::ByteRecord;
use super::error::{ConcretizationError, SpecifiedEvaluationError};
use super::interpreter::{concretize_expression, eval_expression, ConcreteExpr};
use super::parser::parse_named_expressions;
use super::types::HeadersIndex;
#[derive(Clone, Debug)]
pub struct SelectionProgram {
exprs: Vec<(ConcreteExpr, String, bool)>,
headers_index: HeadersIndex,
mask: Vec<Option<usize>>,
rest: Vec<usize>,
}
impl SelectionProgram {
pub fn parse(code: &str, headers: &ByteRecord) -> Result<Self, ConcretizationError> {
let headers_index = HeadersIndex::from_headers(headers);
let mut mask = vec![None; headers.len()];
let mut rest = vec![];
let exprs = match parse_named_expressions(code) {
Err(err) => return Err(ConcretizationError::ParseError(err)),
Ok(parsed_exprs) => parsed_exprs
.into_iter()
.enumerate()
.map(|(expr_i, (expr, expr_name))| {
concretize_expression(expr, headers, None).map(|c| {
let pos = headers_index.get_first_by_name(&expr_name);
if let Some(i) = pos {
mask[i] = Some(expr_i);
} else {
mask.push(None);
rest.push(expr_i);
}
(c, expr_name, pos.is_some())
})
})
.collect::<Result<Vec<_>, _>>(),
}?;
Ok(Self {
exprs,
headers_index,
mask,
rest,
})
}
pub fn headers(&self) -> impl Iterator<Item = &[u8]> {
self.exprs.iter().map(|(_, name, _)| name.as_bytes())
}
pub fn new_headers(&self) -> impl Iterator<Item = &[u8]> {
self.exprs
.iter()
.filter_map(|(_, name, already_exists)| (!already_exists).then_some(name.as_bytes()))
}
pub fn has_something_to_overwrite(&self) -> bool {
self.rest.len() < self.exprs.len()
}
pub fn extend_into(
&self,
index: usize,
record: &ByteRecord,
output_record: &mut ByteRecord,
) -> Result<(), SpecifiedEvaluationError> {
for (expr, _, _) in self.exprs.iter() {
let value = eval_expression(expr, Some(index), record, &self.headers_index)?;
output_record.push_field(&value.serialize_as_bytes());
}
Ok(())
}
pub fn extend(
&self,
index: usize,
record: &mut ByteRecord,
) -> Result<(), SpecifiedEvaluationError> {
for (expr, _, _) in self.exprs.iter() {
let value = eval_expression(expr, Some(index), record, &self.headers_index)?;
record.push_field(&value.serialize_as_bytes());
}
Ok(())
}
pub fn overwrite(
&self,
index: usize,
record: &mut ByteRecord,
) -> Result<ByteRecord, SpecifiedEvaluationError> {
let mut new_record = ByteRecord::new();
for (expr_i_opt, cell) in self.mask.iter().copied().zip(record.iter()) {
if let Some(expr_i) = expr_i_opt {
let expr = &self.exprs[expr_i].0;
let value = eval_expression(expr, Some(index), record, &self.headers_index)?;
new_record.push_field(&value.serialize_as_bytes());
} else {
new_record.push_field(cell);
}
}
for expr_i in self.rest.iter().copied() {
let expr = &self.exprs[expr_i].0;
let value = eval_expression(expr, Some(index), record, &self.headers_index)?;
new_record.push_field(&value.serialize_as_bytes());
}
Ok(new_record)
}
}