leo-ast 4.0.2

// Copyright (C) 2019-2026 Provable Inc.
// This file is part of the Leo library.

// The Leo library is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// The Leo library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.

//! This module contains Visitor trait implementations for the AST.
//! It implements default methods for each node to be made
//! given the type of node its visiting.

use crate::*;

// TODO: The Visitor and Reconstructor patterns need a redesign so that the default implementation can easily be invoked though its implemented in an overriding trait.
// Here is a pattern that seems to work
// trait ProgramVisitor {
//     // The trait method that can be overridden
//     fn visit_program_scope(&mut self);
//
//     // Private helper function containing the default implementation
//     fn default_visit_program_scope(&mut self) {
//         println!("Do default stuff");
//     }
// }
//
// struct YourStruct;
//
// impl ProgramVisitor for YourStruct {
//     fn visit_program_scope(&mut self) {
//         println!("Do custom stuff.");
//         // Call the default implementation
//         self.default_visit_program_scope();
//     }
// }

/// A Visitor trait for types in the AST.
pub trait AstVisitor {
    /* Types */
    fn visit_type(&mut self, input: &Type) {
        match input {
            Type::Array(array_type) => self.visit_array_type(array_type),
            Type::Composite(composite_type) => self.visit_composite_type(composite_type),
            Type::Future(future_type) => self.visit_future_type(future_type),
            Type::Mapping(mapping_type) => self.visit_mapping_type(mapping_type),
            Type::Optional(optional_type) => self.visit_optional_type(optional_type),
            Type::Tuple(tuple_type) => self.visit_tuple_type(tuple_type),
            Type::Vector(array_type) => self.visit_vector_type(array_type),
            Type::Address
            | Type::Boolean
            | Type::Field
            | Type::Group
            | Type::Identifier
            | Type::DynRecord
            | Type::Ident(_)
            | Type::Integer(_)
            | Type::Scalar
            | Type::Signature
            | Type::String
            | Type::Numeric
            | Type::Unit
            | Type::Err => {}
        }
    }

    fn visit_array_type(&mut self, input: &ArrayType) {
        self.visit_type(&input.element_type);
        self.visit_expression(&input.length, &Default::default());
    }

    fn visit_composite_type(&mut self, input: &CompositeType) {
        input.const_arguments.iter().for_each(|expr| {
            self.visit_expression(expr, &Default::default());
        });
    }

    fn visit_future_type(&mut self, input: &FutureType) {
        input.inputs.iter().for_each(|input| self.visit_type(input));
    }

    fn visit_mapping_type(&mut self, input: &MappingType) {
        self.visit_type(&input.key);
        self.visit_type(&input.value);
    }

    fn visit_optional_type(&mut self, input: &OptionalType) {
        self.visit_type(&input.inner);
    }

    fn visit_tuple_type(&mut self, input: &TupleType) {
        input.elements().iter().for_each(|input| self.visit_type(input));
    }

    fn visit_vector_type(&mut self, input: &VectorType) {
        self.visit_type(&input.element_type);
    }

    /* Expressions */
    type AdditionalInput: Default;
    type Output: Default;

    fn visit_expression(&mut self, input: &Expression, additional: &Self::AdditionalInput) -> Self::Output {
        match input {
            Expression::Array(array) => self.visit_array(array, additional),
            Expression::ArrayAccess(access) => self.visit_array_access(access, additional),
            Expression::Async(async_) => self.visit_async(async_, additional),
            Expression::Binary(binary) => self.visit_binary(binary, additional),
            Expression::Call(call) => self.visit_call(call, additional),
            Expression::DynamicCall(dc) => self.visit_dynamic_call(dc, additional),
            Expression::Cast(cast) => self.visit_cast(cast, additional),
            Expression::Composite(composite_) => self.visit_composite_init(composite_, additional),
            Expression::Err(err) => self.visit_err(err, additional),
            Expression::Path(path) => self.visit_path(path, additional),
            Expression::Literal(literal) => self.visit_literal(literal, additional),
            Expression::MemberAccess(access) => self.visit_member_access(access, additional),
            Expression::Repeat(repeat) => self.visit_repeat(repeat, additional),
            Expression::Ternary(ternary) => self.visit_ternary(ternary, additional),
            Expression::Tuple(tuple) => self.visit_tuple(tuple, additional),
            Expression::TupleAccess(access) => self.visit_tuple_access(access, additional),
            Expression::Unary(unary) => self.visit_unary(unary, additional),
            Expression::Unit(unit) => self.visit_unit(unit, additional),
            Expression::Intrinsic(intr) => self.visit_intrinsic(intr, additional),
        }
    }

    fn visit_array_access(&mut self, input: &ArrayAccess, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.array, &Default::default());
        self.visit_expression(&input.index, &Default::default());
        Default::default()
    }

    fn visit_member_access(&mut self, input: &MemberAccess, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.inner, &Default::default());
        Default::default()
    }

    fn visit_tuple_access(&mut self, input: &TupleAccess, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.tuple, &Default::default());
        Default::default()
    }

    fn visit_array(&mut self, input: &ArrayExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        input.elements.iter().for_each(|expr| {
            self.visit_expression(expr, &Default::default());
        });
        Default::default()
    }

    fn visit_async(&mut self, input: &AsyncExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_block(&input.block);
        Default::default()
    }

    fn visit_binary(&mut self, input: &BinaryExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.left, &Default::default());
        self.visit_expression(&input.right, &Default::default());
        Default::default()
    }

    fn visit_call(&mut self, input: &CallExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        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());
        });
        Default::default()
    }

    fn visit_dynamic_call(
        &mut self,
        input: &DynamicCallExpression,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        self.visit_type(&input.interface);
        self.visit_expression(&input.target_program, &Default::default());
        if let Some(ref network) = input.network {
            self.visit_expression(network, &Default::default());
        }
        input.arguments.iter().for_each(|expr| {
            self.visit_expression(expr, &Default::default());
        });
        Default::default()
    }

    fn visit_intrinsic(&mut self, input: &IntrinsicExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        for (ty, _) in &input.type_parameters {
            self.visit_type(ty);
        }
        // `input_types` and `return_types` are derived from `type_parameters` at parse time, but
        // are reconstructed independently during path resolution and must be visited separately.
        for (_, ty, _) in &input.input_types {
            self.visit_type(ty);
        }
        for (_, ty, _) in &input.return_types {
            self.visit_type(ty);
        }
        input.arguments.iter().for_each(|arg| {
            self.visit_expression(arg, &Default::default());
        });
        Default::default()
    }

    fn visit_cast(&mut self, input: &CastExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.expression, &Default::default());
        Default::default()
    }

    fn visit_composite_init(
        &mut self,
        input: &CompositeExpression,
        _additional: &Self::AdditionalInput,
    ) -> Self::Output {
        input.const_arguments.iter().for_each(|expr| {
            self.visit_expression(expr, &Default::default());
        });
        for CompositeFieldInitializer { expression, .. } in input.members.iter() {
            if let Some(expression) = expression {
                self.visit_expression(expression, &Default::default());
            }
        }
        Default::default()
    }

    fn visit_err(&mut self, _input: &ErrExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        panic!("`ErrExpression`s should not be in the AST at this phase of compilation.")
    }

    fn visit_path(&mut self, _input: &Path, _additional: &Self::AdditionalInput) -> Self::Output {
        Default::default()
    }

    fn visit_literal(&mut self, _input: &Literal, _additional: &Self::AdditionalInput) -> Self::Output {
        Default::default()
    }

    fn visit_repeat(&mut self, input: &RepeatExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.expr, &Default::default());
        self.visit_expression(&input.count, &Default::default());
        Default::default()
    }

    fn visit_ternary(&mut self, input: &TernaryExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.condition, &Default::default());
        self.visit_expression(&input.if_true, &Default::default());
        self.visit_expression(&input.if_false, &Default::default());
        Default::default()
    }

    fn visit_tuple(&mut self, input: &TupleExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        input.elements.iter().for_each(|expr| {
            self.visit_expression(expr, &Default::default());
        });
        Default::default()
    }

    fn visit_unary(&mut self, input: &UnaryExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        self.visit_expression(&input.receiver, &Default::default());
        Default::default()
    }

    fn visit_unit(&mut self, _input: &UnitExpression, _additional: &Self::AdditionalInput) -> Self::Output {
        Default::default()
    }

    /* Statements */
    fn visit_statement(&mut self, input: &Statement) {
        match input {
            Statement::Assert(stmt) => self.visit_assert(stmt),
            Statement::Assign(stmt) => self.visit_assign(stmt),
            Statement::Block(stmt) => self.visit_block(stmt),
            Statement::Conditional(stmt) => self.visit_conditional(stmt),
            Statement::Const(stmt) => self.visit_const(stmt),
            Statement::Definition(stmt) => self.visit_definition(stmt),
            Statement::Expression(stmt) => self.visit_expression_statement(stmt),
            Statement::Iteration(stmt) => self.visit_iteration(stmt),
            Statement::Return(stmt) => self.visit_return(stmt),
        }
    }

    fn visit_assert(&mut self, input: &AssertStatement) {
        match &input.variant {
            AssertVariant::Assert(expr) => self.visit_expression(expr, &Default::default()),
            AssertVariant::AssertEq(left, right) | AssertVariant::AssertNeq(left, right) => {
                self.visit_expression(left, &Default::default());
                self.visit_expression(right, &Default::default())
            }
        };
    }

    fn visit_assign(&mut self, input: &AssignStatement) {
        self.visit_expression(&input.place, &Default::default());
        self.visit_expression(&input.value, &Default::default());
    }

    fn visit_block(&mut self, input: &Block) {
        input.statements.iter().for_each(|stmt| self.visit_statement(stmt));
    }

    fn visit_conditional(&mut self, input: &ConditionalStatement) {
        self.visit_expression(&input.condition, &Default::default());
        self.visit_block(&input.then);
        if let Some(stmt) = input.otherwise.as_ref() {
            self.visit_statement(stmt);
        }
    }

    fn visit_const(&mut self, input: &ConstDeclaration) {
        self.visit_type(&input.type_);
        self.visit_expression(&input.value, &Default::default());
    }

    fn visit_definition(&mut self, input: &DefinitionStatement) {
        if let Some(ty) = input.type_.as_ref() {
            self.visit_type(ty)
        }
        self.visit_expression(&input.value, &Default::default());
    }

    fn visit_expression_statement(&mut self, input: &ExpressionStatement) {
        self.visit_expression(&input.expression, &Default::default());
    }

    fn visit_iteration(&mut self, input: &IterationStatement) {
        if let Some(ty) = input.type_.as_ref() {
            self.visit_type(ty)
        }
        self.visit_expression(&input.start, &Default::default());
        self.visit_expression(&input.stop, &Default::default());
        self.visit_block(&input.block);
    }

    fn visit_return(&mut self, input: &ReturnStatement) {
        self.visit_expression(&input.expression, &Default::default());
    }
}

/// A Visitor trait for the program represented by the AST.
pub trait ProgramVisitor: AstVisitor {
    fn visit_program(&mut self, input: &Program) {
        input.program_scopes.values().for_each(|scope| self.visit_program_scope(scope));
        input.modules.values().for_each(|module| self.visit_module(module));
        input.stubs.values().for_each(|stub| self.visit_stub(stub));
    }

    fn visit_aleo_program(&mut self, _input: &AleoProgram) {}

    fn visit_library(&mut self, input: &Library) {
        input.interfaces.iter().for_each(|(_, i)| self.visit_interface(i));
        input.consts.iter().for_each(|(_, c)| self.visit_const(c));
        input.structs.iter().for_each(|(_, s)| self.visit_composite(s));
        input.functions.iter().for_each(|(_, f)| self.visit_function(f));
        input.modules.values().for_each(|m| self.visit_module(m));
    }

    fn visit_stub(&mut self, input: &Stub) {
        match input {
            Stub::FromLeo { program, .. } => self.visit_program(program),
            Stub::FromAleo { program, .. } => self.visit_aleo_program(program),
            Stub::FromLibrary { library, .. } => self.visit_library(library),
        }
    }

    fn visit_program_scope(&mut self, input: &ProgramScope) {
        input.consts.iter().for_each(|(_, c)| self.visit_const(c));
        input.composites.iter().for_each(|(_, c)| self.visit_composite(c));
        input.interfaces.iter().for_each(|(_, c)| self.visit_interface(c));
        input.mappings.iter().for_each(|(_, c)| self.visit_mapping(c));
        input.storage_variables.iter().for_each(|(_, c)| self.visit_storage_variable(c));
        input.functions.iter().for_each(|(_, c)| self.visit_function(c));
        if let Some(c) = input.constructor.as_ref() {
            self.visit_constructor(c);
        }
    }

    fn visit_module(&mut self, input: &Module) {
        input.consts.iter().for_each(|(_, c)| self.visit_const(c));
        input.composites.iter().for_each(|(_, c)| self.visit_composite(c));
        input.interfaces.iter().for_each(|(_, c)| self.visit_interface(c));
        input.functions.iter().for_each(|(_, c)| self.visit_function(c));
    }

    fn visit_composite(&mut self, input: &Composite) {
        input.const_parameters.iter().for_each(|input| self.visit_type(&input.type_));
        input.members.iter().for_each(|member| self.visit_type(&member.type_));
    }

    fn visit_mapping(&mut self, input: &Mapping) {
        self.visit_type(&input.key_type);
        self.visit_type(&input.value_type);
    }

    fn visit_storage_variable(&mut self, input: &StorageVariable) {
        self.visit_type(&input.type_);
    }

    fn visit_function(&mut self, input: &Function) {
        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);
    }

    fn visit_interface(&mut self, input: &Interface) {
        input.functions.iter().for_each(|(_, f)| self.visit_function_prototype(f));
        input.records.iter().for_each(|(_, r)| self.visit_record_prototype(r));
        input.mappings.iter().for_each(|m| self.visit_mapping_prototype(m));
        input.storages.iter().for_each(|s| self.visit_storage_variable_prototype(s));
    }

    fn visit_function_prototype(&mut self, input: &FunctionPrototype) {
        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);
    }

    fn visit_record_prototype(&mut self, input: &RecordPrototype) {
        input.members.iter().for_each(|member| self.visit_type(&member.type_));
    }

    fn visit_mapping_prototype(&mut self, input: &MappingPrototype) {
        self.visit_type(&input.key_type);
        self.visit_type(&input.value_type);
    }

    fn visit_storage_variable_prototype(&mut self, input: &StorageVariablePrototype) {
        self.visit_type(&input.type_);
    }

    fn visit_constructor(&mut self, input: &Constructor) {
        self.visit_block(&input.block);
    }

    fn visit_function_stub(&mut self, _input: &FunctionStub) {}

    fn visit_composite_stub(&mut self, _input: &Composite) {}
}