use indexmap::IndexSet;
use leo_ast::{AstVisitor, AsyncExpression, CallExpression, ErrExpression, Function, ProgramVisitor, Variant};
use leo_span::Symbol;
#[derive(Debug)]
pub struct AnalysisVisitor {
current_variant: Option<Variant>,
in_constructor: bool,
pub functions_to_inline: IndexSet<Vec<Symbol>>,
}
impl AnalysisVisitor {
pub fn new() -> Self {
Self { current_variant: None, in_constructor: false, functions_to_inline: IndexSet::new() }
}
}
impl AstVisitor for AnalysisVisitor {
type AdditionalInput = ();
type Output = ();
fn visit_call(&mut self, input: &CallExpression, _additional: &Self::AdditionalInput) -> Self::Output {
let callee_path = input.function.segments();
if matches!(self.current_variant, Some(Variant::Finalize | Variant::FinalFn | Variant::Fn))
|| self.in_constructor
{
self.functions_to_inline.insert(callee_path.clone());
}
input.const_arguments.iter().for_each(|expr| {
self.visit_expression(expr, &Default::default());
});
input.arguments.iter().for_each(|expr| {
self.visit_expression(expr, &Default::default());
});
}
fn visit_async(&mut self, _input: &AsyncExpression, _additional: &Self::AdditionalInput) -> Self::Output {
panic!("Async blocks should no longer exist at this point in compilation");
}
fn visit_err(&mut self, _input: &ErrExpression, _additional: &Self::AdditionalInput) -> Self::Output {
}
}
impl ProgramVisitor for AnalysisVisitor {
fn visit_function(&mut self, input: &Function) {
let prev_variant = self.current_variant;
self.current_variant = Some(input.variant);
input.const_parameters.iter().for_each(|input| self.visit_type(&input.type_));
input.input.iter().for_each(|input| self.visit_type(&input.type_));
input.output.iter().for_each(|output| self.visit_type(&output.type_));
self.visit_type(&input.output_type);
self.visit_block(&input.block);
self.current_variant = prev_variant;
}
fn visit_constructor(&mut self, input: &leo_ast::Constructor) {
let was_in_constructor = self.in_constructor;
self.in_constructor = true;
self.visit_block(&input.block);
self.in_constructor = was_in_constructor;
}
}