ir.rs

use hercules_ir::ir::{Device, Schedule};

use crate::parser;

#[derive(Debug, Copy, Clone)]
pub enum Pass {
    ArraySLF,
    ArrayToProduct,
    AutoOutline,
    CCP,
    CRC,
    CleanMonoidReduces,
    ConstInline,
    DCE,
    FloatCollections,
    ForkChunk,
    ForkCoalesce,
    ForkDimMerge,
    ForkExtend,
    ForkFissionBufferize,
    ForkFusion,
    ForkGuardElim,
    ForkInterchange,
    ForkReshape,
    ForkSplit,
    ForkUnroll,
    Forkify,
    GCM,
    GVN,
    InferSchedules,
    Inline,
    InterproceduralSROA,
    LiftDCMath,
    LoopBoundCanon,
    Outline,
    PhiElim,
    Predication,
    Print,
    ReduceSLF,
    Rename,
    ReuseProducts,
    RewriteMathExpressions,
    SLF,
    SROA,
    Serialize,
    SimplifyCFG,
    Unforkify,
    UnforkifyOne,
    Verify,
    WritePredication,
    Xdot,
}

impl Pass {
    pub fn is_valid_num_args(&self, num: usize) -> bool {
        match self {
            Pass::ArrayToProduct => num == 0 || num == 1,
            Pass::ConstInline => num == 0 || num == 1,
            Pass::ForkChunk => num == 4,
            Pass::ForkExtend => num == 1,
            Pass::ForkFissionBufferize => num == 2 || num == 1,
            Pass::ForkInterchange => num == 2,
            Pass::ForkReshape => true,
            Pass::InterproceduralSROA => num == 0 || num == 1,
            Pass::Print => num == 1,
            Pass::Rename => num == 1,
            Pass::SROA => num == 0 || num == 1,
            Pass::Xdot => num == 0 || num == 1,
            _ => num == 0,
        }
    }

    pub fn valid_arg_nums(&self) -> &'static str {
        match self {
            Pass::ArrayToProduct => "0 or 1",
            Pass::ConstInline => "0 or 1",
            Pass::ForkChunk => "4",
            Pass::ForkExtend => "1",
            Pass::ForkFissionBufferize => "1 or 2",
            Pass::ForkInterchange => "2",
            Pass::ForkReshape => "any",
            Pass::InterproceduralSROA => "0 or 1",
            Pass::Print => "1",
            Pass::Rename => "1",
            Pass::SROA => "0 or 1",
            Pass::Xdot => "0 or 1",
            _ => "0",
        }
    }
}

#[derive(Debug, Clone)]
pub enum Selector {
    Everything(),
    Selection(Vec<ScheduleExp>),
}

#[derive(Debug, Clone)]
pub enum ScheduleExp {
    Variable {
        var: String,
    },
    Integer {
        val: usize,
    },
    Boolean {
        val: bool,
    },
    String {
        val: String,
    },
    Field {
        collect: Box<ScheduleExp>,
        field: String,
    },
    RunPass {
        pass: Pass,
        args: Vec<ScheduleExp>,
        on: Selector,
    },
    DeleteUncalled {
        on: Selector,
    },
    Feature {
        feature: Box<ScheduleExp>,
    },
    Record {
        fields: Vec<(String, ScheduleExp)>,
    },
    Block {
        body: Vec<ScheduleStmt>,
        res: Box<ScheduleExp>,
    },
    SetOp {
        op: parser::SetOp,
        lhs: Box<ScheduleExp>,
        rhs: Box<ScheduleExp>,
    },
    Tuple {
        exprs: Vec<ScheduleExp>,
    },
    TupleField {
        lhs: Box<ScheduleExp>,
        field: usize,
    },
    // This is used to "box" a selection by evaluating it at one point and then
    // allowing it to be used as a selector later on
    Selection {
        selection: Selector,
    },
}

#[derive(Debug, Copy, Clone)]
pub enum FixpointLimit {
    NoLimit(),
    PrintIter(),
    StopAfter(usize),
    PanicAfter(usize),
}

#[derive(Debug, Clone)]
pub enum ScheduleStmt {
    Fixpoint {
        body: Box<ScheduleStmt>,
        limit: FixpointLimit,
    },
    Block {
        body: Vec<ScheduleStmt>,
    },
    Let {
        var: String,
        exp: ScheduleExp,
    },
    Assign {
        var: String,
        exp: ScheduleExp,
    },
    AddSchedule {
        sched: Schedule,
        on: Selector,
    },
    AddDevice {
        device: Device,
        on: Selector,
    },
    IfThenElse {
        cond: ScheduleExp,
        thn: Box<ScheduleStmt>,
        els: Box<ScheduleStmt>,
    },
}