joule_mir/dataflow/

extract.rs

//! MIR to Dataflow Graph Extraction
//!
//! This module converts MIR (basic block form) to a dataflow graph representation.
//! The extraction process:
//!
//! 1. Creates channels for each MIR local variable
//! 2. Converts statements to Compute/Memory operators
//! 3. Converts terminators to control-flow operators (Steer, Stream, etc.)
//! 4. Detects loop patterns and generates Stream operators where beneficial
//!
//! # Example
//!
//! ```ignore
//! let mir = ...;  // MIR for a function
//! let dfg = DfgExtractor::extract(&mir);
//! ```

use super::{ChannelId, ComputeOp, DataflowGraph, DfOperator, MemoryOp, TokenType, TokenValue};
use crate::{
    BasicBlock, BasicBlockId, BinOp, FunctionMIR, Local, Operand, Place, Rvalue, Statement,
    Terminator, Ty, UnOp,
};
use std::collections::HashMap;

/// Extracts a dataflow graph from MIR
pub struct DfgExtractor {
    dfg: DataflowGraph,
    /// Map from MIR local to channel ID
    local_to_channel: HashMap<Local, ChannelId>,
    /// Map from basic block to entry channel (for control flow)
    block_entry: HashMap<BasicBlockId, ChannelId>,
    /// Current block being processed
    current_block: Option<BasicBlockId>,
    /// Loop detection: back edges
    back_edges: Vec<(BasicBlockId, BasicBlockId)>,
}

impl DfgExtractor {
    /// Create a new extractor
    pub fn new() -> Self {
        Self {
            dfg: DataflowGraph::new(),
            local_to_channel: HashMap::new(),
            block_entry: HashMap::new(),
            current_block: None,
            back_edges: Vec::new(),
        }
    }

    /// Extract a dataflow graph from a MIR function
    pub fn extract(mir: &FunctionMIR) -> DataflowGraph {
        let mut extractor = Self::new();
        extractor.dfg.name = Some(mir.name.as_str().to_string());

        // Phase 1: Create channels for all locals
        extractor.create_local_channels(mir);

        // Phase 2: Detect loops (back edges)
        extractor.detect_loops(mir);

        // Phase 3: Process each basic block
        for (bb_idx, block) in mir.basic_blocks.iter().enumerate() {
            let bb_id = BasicBlockId::new(bb_idx as u32);
            extractor.current_block = Some(bb_id);
            extractor.process_block(bb_id, block, mir);
        }

        // Phase 4: Mark external inputs/outputs
        extractor.mark_io(mir);

        extractor.dfg
    }

    /// Create channels for all MIR locals
    fn create_local_channels(&mut self, mir: &FunctionMIR) {
        for (idx, local_decl) in mir.locals.iter().enumerate() {
            let local = Local::new(idx as u32);
            let token_type = Self::mir_ty_to_token_type(&local_decl.ty);
            let channel = self.dfg.add_channel(token_type);
            self.local_to_channel.insert(local, channel);
        }
    }

    /// Detect loop back edges using DFS
    fn detect_loops(&mut self, mir: &FunctionMIR) {
        if mir.basic_blocks.is_empty() {
            return;
        }

        let mut visited = vec![false; mir.basic_blocks.len()];
        let mut in_stack = vec![false; mir.basic_blocks.len()];
        self.dfs_detect_loops(mir, BasicBlockId::START, &mut visited, &mut in_stack);
    }

    fn dfs_detect_loops(
        &mut self,
        mir: &FunctionMIR,
        bb: BasicBlockId,
        visited: &mut [bool],
        in_stack: &mut [bool],
    ) {
        let idx = bb.0 as usize;
        if idx >= mir.basic_blocks.len() {
            return;
        }

        visited[idx] = true;
        in_stack[idx] = true;

        let successors = self.get_successors(&mir.basic_blocks[idx].terminator);
        for succ in successors {
            let succ_idx = succ.0 as usize;
            if succ_idx >= mir.basic_blocks.len() {
                continue;
            }

            if !visited[succ_idx] {
                self.dfs_detect_loops(mir, succ, visited, in_stack);
            } else if in_stack[succ_idx] {
                // Back edge found - this is a loop
                self.back_edges.push((bb, succ));
            }
        }

        in_stack[idx] = false;
    }

    fn get_successors(&self, terminator: &Terminator) -> Vec<BasicBlockId> {
        match terminator {
            Terminator::Return { .. }
            | Terminator::Abort { .. }
            | Terminator::Unreachable { .. }
            | Terminator::Cancel { .. } => {
                vec![]
            }
            Terminator::Goto { target, .. } => vec![*target],
            Terminator::SwitchInt { targets, .. } => {
                let mut succs: Vec<_> = targets.branches.iter().map(|(_, bb)| *bb).collect();
                succs.push(targets.otherwise);
                succs
            }
            Terminator::Call { target, .. } => vec![*target],
            Terminator::Spawn { target, .. } => vec![*target],
            Terminator::TaskAwait { target, .. } => vec![*target],
            Terminator::TaskGroupEnter {
                body, join_block, ..
            } => vec![*body, *join_block],
            Terminator::TaskGroupExit { target, .. } => vec![*target],
            Terminator::ChannelRecv {
                target,
                closed_target,
                ..
            } => vec![*target, *closed_target],
            Terminator::ChannelSend {
                target,
                closed_target,
                ..
            } => vec![*target, *closed_target],
            Terminator::Select { arms, default, .. } => {
                let mut succs: Vec<_> = arms.iter().map(|arm| arm.target).collect();
                if let Some(def) = default {
                    succs.push(*def);
                }
                succs
            }
        }
    }

    /// Process a basic block
    fn process_block(&mut self, bb_id: BasicBlockId, block: &BasicBlock, mir: &FunctionMIR) {
        // Create block entry channel if needed
        if !self.block_entry.contains_key(&bb_id) {
            let entry = self.dfg.add_channel(TokenType::Unit);
            self.block_entry.insert(bb_id, entry);
        }

        // Process statements
        for stmt in &block.statements {
            self.process_statement(stmt);
        }

        // Process terminator
        self.process_terminator(&block.terminator, mir);
    }

    /// Process a statement
    fn process_statement(&mut self, stmt: &Statement) {
        match stmt {
            Statement::Assign { place, rvalue, .. } => {
                self.process_assignment(place, rvalue);
            }
            Statement::StorageLive { .. } | Statement::StorageDead { .. } | Statement::Nop => {
                // No dataflow impact
            }
        }
    }

    /// Process an assignment
    fn process_assignment(&mut self, place: &Place, rvalue: &Rvalue) {
        let dest_channel = self.place_to_channel(place);

        match rvalue {
            Rvalue::Use(operand) => {
                // Simple copy/move - create identity operator or direct connection
                if let Some(src_channel) = self.operand_to_channel(operand) {
                    // Identity: just connect channels
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add, // Add 0 is identity
                        inputs: vec![src_channel],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::BinaryOp { op, left, right } => {
                let compute_op = Self::binop_to_compute_op(op);
                let left_ch = self.operand_to_channel(left);
                let right_ch = self.operand_to_channel(right);

                if let (Some(l), Some(r)) = (left_ch, right_ch) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: compute_op,
                        inputs: vec![l, r],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::UnaryOp { op, operand } => {
                let compute_op = Self::unop_to_compute_op(op);
                if let Some(src) = self.operand_to_channel(operand) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: compute_op,
                        inputs: vec![src],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::Ref { place, .. } => {
                // Reference creation - treat as address computation
                if let Some(src) = self.local_to_channel.get(&place.local) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add, // Identity
                        inputs: vec![*src],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::Aggregate { operands, .. } => {
                let inputs: Vec<_> = operands
                    .iter()
                    .filter_map(|op| self.operand_to_channel(op))
                    .collect();

                if !inputs.is_empty() {
                    // For now, treat as merge
                    self.dfg.add_operator(DfOperator::Merge {
                        inputs,
                        output: dest_channel,
                    });
                }
            }

            Rvalue::Cast { operand, kind, .. } => {
                if let Some(src) = self.operand_to_channel(operand) {
                    let op = match kind {
                        crate::CastKind::IntToFloat => ComputeOp::IntToFloat,
                        crate::CastKind::FloatToInt => ComputeOp::FloatToInt,
                        crate::CastKind::IntToInt => ComputeOp::SignExtend,
                        crate::CastKind::FloatToFloat => ComputeOp::Truncate,
                        crate::CastKind::PtrToPtr => ComputeOp::ZeroExtend,
                        crate::CastKind::Bitcast => ComputeOp::ZeroExtend,
                    };
                    self.dfg.add_operator(DfOperator::Compute {
                        op,
                        inputs: vec![src],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::Len { place } => {
                // Length extraction reads the length field from an array/slice
                // structure. This is semantically a memory load (reading a
                // field), modeled as a Load operation in the dataflow graph.
                if let Some(src) = self.local_to_channel.get(&place.local) {
                    self.dfg.add_operator(DfOperator::Memory {
                        op: MemoryOp::Load,
                        address: *src,
                        data: dest_channel,
                        ordering: vec![],
                    });
                }
            }

            Rvalue::Discriminant { place } => {
                // Discriminant extraction reads the tag field of an enum and
                // truncates it to the discriminant integer type. This is a
                // load + truncate operation; we model the dominant cost as the
                // memory read of the tag field.
                if let Some(src) = self.local_to_channel.get(&place.local) {
                    self.dfg.add_operator(DfOperator::Memory {
                        op: MemoryOp::Load,
                        address: *src,
                        data: dest_channel,
                        ordering: vec![],
                    });
                }
            }

            // SIMD operations
            Rvalue::SimdBinaryOp {
                op, left, right, ..
            } => {
                let compute_op = Self::simd_op_to_compute_op(op);
                let left_ch = self.operand_to_channel(left);
                let right_ch = self.operand_to_channel(right);

                if let (Some(l), Some(r)) = (left_ch, right_ch) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: compute_op,
                        inputs: vec![l, r],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::SimdLoad { source, .. } => {
                if let Some(src) = self.local_to_channel.get(&source.local) {
                    self.dfg.add_operator(DfOperator::Memory {
                        op: MemoryOp::Load,
                        address: *src,
                        data: dest_channel,
                        ordering: vec![],
                    });
                }
            }

            Rvalue::SimdStore { value, dest, .. } => {
                let val_ch = self.operand_to_channel(value);
                if let (Some(val), Some(dst)) = (val_ch, self.local_to_channel.get(&dest.local)) {
                    self.dfg.add_operator(DfOperator::Memory {
                        op: MemoryOp::Store,
                        address: *dst,
                        data: val,
                        ordering: vec![],
                    });
                }
            }

            Rvalue::SimdSplat { value, .. } => {
                if let Some(src) = self.operand_to_channel(value) {
                    // Splat is just broadcasting
                    self.dfg.add_operator(DfOperator::Split {
                        input: src,
                        outputs: vec![dest_channel],
                    });
                }
            }

            // Channel operations - convert to operators
            Rvalue::ChannelCreate { .. } => {
                // Channel creation is a source of values
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_channel.0,
                    output: dest_channel,
                });
            }

            Rvalue::ChannelTryRecv { channel } => {
                if let Some(ch) = self.operand_to_channel(channel) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add,
                        inputs: vec![ch],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::ChannelTrySend { channel, value } => {
                let ch = self.operand_to_channel(channel);
                let val = self.operand_to_channel(value);
                if let (Some(c), Some(v)) = (ch, val) {
                    self.dfg.add_operator(DfOperator::Merge {
                        inputs: vec![c, v],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::ChannelSender { channel } | Rvalue::ChannelReceiver { channel } => {
                if let Some(ch) = self.operand_to_channel(channel) {
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add,
                        inputs: vec![ch],
                        output: dest_channel,
                    });
                }
            }

            Rvalue::ChannelClose { channel } => {
                if let Some(ch) = self.operand_to_channel(channel) {
                    self.dfg.add_operator(DfOperator::Sink {
                        input: ch,
                        external_id: dest_channel.0,
                    });
                }
            }

            Rvalue::IsCancelled | Rvalue::CurrentTask => {
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_channel.0,
                    output: dest_channel,
                });
            }
            Rvalue::Try(_operand) => {
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_channel.0,
                    output: dest_channel,
                });
            }
            Rvalue::EnumVariant { .. } => {
                // Enum variant is a source (constant value)
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_channel.0,
                    output: dest_channel,
                });
            }
        }
    }

    /// Process a terminator
    fn process_terminator(&mut self, terminator: &Terminator, _mir: &FunctionMIR) {
        match terminator {
            Terminator::Return { .. } => {
                // Return value goes to sink
                if let Some(&ret_ch) = self.local_to_channel.get(&Local::RETURN) {
                    self.dfg.add_operator(DfOperator::Sink {
                        input: ret_ch,
                        external_id: 0, // Return is always external_id 0
                    });
                    self.dfg.add_output(ret_ch);
                }
            }

            Terminator::Goto { target, .. } => {
                // Unconditional jump - connect to target block
                let _entry = self.get_or_create_block_entry(*target);
                // Goto doesn't need explicit dataflow - control flows implicitly
            }

            Terminator::SwitchInt {
                discriminant,
                targets,
                ..
            } => {
                // Conditional branch - use Steer
                if let Some(cond_ch) = self.operand_to_channel(discriminant) {
                    // For simple if-else (one branch + otherwise)
                    if targets.branches.len() == 1 {
                        let true_target = targets.branches[0].1;
                        let false_target = targets.otherwise;

                        let true_entry = self.get_or_create_block_entry(true_target);
                        let false_entry = self.get_or_create_block_entry(false_target);

                        // Create data channel to route
                        let data = self.dfg.add_channel(TokenType::Unit);
                        self.dfg.add_operator(DfOperator::Constant {
                            value: TokenValue::Unit,
                            output: data,
                            repeat: None,
                        });

                        self.dfg.add_operator(DfOperator::Steer {
                            decider: cond_ch,
                            data,
                            true_out: true_entry,
                            false_out: false_entry,
                        });
                    } else {
                        // Multi-way switch - use Select
                        // Collect branch entries first to avoid borrow checker issues
                        let branch_bbs: Vec<_> =
                            targets.branches.iter().map(|(_, bb)| *bb).collect();
                        let mut outputs = Vec::with_capacity(branch_bbs.len() + 1);
                        for bb in branch_bbs {
                            outputs.push(self.get_or_create_block_entry(bb));
                        }
                        outputs.push(self.get_or_create_block_entry(targets.otherwise));

                        let data = self.dfg.add_channel(TokenType::Unit);
                        self.dfg.add_operator(DfOperator::Constant {
                            value: TokenValue::Unit,
                            output: data,
                            repeat: None,
                        });

                        self.dfg.add_operator(DfOperator::Split {
                            input: data,
                            outputs,
                        });
                    }
                }
            }

            Terminator::Call {
                func: _,
                args,
                destination,
                target: _,
                func_name: _,
                ..
            } => {
                // Function call - create compute node
                let inputs: Vec<_> = args
                    .iter()
                    .filter_map(|arg| self.operand_to_channel(arg))
                    .collect();

                let dest_ch = self.place_to_channel(destination);

                // Model the function call as a Merge: the result depends on all
                // arguments. This correctly captures data dependencies (the output
                // is available only after all inputs are ready) and the energy cost
                // of the call is dominated by the callee's body, not the call
                // instruction itself. For inter-procedural energy analysis, the
                // callee's DFG would be inlined here.
                if !inputs.is_empty() {
                    self.dfg.add_operator(DfOperator::Merge {
                        inputs,
                        output: dest_ch,
                    });
                }
            }

            Terminator::Spawn {
                func: _,
                args: _,
                destination,
                ..
            } => {
                // Spawn - create a fork in the dataflow
                // Spawned computation is represented as an async source
                // (the actual args would be processed by the spawned task's DFG)
                let dest_ch = self.place_to_channel(destination);
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_ch.0,
                    output: dest_ch,
                });
            }

            Terminator::TaskAwait {
                task, destination, ..
            } => {
                // Await - synchronization point
                if let Some(task_ch) = self.operand_to_channel(task) {
                    let dest_ch = self.place_to_channel(destination);
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add,
                        inputs: vec![task_ch],
                        output: dest_ch,
                    });
                }
            }

            Terminator::ChannelRecv {
                channel,
                destination,
                ..
            } => {
                if let Some(ch) = self.operand_to_channel(channel) {
                    let dest_ch = self.place_to_channel(destination);
                    self.dfg.add_operator(DfOperator::Compute {
                        op: ComputeOp::Add,
                        inputs: vec![ch],
                        output: dest_ch,
                    });
                }
            }

            Terminator::ChannelSend { channel, value, .. } => {
                let ch = self.operand_to_channel(channel);
                let val = self.operand_to_channel(value);
                if let (Some(c), Some(v)) = (ch, val) {
                    // Send merges channel and value
                    let out = self.dfg.add_channel(TokenType::Unit);
                    self.dfg.add_operator(DfOperator::Merge {
                        inputs: vec![c, v],
                        output: out,
                    });
                }
            }

            // Other terminators - basic handling
            Terminator::Abort { .. }
            | Terminator::Unreachable { .. }
            | Terminator::Cancel { .. } => {
                // Terminal nodes - no dataflow continuation
            }

            Terminator::TaskGroupEnter { destination, .. } => {
                let dest_ch = self.place_to_channel(destination);
                self.dfg.add_operator(DfOperator::Source {
                    external_id: dest_ch.0,
                    output: dest_ch,
                });
            }

            Terminator::TaskGroupExit { group, .. } => {
                if let Some(ch) = self.operand_to_channel(group) {
                    self.dfg.add_operator(DfOperator::Sink {
                        input: ch,
                        external_id: ch.0,
                    });
                }
            }

            Terminator::Select {
                arms, destination, ..
            } => {
                let inputs: Vec<_> = arms
                    .iter()
                    .filter_map(|arm| match &arm.operation {
                        crate::ChannelOp::Recv { channel } => self.operand_to_channel(channel),
                        crate::ChannelOp::Send { channel, .. } => self.operand_to_channel(channel),
                        crate::ChannelOp::Timeout { .. } => None,
                    })
                    .collect();

                let dest_ch = self.place_to_channel(destination);
                if !inputs.is_empty() {
                    self.dfg.add_operator(DfOperator::Merge {
                        inputs,
                        output: dest_ch,
                    });
                }
            }
        }
    }

    fn get_or_create_block_entry(&mut self, bb: BasicBlockId) -> ChannelId {
        if let Some(&ch) = self.block_entry.get(&bb) {
            ch
        } else {
            let ch = self.dfg.add_channel(TokenType::Unit);
            self.block_entry.insert(bb, ch);
            ch
        }
    }

    /// Mark external inputs and outputs
    fn mark_io(&mut self, mir: &FunctionMIR) {
        // Parameters are inputs
        for param in &mir.params {
            if let Some(&ch) = self.local_to_channel.get(param) {
                self.dfg.add_input(ch);

                // Add source operator for parameters
                self.dfg.add_operator(DfOperator::Source {
                    external_id: ch.0,
                    output: ch,
                });
            }
        }

        // Return is output (already handled in Return terminator)
    }

    /// Convert a place to its channel
    fn place_to_channel(&self, place: &Place) -> ChannelId {
        // For now, just use the base local
        // More complex projections would need additional channels
        *self
            .local_to_channel
            .get(&place.local)
            .unwrap_or(&ChannelId::new(0))
    }

    /// Convert an operand to a channel
    fn operand_to_channel(&mut self, operand: &Operand) -> Option<ChannelId> {
        match operand {
            Operand::Copy(place) | Operand::Move(place) => {
                self.local_to_channel.get(&place.local).copied()
            }
            Operand::Constant(constant) => {
                // Create a constant operator
                let token_type = Self::mir_ty_to_token_type(&constant.ty);
                let output = self.dfg.add_channel(token_type);
                let value = Self::literal_to_token_value(&constant.literal);

                self.dfg.add_operator(DfOperator::Constant {
                    value,
                    output,
                    repeat: None,
                });

                Some(output)
            }
        }
    }

    /// Convert MIR type to token type
    fn mir_ty_to_token_type(ty: &Ty) -> TokenType {
        match ty {
            Ty::Bool => TokenType::Bool,
            Ty::Int(int_ty) => {
                let bits = match int_ty {
                    crate::IntTy::I8 => 8,
                    crate::IntTy::I16 => 16,
                    crate::IntTy::I32 => 32,
                    crate::IntTy::I64 => 64,
                    crate::IntTy::Isize => 64,
                };
                TokenType::Int { bits, signed: true }
            }
            Ty::Uint(uint_ty) => {
                let bits = match uint_ty {
                    crate::UintTy::U8 => 8,
                    crate::UintTy::U16 => 16,
                    crate::UintTy::U32 => 32,
                    crate::UintTy::U64 => 64,
                    crate::UintTy::Usize => 64,
                };
                TokenType::Int {
                    bits,
                    signed: false,
                }
            }
            Ty::Float(float_ty) => {
                let bits = match float_ty {
                    crate::FloatTy::F16 | crate::FloatTy::BF16 => 16,
                    crate::FloatTy::F32 => 32,
                    crate::FloatTy::F64 => 64,
                };
                TokenType::Float { bits }
            }
            Ty::Ref { .. } | Ty::RawPtr { .. } => TokenType::Ptr,
            Ty::Unit => TokenType::Unit,
            Ty::Tuple(tys) => {
                TokenType::Tuple(tys.iter().map(Self::mir_ty_to_token_type).collect())
            }
            Ty::Array { element, size } => TokenType::Array {
                element: Box::new(Self::mir_ty_to_token_type(element)),
                size: *size as usize,
            },
            _ => TokenType::Unit, // Default for complex types
        }
    }

    /// Convert literal to token value
    fn literal_to_token_value(literal: &crate::Literal) -> TokenValue {
        match literal {
            crate::Literal::Bool(b) => TokenValue::Bool(*b),
            crate::Literal::Int(i, _) => TokenValue::Int(*i),
            crate::Literal::Uint(u, _) => TokenValue::Uint(*u),
            crate::Literal::Float(f, _) => TokenValue::Float(*f),
            crate::Literal::Char(c) => TokenValue::Uint(*c as u64),
            // String literals are heap-allocated data referenced by pointer.
            // In the dataflow graph, they are modeled as a pointer token with
            // address 0 (the actual address is assigned at link time). This
            // correctly represents the energy cost: reading a string literal
            // involves a pointer load from the data section.
            crate::Literal::String(_) => TokenValue::Ptr(0),
        }
    }

    /// Convert BinOp to ComputeOp
    fn binop_to_compute_op(op: &BinOp) -> ComputeOp {
        match op {
            BinOp::Add => ComputeOp::Add,
            BinOp::Sub => ComputeOp::Sub,
            BinOp::Mul => ComputeOp::Mul,
            BinOp::Div => ComputeOp::Div,
            BinOp::Rem => ComputeOp::Rem,
            BinOp::BitAnd => ComputeOp::BitAnd,
            BinOp::BitOr => ComputeOp::BitOr,
            BinOp::BitXor => ComputeOp::BitXor,
            BinOp::Shl => ComputeOp::Shl,
            BinOp::Shr => ComputeOp::Shr,
            BinOp::Eq => ComputeOp::Eq,
            BinOp::Ne => ComputeOp::Ne,
            BinOp::Lt => ComputeOp::Lt,
            BinOp::Le => ComputeOp::Le,
            BinOp::Gt => ComputeOp::Gt,
            BinOp::Ge => ComputeOp::Ge,
            BinOp::And => ComputeOp::And,
            BinOp::Or => ComputeOp::Or,
        }
    }

    /// Convert UnOp to ComputeOp
    fn unop_to_compute_op(op: &UnOp) -> ComputeOp {
        match op {
            UnOp::Neg => ComputeOp::Neg,
            UnOp::Not => ComputeOp::Not,
        }
    }

    /// Convert SIMD op to ComputeOp
    fn simd_op_to_compute_op(op: &crate::SimdOp) -> ComputeOp {
        match op {
            crate::SimdOp::Add => ComputeOp::Add,
            crate::SimdOp::Sub => ComputeOp::Sub,
            crate::SimdOp::Mul => ComputeOp::Mul,
            crate::SimdOp::Div => ComputeOp::Div,
            crate::SimdOp::Fma => ComputeOp::Fma,
        }
    }
}

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

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

    fn dummy_span() -> Span {
        Span::dummy()
    }

    #[test]
    fn test_simple_extraction() {
        // Create a simple MIR function: fn add(a: i32, b: i32) -> i32 { a + b }
        let mut mir = FunctionMIR::new(
            crate::FunctionId::new(0),
            Symbol::intern("add"),
            Ty::Int(crate::IntTy::I32),
            dummy_span(),
        );

        // Add parameters
        let a = mir.add_local(LocalDecl::new(
            Some(Symbol::intern("a")),
            Ty::Int(crate::IntTy::I32),
            false,
            dummy_span(),
        ));
        let b = mir.add_local(LocalDecl::new(
            Some(Symbol::intern("b")),
            Ty::Int(crate::IntTy::I32),
            false,
            dummy_span(),
        ));
        mir.params = vec![a, b];

        // Create basic block with add and return
        let mut block = BasicBlock::new(Terminator::Return { span: dummy_span() });
        block.push_statement(Statement::Assign {
            place: Place::return_place(),
            rvalue: Rvalue::BinaryOp {
                op: BinOp::Add,
                left: Operand::from_local(a),
                right: Operand::from_local(b),
            },
            span: dummy_span(),
        });
        mir.add_block(block);

        // Extract DFG
        let dfg = DfgExtractor::extract(&mir);

        assert!(dfg.operators.len() >= 1); // At least the add operator
        assert!(!dfg.channels.is_empty());
    }
}