SeaLang/src/ast2/visitor.rs

use super::*;

pub trait AstExt {
    fn get_node_children(&self, node: Index) -> Vec<Index>;
    fn get_node_tag_and_data(&self, node: Index) -> (Tag, Data);

    fn get_node_data_for_tag(&self, idx: Index, tag: Tag) -> Option<Data> {
        let (t, data) = self.get_node_tag_and_data(idx);
        if t == tag {
            Some(data)
        } else {
            None
        }
    }

    fn expect_node_data_for_tag(&self, idx: Index, tag: Tag) -> Data {
        self.get_node_data_for_tag(idx, tag)
            .expect(&format!("node {idx} is not a {tag:?}"))
    }
}

impl AstExt for &Ast {
    fn get_node_children(&self, node: Index) -> Vec<Index> {
        Ast::get_node_children(self, node)
    }

    fn get_node_tag_and_data(&self, node: Index) -> (Tag, Data) {
        (self.tags[node], self.datas[node])
    }
}

impl AstExt for &mut Ast {
    fn get_node_children(&self, node: Index) -> Vec<Index> {
        Ast::get_node_children(self, node)
    }

    fn get_node_tag_and_data(&self, node: Index) -> (Tag, Data) {
        (self.tags[node], self.datas[node])
    }
}

pub struct AstVisitor<AstT> {
    pub(super) ast: AstT,
    pub(super) scopes: Vec<Index>,
    pub(super) nodes: Vec<A>,
    pub(super) rev: bool,
}

impl<AstT> AstVisitor<AstT>
where
    AstT: AstExt,
{
    pub fn visit<
        F: FnMut(&mut AstT, &[Index], Index, Tag, Data),
        G: FnMut(&mut AstT, &[Index], Index, Tag, Data),
    >(
        &mut self,
        mut pre: F,
        mut post: G,
    ) {
        while let Some(node) = self.nodes.pop() {
            match node {
                A::PushChildren(i) => {
                    self.nodes.push(A::PopSelf(i));
                    let children_iter = self
                        .ast
                        .get_node_children(i)
                        .into_iter()
                        .map(|i| A::PushChildren(i));
                    // inverse because we are popping from the end
                    if !self.rev {
                        self.nodes.extend(children_iter.rev())
                    } else {
                        self.nodes.extend(children_iter)
                    };

                    let (tag, data) = self.ast.get_node_tag_and_data(i);

                    let _ = pre(&mut self.ast, &self.scopes, i, tag, data);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.push(i);
                        }
                        _ => {}
                    }
                }
                A::PopSelf(i) => {
                    // already popped.
                    let (tag, data) = self.ast.get_node_tag_and_data(i);

                    let _ = post(&mut self.ast, &self.scopes, i, tag, data);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.pop();
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    pub fn visit_pre<F: FnMut(&mut AstT, &[Index], Index, Tag, Data)>(&mut self, mut cb: F) {
        while let Some(node) = self.nodes.pop() {
            match node {
                A::PushChildren(i) => {
                    self.nodes.push(A::PopSelf(i));
                    let children_iter = self
                        .ast
                        .get_node_children(i)
                        .into_iter()
                        .map(|i| A::PushChildren(i));
                    // inverse because we are popping from the end
                    if !self.rev {
                        self.nodes.extend(children_iter.rev())
                    } else {
                        self.nodes.extend(children_iter)
                    };

                    let (tag, data) = self.ast.get_node_tag_and_data(i);

                    let _ = cb(&mut self.ast, &self.scopes, i, tag, data);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.push(i);
                        }
                        _ => {}
                    }
                }
                A::PopSelf(i) => {
                    // already popped.
                    let (tag, _data) = self.ast.get_node_tag_and_data(i);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.pop();
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    pub fn visit_post<F: FnMut(&mut AstT, &[Index], Index, Tag, Data)>(&mut self, mut cb: F) {
        while let Some(node) = self.nodes.pop() {
            match node {
                A::PushChildren(i) => {
                    self.nodes.push(A::PopSelf(i));
                    let children_iter = self
                        .ast
                        .get_node_children(i)
                        .into_iter()
                        .map(|i| A::PushChildren(i));
                    if self.rev {
                        self.nodes.extend(children_iter.rev())
                    } else {
                        self.nodes.extend(children_iter)
                    };

                    let (tag, _data) = self.ast.get_node_tag_and_data(i);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.push(i);
                        }
                        _ => {}
                    }
                }
                A::PopSelf(i) => {
                    // already popped.
                    let (tag, data) = self.ast.get_node_tag_and_data(i);

                    let _ = cb(&mut self.ast, &self.scopes, i, tag, data);

                    match tag {
                        Tag::File
                        | Tag::FunctionDecl
                        | Tag::GlobalDecl
                        | Tag::Block
                        | Tag::BlockTrailingExpr => {
                            self.scopes.pop();
                        }
                        _ => {}
                    }
                }
            }
        }
    }
}

pub enum A {
    PushChildren(Index),
    PopSelf(Index),
}

impl Ast {
    pub fn visitor_mut(&mut self) -> AstVisitor<&mut Self> {
        AstVisitor {
            scopes: vec![],
            nodes: self
                .get_root_file_indices()
                .map(|i| visitor::A::PushChildren(i))
                .collect(),
            ast: self,
            rev: false,
        }
    }
    pub fn visitor_rev_mut(&mut self) -> AstVisitor<&mut Self> {
        AstVisitor {
            scopes: vec![],
            nodes: self
                .get_root_file_indices()
                .map(|i| visitor::A::PushChildren(i))
                .collect(),
            ast: self,
            rev: true,
        }
    }

    pub fn visitor(&self) -> AstVisitor<&Self> {
        AstVisitor {
            scopes: vec![],
            nodes: self
                .get_root_file_indices()
                .map(|i| visitor::A::PushChildren(i))
                .collect(),
            ast: self,
            rev: false,
        }
    }
}

pub trait AstVisitorTrait<Ast: AstExt> {
    type Error;
    type Value;
    const UNIMPL: Self::Error;

    fn visit_function_decl(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_parameter(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_parameter_list(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_block(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_block_trailing_expr(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_block_maybe_trailing(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        match ast.get_node_tag_and_data(idx).0 {
            Tag::BlockTrailingExpr => self.visit_block_trailing_expr(ast, idx),
            Tag::Block => self.visit_block(ast, idx),
            _ => unreachable!(),
        }
    }
    fn visit_function_proto(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_call_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_add_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_sub_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_mul_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_div_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_rem_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_eq_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_neq_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_lt_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_gt_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_le_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_ge_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_shl_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_shr_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_bitor_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_bitxor_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_bitand_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_or_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_and_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_not_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_negate_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_deref_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_address_of_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_explicit_cast_expr(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_assign(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_subscript_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_if_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_if_else_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_argument(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }

    fn visit_any_argument(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        match ast.get_node_tag_and_data(idx).0 {
            Tag::Argument => self.visit_argument(ast, idx),
            Tag::NamedArgument => self.visit_named_argument(ast, idx),
            _ => unreachable!(),
        }
    }

    fn visit_named_argument(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }

    fn visit_argument_list(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }

    fn visit_constant(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }

    fn visit_return(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_return_expr(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_global_decl(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_var_decl(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_mut_var_decl(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_var_assign_decl(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_mut_var_assign_decl(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_decl_ref(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        _ = (ast, idx);
        Err(Self::UNIMPL)
    }
    fn visit_value_to_place_conversion(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        let idx = ast
            .get_node_data_for_tag(idx, Tag::ValueToPlaceConversion)
            .unwrap()
            .as_index();
        self.visit_any(ast, idx)
    }
    fn visit_place_to_value_conversion(
        &mut self,
        ast: Ast,
        idx: Index,
    ) -> Result<Self::Value, Self::Error> {
        let idx = ast
            .get_node_data_for_tag(idx, Tag::PlaceToValueConversion)
            .unwrap()
            .as_index();
        self.visit_any(ast, idx)
    }

    fn visit_any(&mut self, ast: Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        match ast.get_node_tag_and_data(idx).0 {
            Tag::FunctionProto => self.visit_function_proto(ast, idx),
            Tag::FunctionDecl => self.visit_function_decl(ast, idx),
            Tag::ParameterList => self.visit_parameter_list(ast, idx),
            Tag::Parameter => self.visit_parameter(ast, idx),
            Tag::Block => self.visit_block(ast, idx),
            Tag::BlockTrailingExpr => self.visit_block_trailing_expr(ast, idx),
            Tag::Constant => self.visit_constant(ast, idx),
            Tag::ExprStmt => todo!(),
            Tag::ReturnStmt => self.visit_return(ast, idx),
            Tag::ReturnExprStmt => self.visit_return_expr(ast, idx),
            Tag::VarDecl => self.visit_var_decl(ast, idx),
            Tag::MutVarDecl => self.visit_mut_var_decl(ast, idx),
            Tag::VarDeclAssignment => self.visit_var_assign_decl(ast, idx),
            Tag::MutVarDeclAssignment => self.visit_mut_var_assign_decl(ast, idx),
            Tag::GlobalDecl => self.visit_global_decl(ast, idx),
            Tag::StructDecl => todo!(),
            Tag::FieldDecl => todo!(),
            Tag::DeclRef => self.visit_decl_ref(ast, idx),
            Tag::DeclRefUnresolved => todo!(),
            Tag::InternedType => todo!(),
            Tag::TypeDeclRef => todo!(),
            Tag::TypeDeclRefUnresolved => todo!(),
            Tag::PointerType => todo!(),
            Tag::ArrayType => todo!(),
            Tag::CallExpr => self.visit_call_expr(ast, idx),
            Tag::FieldAccess => todo!(),
            Tag::ArgumentList => self.visit_argument_list(ast, idx),
            Tag::Argument => self.visit_argument(ast, idx),
            Tag::NamedArgument => self.visit_named_argument(ast, idx),
            Tag::ExplicitCast => self.visit_explicit_cast_expr(ast, idx),
            Tag::Deref => self.visit_deref_expr(ast, idx),
            Tag::AddressOf => self.visit_add_expr(ast, idx),
            Tag::Not => self.visit_not_expr(ast, idx),
            Tag::Negate => self.visit_negate_expr(ast, idx),
            Tag::Or => self.visit_or_expr(ast, idx),
            Tag::And => self.visit_and_expr(ast, idx),
            Tag::BitOr => self.visit_bitor_expr(ast, idx),
            Tag::BitXOr => self.visit_bitxor_expr(ast, idx),
            Tag::BitAnd => self.visit_bitand_expr(ast, idx),
            Tag::Eq => self.visit_eq_expr(ast, idx),
            Tag::NEq => self.visit_neq_expr(ast, idx),
            Tag::Lt => self.visit_lt_expr(ast, idx),
            Tag::Gt => self.visit_gt_expr(ast, idx),
            Tag::Le => self.visit_le_expr(ast, idx),
            Tag::Ge => self.visit_ge_expr(ast, idx),
            Tag::Shl => self.visit_shl_expr(ast, idx),
            Tag::Shr => self.visit_shr_expr(ast, idx),
            Tag::Add => self.visit_add_expr(ast, idx),
            Tag::Sub => self.visit_sub_expr(ast, idx),
            Tag::Mul => self.visit_mul_expr(ast, idx),
            Tag::Div => self.visit_div_expr(ast, idx),
            Tag::Rem => self.visit_rem_expr(ast, idx),
            Tag::Assign => self.visit_assign(ast, idx),
            Tag::SubscriptExpr => self.visit_subscript_expr(ast, idx),
            Tag::IfExpr => self.visit_if_expr(ast, idx),
            Tag::IfElseExpr => self.visit_if_else_expr(ast, idx),
            Tag::PlaceToValueConversion => self.visit_place_to_value_conversion(ast, idx),
            Tag::ValueToPlaceConversion => self.visit_value_to_place_conversion(ast, idx),
            Tag::Error => todo!(),
            Tag::Undefined => todo!(),
            Tag::Root => todo!(),
            Tag::File => todo!(),
        }
    }
}