joule_mir/

visit.rs

//! Visitor pattern for traversing MIR
//!
//! This module provides visitor traits for traversing and transforming MIR.
//! Used by optimization passes, analysis passes, and code generation.

use crate::{
    BasicBlock, BasicBlockId, FunctionMIR, Local, MirContext, Operand, Place, Rvalue, Statement,
    Terminator,
};

/// Visitor trait for MIR
///
/// Implement this trait to traverse MIR structures.
/// All methods have default implementations that do nothing.
pub trait MirVisitor {
    /// Visit a MIR context
    fn visit_context(&mut self, ctx: &MirContext) {
        for (_, func) in &ctx.functions {
            self.visit_function(func);
        }
    }

    /// Visit a function
    fn visit_function(&mut self, func: &FunctionMIR) {
        for (idx, block) in func.basic_blocks.iter().enumerate() {
            self.visit_basic_block(BasicBlockId::new(idx as u32), block);
        }
    }

    /// Visit a basic block
    fn visit_basic_block(&mut self, _id: BasicBlockId, block: &BasicBlock) {
        for stmt in &block.statements {
            self.visit_statement(stmt);
        }
        self.visit_terminator(&block.terminator);
    }

    /// Visit a statement
    fn visit_statement(&mut self, stmt: &Statement) {
        match stmt {
            Statement::Assign { place, rvalue, .. } => {
                self.visit_place(place);
                self.visit_rvalue(rvalue);
            }
            Statement::StorageLive { local, .. } | Statement::StorageDead { local, .. } => {
                self.visit_local(*local);
            }
            Statement::Nop => {}
        }
    }

    /// Visit a terminator
    fn visit_terminator(&mut self, term: &Terminator) {
        match term {
            Terminator::Return { .. } => {}
            Terminator::Goto { target, .. } => {
                self.visit_target(*target);
            }
            Terminator::SwitchInt {
                discriminant,
                targets,
                ..
            } => {
                self.visit_operand(discriminant);
                for (_, target) in &targets.branches {
                    self.visit_target(*target);
                }
                self.visit_target(targets.otherwise);
            }
            Terminator::Call {
                func,
                args,
                destination,
                target,
                ..
            } => {
                self.visit_operand(func);
                for arg in args {
                    self.visit_operand(arg);
                }
                self.visit_place(destination);
                self.visit_target(*target);
            }
            Terminator::Abort { .. }
            | Terminator::Unreachable { .. }
            | Terminator::Cancel { .. } => {}

            // Concurrency terminators
            Terminator::Spawn {
                func,
                args,
                destination,
                target,
                ..
            } => {
                self.visit_operand(func);
                for arg in args {
                    self.visit_operand(arg);
                }
                self.visit_place(destination);
                self.visit_target(*target);
            }
            Terminator::TaskAwait {
                task,
                destination,
                target,
                ..
            } => {
                self.visit_operand(task);
                self.visit_place(destination);
                self.visit_target(*target);
            }
            Terminator::TaskGroupEnter {
                destination,
                body,
                join_block,
                ..
            } => {
                self.visit_place(destination);
                self.visit_target(*body);
                self.visit_target(*join_block);
            }
            Terminator::TaskGroupExit { group, target, .. } => {
                self.visit_operand(group);
                self.visit_target(*target);
            }
            Terminator::ChannelRecv {
                channel,
                destination,
                target,
                closed_target,
                ..
            } => {
                self.visit_operand(channel);
                self.visit_place(destination);
                self.visit_target(*target);
                self.visit_target(*closed_target);
            }
            Terminator::ChannelSend {
                channel,
                value,
                target,
                closed_target,
                ..
            } => {
                self.visit_operand(channel);
                self.visit_operand(value);
                self.visit_target(*target);
                self.visit_target(*closed_target);
            }
            Terminator::Select {
                arms,
                default,
                destination,
                selected_arm,
                ..
            } => {
                for arm in arms {
                    match &arm.operation {
                        crate::ChannelOp::Recv { channel } => {
                            self.visit_operand(channel);
                        }
                        crate::ChannelOp::Send { channel, value } => {
                            self.visit_operand(channel);
                            self.visit_operand(value);
                        }
                        crate::ChannelOp::Timeout { .. } => {}
                    }
                    self.visit_target(arm.target);
                }
                if let Some(default) = default {
                    self.visit_target(*default);
                }
                self.visit_place(destination);
                self.visit_place(selected_arm);
            }
        }
    }

    /// Visit a place
    fn visit_place(&mut self, place: &Place) {
        self.visit_local(place.local);
        for elem in &place.projection {
            if let crate::PlaceElem::Index(local) = elem {
                self.visit_local(*local);
            }
        }
    }

    /// Visit an rvalue
    fn visit_rvalue(&mut self, rvalue: &Rvalue) {
        match rvalue {
            Rvalue::Use(operand) => {
                self.visit_operand(operand);
            }
            Rvalue::BinaryOp { left, right, .. } => {
                self.visit_operand(left);
                self.visit_operand(right);
            }
            Rvalue::UnaryOp { operand, .. } => {
                self.visit_operand(operand);
            }
            Rvalue::Ref { place, .. } => {
                self.visit_place(place);
            }
            Rvalue::Aggregate { operands, .. } => {
                for operand in operands {
                    self.visit_operand(operand);
                }
            }
            Rvalue::Cast { operand, .. } => {
                self.visit_operand(operand);
            }
            Rvalue::Discriminant { place } | Rvalue::Len { place } => {
                self.visit_place(place);
            }
            Rvalue::SimdBinaryOp { left, right, .. } => {
                self.visit_operand(left);
                self.visit_operand(right);
            }
            Rvalue::SimdLoad { source, .. } => {
                self.visit_place(source);
            }
            Rvalue::SimdStore { value, dest, .. } => {
                self.visit_operand(value);
                self.visit_place(dest);
            }
            Rvalue::SimdSplat { value, .. } => {
                self.visit_operand(value);
            }

            // Concurrency rvalues
            Rvalue::ChannelCreate { .. } => {}
            Rvalue::ChannelTryRecv { channel } => {
                self.visit_operand(channel);
            }
            Rvalue::ChannelTrySend { channel, value } => {
                self.visit_operand(channel);
                self.visit_operand(value);
            }
            Rvalue::ChannelSender { channel }
            | Rvalue::ChannelReceiver { channel }
            | Rvalue::ChannelClose { channel } => {
                self.visit_operand(channel);
            }
            Rvalue::IsCancelled | Rvalue::CurrentTask => {}
            Rvalue::Try(operand) => {
                self.visit_operand(operand);
            }
            Rvalue::EnumVariant { fields, .. } => {
                for op in fields {
                    self.visit_operand(op);
                }
            }
        }
    }

    /// Visit an operand
    fn visit_operand(&mut self, operand: &Operand) {
        match operand {
            Operand::Copy(place) | Operand::Move(place) => {
                self.visit_place(place);
            }
            Operand::Constant(_) => {}
        }
    }

    /// Visit a local
    fn visit_local(&mut self, _local: Local) {}

    /// Visit a target block
    fn visit_target(&mut self, _target: BasicBlockId) {}
}

/// Mutable visitor trait for MIR
///
/// Implement this trait to transform MIR structures.
/// All methods have default implementations that do nothing.
pub trait MirVisitorMut {
    /// Visit a MIR context
    fn visit_context_mut(&mut self, ctx: &mut MirContext) {
        for (_, func) in &mut ctx.functions {
            self.visit_function_mut(func);
        }
    }

    /// Visit a function
    fn visit_function_mut(&mut self, func: &mut FunctionMIR) {
        for idx in 0..func.basic_blocks.len() {
            let block = &mut func.basic_blocks[idx];
            self.visit_basic_block_mut(BasicBlockId::new(idx as u32), block);
        }
    }

    /// Visit a basic block
    fn visit_basic_block_mut(&mut self, _id: BasicBlockId, block: &mut BasicBlock) {
        for stmt in &mut block.statements {
            self.visit_statement_mut(stmt);
        }
        self.visit_terminator_mut(&mut block.terminator);
    }

    /// Visit a statement
    fn visit_statement_mut(&mut self, stmt: &mut Statement) {
        match stmt {
            Statement::Assign { place, rvalue, .. } => {
                self.visit_place_mut(place);
                self.visit_rvalue_mut(rvalue);
            }
            Statement::StorageLive { local, .. } | Statement::StorageDead { local, .. } => {
                self.visit_local_mut(local);
            }
            Statement::Nop => {}
        }
    }

    /// Visit a terminator
    fn visit_terminator_mut(&mut self, term: &mut Terminator) {
        match term {
            Terminator::Return { .. } => {}
            Terminator::Goto { target, .. } => {
                self.visit_target_mut(target);
            }
            Terminator::SwitchInt {
                discriminant,
                targets,
                ..
            } => {
                self.visit_operand_mut(discriminant);
                for (_, target) in &mut targets.branches {
                    self.visit_target_mut(target);
                }
                self.visit_target_mut(&mut targets.otherwise);
            }
            Terminator::Call {
                func,
                args,
                destination,
                target,
                ..
            } => {
                self.visit_operand_mut(func);
                for arg in args {
                    self.visit_operand_mut(arg);
                }
                self.visit_place_mut(destination);
                self.visit_target_mut(target);
            }
            Terminator::Abort { .. }
            | Terminator::Unreachable { .. }
            | Terminator::Cancel { .. } => {}

            // Concurrency terminators
            Terminator::Spawn {
                func,
                args,
                destination,
                target,
                ..
            } => {
                self.visit_operand_mut(func);
                for arg in args {
                    self.visit_operand_mut(arg);
                }
                self.visit_place_mut(destination);
                self.visit_target_mut(target);
            }
            Terminator::TaskAwait {
                task,
                destination,
                target,
                ..
            } => {
                self.visit_operand_mut(task);
                self.visit_place_mut(destination);
                self.visit_target_mut(target);
            }
            Terminator::TaskGroupEnter {
                destination,
                body,
                join_block,
                ..
            } => {
                self.visit_place_mut(destination);
                self.visit_target_mut(body);
                self.visit_target_mut(join_block);
            }
            Terminator::TaskGroupExit { group, target, .. } => {
                self.visit_operand_mut(group);
                self.visit_target_mut(target);
            }
            Terminator::ChannelRecv {
                channel,
                destination,
                target,
                closed_target,
                ..
            } => {
                self.visit_operand_mut(channel);
                self.visit_place_mut(destination);
                self.visit_target_mut(target);
                self.visit_target_mut(closed_target);
            }
            Terminator::ChannelSend {
                channel,
                value,
                target,
                closed_target,
                ..
            } => {
                self.visit_operand_mut(channel);
                self.visit_operand_mut(value);
                self.visit_target_mut(target);
                self.visit_target_mut(closed_target);
            }
            Terminator::Select {
                arms,
                default,
                destination,
                selected_arm,
                ..
            } => {
                for arm in arms {
                    match &mut arm.operation {
                        crate::ChannelOp::Recv { channel } => {
                            self.visit_operand_mut(channel);
                        }
                        crate::ChannelOp::Send { channel, value } => {
                            self.visit_operand_mut(channel);
                            self.visit_operand_mut(value);
                        }
                        crate::ChannelOp::Timeout { .. } => {}
                    }
                    self.visit_target_mut(&mut arm.target);
                }
                if let Some(default) = default {
                    self.visit_target_mut(default);
                }
                self.visit_place_mut(destination);
                self.visit_place_mut(selected_arm);
            }
        }
    }

    /// Visit a place
    fn visit_place_mut(&mut self, place: &mut Place) {
        self.visit_local_mut(&mut place.local);
        for elem in &mut place.projection {
            if let crate::PlaceElem::Index(local) = elem {
                self.visit_local_mut(local);
            }
        }
    }

    /// Visit an rvalue
    fn visit_rvalue_mut(&mut self, rvalue: &mut Rvalue) {
        match rvalue {
            Rvalue::Use(operand) => {
                self.visit_operand_mut(operand);
            }
            Rvalue::BinaryOp { left, right, .. } => {
                self.visit_operand_mut(left);
                self.visit_operand_mut(right);
            }
            Rvalue::UnaryOp { operand, .. } => {
                self.visit_operand_mut(operand);
            }
            Rvalue::Ref { place, .. } => {
                self.visit_place_mut(place);
            }
            Rvalue::Aggregate { operands, .. } => {
                for operand in operands {
                    self.visit_operand_mut(operand);
                }
            }
            Rvalue::Cast { operand, .. } => {
                self.visit_operand_mut(operand);
            }
            Rvalue::Discriminant { place } | Rvalue::Len { place } => {
                self.visit_place_mut(place);
            }
            Rvalue::SimdBinaryOp { left, right, .. } => {
                self.visit_operand_mut(left);
                self.visit_operand_mut(right);
            }
            Rvalue::SimdLoad { source, .. } => {
                self.visit_place_mut(source);
            }
            Rvalue::SimdStore { value, dest, .. } => {
                self.visit_operand_mut(value);
                self.visit_place_mut(dest);
            }
            Rvalue::SimdSplat { value, .. } => {
                self.visit_operand_mut(value);
            }

            // Concurrency rvalues
            Rvalue::ChannelCreate { .. } => {}
            Rvalue::ChannelTryRecv { channel } => {
                self.visit_operand_mut(channel);
            }
            Rvalue::ChannelTrySend { channel, value } => {
                self.visit_operand_mut(channel);
                self.visit_operand_mut(value);
            }
            Rvalue::ChannelSender { channel }
            | Rvalue::ChannelReceiver { channel }
            | Rvalue::ChannelClose { channel } => {
                self.visit_operand_mut(channel);
            }
            Rvalue::IsCancelled | Rvalue::CurrentTask => {}
            Rvalue::Try(operand) => {
                self.visit_operand_mut(operand);
            }
            Rvalue::EnumVariant { fields, .. } => {
                for op in fields {
                    self.visit_operand_mut(op);
                }
            }
        }
    }

    /// Visit an operand
    fn visit_operand_mut(&mut self, operand: &mut Operand) {
        match operand {
            Operand::Copy(place) | Operand::Move(place) => {
                self.visit_place_mut(place);
            }
            Operand::Constant(_) => {}
        }
    }

    /// Visit a local
    fn visit_local_mut(&mut self, _local: &mut Local) {}

    /// Visit a target block
    fn visit_target_mut(&mut self, _target: &mut BasicBlockId) {}
}

/// Example visitor that counts locals
pub struct LocalCounter {
    pub count: usize,
}

impl LocalCounter {
    pub fn new() -> Self {
        Self { count: 0 }
    }
}

impl Default for LocalCounter {
    fn default() -> Self {
        Self::new()
    }
}

impl MirVisitor for LocalCounter {
    fn visit_local(&mut self, _local: Local) {
        self.count += 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{FunctionId, LocalDecl, Symbol, Ty};
    use joule_common::Span;

    #[test]
    fn test_local_counter() {
        let mut func = FunctionMIR::new(
            FunctionId::new(0),
            Symbol::from_u32(0),
            Ty::Unit,
            Span::dummy(),
        );

        // Add some locals
        func.add_local(LocalDecl::new(None, Ty::Unit, false, Span::dummy()));
        func.add_local(LocalDecl::new(None, Ty::Unit, false, Span::dummy()));

        // Add a block with a statement
        let mut block = BasicBlock::new(Terminator::Return {
            span: Span::dummy(),
        });
        block.push_statement(Statement::Assign {
            place: Place::from_local(Local::new(1)),
            rvalue: Rvalue::Use(Operand::Copy(Place::from_local(Local::new(2)))),
            span: Span::dummy(),
        });
        func.add_block(block);

        let mut counter = LocalCounter::new();
        counter.visit_function(&func);

        // Should count: local 1 (destination) + local 2 (source) = 2
        assert_eq!(counter.count, 2);
    }

    #[test]
    fn test_visitor_traversal() {
        struct TestVisitor {
            visited_blocks: Vec<BasicBlockId>,
        }

        impl TestVisitor {
            fn new() -> Self {
                Self {
                    visited_blocks: Vec::new(),
                }
            }
        }

        impl MirVisitor for TestVisitor {
            fn visit_basic_block(&mut self, id: BasicBlockId, block: &BasicBlock) {
                self.visited_blocks.push(id);
                // Continue with default behavior
                for stmt in &block.statements {
                    self.visit_statement(stmt);
                }
                self.visit_terminator(&block.terminator);
            }
        }

        let mut func = FunctionMIR::new(
            FunctionId::new(0),
            Symbol::from_u32(0),
            Ty::Unit,
            Span::dummy(),
        );

        // Add three blocks
        func.add_block(BasicBlock::new(Terminator::Return {
            span: Span::dummy(),
        }));
        func.add_block(BasicBlock::new(Terminator::Return {
            span: Span::dummy(),
        }));
        func.add_block(BasicBlock::new(Terminator::Return {
            span: Span::dummy(),
        }));

        let mut visitor = TestVisitor::new();
        visitor.visit_function(&func);

        assert_eq!(visitor.visited_blocks.len(), 3);
        assert_eq!(visitor.visited_blocks[0], BasicBlockId::new(0));
        assert_eq!(visitor.visited_blocks[1], BasicBlockId::new(1));
        assert_eq!(visitor.visited_blocks[2], BasicBlockId::new(2));
    }
}