use itertools::Itertools;
use crate::{
common::{Spanned, error::TypeQLError},
parser::{
IntoChildNodes, Node, Rule, RuleMatcher,
annotation::visit_annotations,
pipeline::{visit_query_stage, visit_reducer},
type_::visit_named_type_any,
visit_identifier, visit_var, visit_vars,
},
schema::definable::{
Function,
function::{
Argument, Check, FunctionBlock, Output, ReturnReduction, ReturnSingle, ReturnStatement, ReturnStream,
Signature, Single, SingleSelector, Stream,
},
},
};
pub(in crate::parser) fn visit_definition_function(node: Node<'_>) -> Function {
debug_assert_eq!(node.as_rule(), Rule::definition_function);
let span = node.span();
let unparsed = node.as_span().as_str().to_owned();
let mut children = node.into_children();
children.skip_expected(Rule::FUN);
let signature = visit_function_signature(children.consume_expected(Rule::function_signature));
let annotations = children.try_consume_expected(Rule::annotations).map(visit_annotations).unwrap_or_default();
let block = visit_function_block(children.consume_expected(Rule::function_block));
debug_assert_eq!(children.try_consume_any(), None);
Function::new(span, signature, annotations, block, unparsed)
}
pub fn visit_function_block(node: Node<'_>) -> FunctionBlock {
debug_assert_eq!(node.as_rule(), Rule::function_block);
let span = node.span();
let mut children = node.into_children();
let stages = children.take_while_ref(|node| node.as_rule() == Rule::query_stage).map(visit_query_stage).collect();
let return_stmt = visit_return_statement(children.consume_expected(Rule::return_statement));
debug_assert_eq!(children.try_consume_any(), None);
FunctionBlock::new(span, stages, return_stmt)
}
fn visit_return_statement(node: Node<'_>) -> ReturnStatement {
debug_assert_eq!(node.as_rule(), Rule::return_statement);
let mut children = node.into_children();
let child = children.consume_any();
let return_stmt = match child.as_rule() {
Rule::return_stream => ReturnStatement::Stream(visit_return_stream(child)),
Rule::return_single => ReturnStatement::Single(visit_return_single(child)),
Rule::return_reduce => ReturnStatement::Reduce(visit_return_reduce(child)),
_ => unreachable!("{}", TypeQLError::IllegalGrammar { input: child.as_str().to_owned() }),
};
debug_assert_eq!(children.try_consume_any(), None);
return_stmt
}
fn visit_return_stream(node: Node<'_>) -> ReturnStream {
debug_assert_eq!(node.as_rule(), Rule::return_stream);
let span = node.span();
let mut children = node.into_children();
children.skip_expected(Rule::RETURN);
let vars = visit_vars(children.consume_any());
debug_assert_eq!(children.try_consume_any(), None);
ReturnStream::new(span, vars)
}
fn visit_return_single(node: Node<'_>) -> ReturnSingle {
debug_assert_eq!(node.as_rule(), Rule::return_single);
let span = node.span();
let mut children = node.into_children();
children.skip_expected(Rule::RETURN);
let selector = visit_return_single_selector(children.consume_expected(Rule::return_single_selector));
let vars = visit_vars(children.consume_any());
debug_assert_eq!(children.try_consume_any(), None);
ReturnSingle::new(span, selector, vars)
}
fn visit_return_single_selector(node: Node<'_>) -> SingleSelector {
debug_assert_eq!(node.as_rule(), Rule::return_single_selector);
let child = node.into_child();
match child.as_rule() {
Rule::FIRST => SingleSelector::First,
Rule::LAST => SingleSelector::Last,
_ => unreachable!("{}", TypeQLError::IllegalGrammar { input: child.as_str().to_owned() }),
}
}
pub(super) fn visit_return_reduce(node: Node<'_>) -> ReturnReduction {
debug_assert_eq!(node.as_rule(), Rule::return_reduce);
let mut children = node.into_children();
children.skip_expected(Rule::RETURN);
let return_reduce_reduction =
visit_return_reduce_reduction(children.consume_expected(Rule::return_reduce_reduction));
debug_assert!(children.try_consume_any().is_none());
return_reduce_reduction
}
fn visit_return_reduce_reduction(node: Node<'_>) -> ReturnReduction {
debug_assert_eq!(node.as_rule(), Rule::return_reduce_reduction);
let span = node.span();
let mut children = node.into_children();
let reduction = match children.peek_rule().unwrap() {
Rule::CHECK => ReturnReduction::Check(Check::new(children.consume_expected(Rule::CHECK).span())),
Rule::reducer => ReturnReduction::Value(children.by_ref().map(visit_reducer).collect(), span),
_ => unreachable!("{}", TypeQLError::IllegalGrammar { input: children.as_str().to_owned() }),
};
debug_assert!(children.try_consume_any().is_none());
reduction
}
fn visit_function_signature(node: Node<'_>) -> Signature {
debug_assert_eq!(node.as_rule(), Rule::function_signature);
let span = node.span();
let mut children = node.into_children();
let sigil = visit_identifier(children.consume_expected(Rule::identifier));
let args = visit_pipeline_arguments(children.consume_expected(Rule::pipeline_arguments));
let return_types = visit_function_output(children.consume_expected(Rule::function_output));
debug_assert_eq!(children.try_consume_any(), None);
Signature::new(span, sigil, args, return_types)
}
fn visit_function_output(node: Node<'_>) -> Output {
debug_assert_eq!(node.as_rule(), Rule::function_output);
let child = node.into_child();
match child.as_rule() {
Rule::function_output_stream => Output::Stream(visit_function_output_stream(child)),
Rule::function_output_single => Output::Single(visit_function_output_single(child)),
_ => unreachable!("{}", TypeQLError::IllegalGrammar { input: child.as_str().to_owned() }),
}
}
fn visit_function_output_stream(node: Node<'_>) -> Stream {
debug_assert_eq!(node.as_rule(), Rule::function_output_stream);
let span = node.span();
Stream::new(span, node.into_children().map(visit_named_type_any).collect())
}
fn visit_function_output_single(node: Node<'_>) -> Single {
debug_assert_eq!(node.as_rule(), Rule::function_output_single);
let span = node.span();
Single::new(span, node.into_children().map(visit_named_type_any).collect())
}
pub(crate) fn visit_pipeline_arguments(node: Node<'_>) -> Vec<Argument> {
debug_assert_eq!(node.as_rule(), Rule::pipeline_arguments);
node.into_children().map(visit_pipeline_argument).collect()
}
fn visit_pipeline_argument(node: Node<'_>) -> Argument {
debug_assert_eq!(node.as_rule(), Rule::pipeline_argument);
let span = node.span();
let mut children = node.into_children();
let var = visit_var(children.consume_expected(Rule::var));
let type_ = visit_named_type_any(children.consume_expected(Rule::named_type_any));
debug_assert_eq!(children.try_consume_any(), None);
Argument::new(span, var, type_)
}