cubecl_core/frontend/

branch.rs

use alloc::vec::Vec;
use cubecl_ir::{
    ExpandState, ExpandValue, OpInserter,
    dialect::{
        branch::{ConditionOp, IfOp, RangeLoopOp, SwitchOp, UnreachableOp, WhileOp},
        general::BoolAndOp,
    },
    pliron::{irbuild::inserter::Inserter, r#type::TypedHandle},
    read_value,
    rewrite::{op_insertion_point_in_block, set_inserter_before_terminator},
};
use pliron::{
    basic_block::BasicBlock,
    builtin::{
        attributes::IntegerAttr,
        op_interfaces::OneRegionInterface,
        types::{IntegerType, Signedness},
    },
    irbuild::{
        listener::DummyListener,
        rewriter::{IRRewriter, Rewriter},
    },
    op::Op,
    region::Region,
    r#type::Typed,
    utils::apint::{APInt, bw},
    value::Value,
};

use crate::{
    IntoRuntime,
    frontend::{
        CubePrimitive, ReadValue, RuntimeAssign, assign, assign_binop_expand, binary_expand,
        init_mut_of_type,
    },
    prelude::{CubeEnum, ExpandTypeClone},
};
use crate::{ir::Scope, prelude::Assign};

use super::{Int, NativeExpand};

/// Something that can be iterated on by a for loop. Currently only includes `Range`, `StepBy` and
/// `Sequence`.
pub trait Iterable: Sized {
    type Item;

    /// Expand a runtime loop without unrolling
    ///
    /// # Arguments
    /// * `scope` - the expansion scope
    /// * `body` - the loop body to be executed repeatedly
    ///
    /// The body is a trait object rather than an `impl FnMut` so that this
    /// expands once per iterable instead of once per loop body. See
    /// [`for_expand`].
    fn expand(self, scope: &Scope, body: &mut dyn FnMut(&Scope, Self::Item));
    /// Expand an unrolled loop. The body should be invoced `n` times, where `n` is the number of
    /// iterations.
    ///
    /// # Arguments
    /// * `scope` - the expansion scope
    /// * `body` - the loop body to be executed repeatedly
    fn expand_unroll(self, scope: &Scope, body: &mut dyn FnMut(&Scope, Self::Item));
    /// Return the comptime length of this iterable, if possible
    fn const_len(&self) -> Option<usize> {
        None
    }
}

/// Expand a `for` loop over `range`.
///
/// Every `for` in every kernel reaches here with a body of its own anonymous
/// closure type. Staying generic over that type all the way down would expand
/// the whole loop machinery once per loop written anywhere — thousands of
/// near-identical copies in a kernel-heavy crate. So this is a shim: it erases
/// the body to a trait object straight away and hands off to
/// [`for_expand_dyn`], which expands once per iterable. Expansion runs on the
/// host while building a kernel, never on the device, so the indirection costs
/// nothing that reaches the GPU.
pub fn for_expand<I: Iterable>(
    scope: &Scope,
    range: I,
    unroll: bool,
    mut body: impl FnMut(&Scope, I::Item),
) {
    for_expand_dyn(scope, range, unroll, &mut body);
}

fn for_expand_dyn<I: Iterable>(
    scope: &Scope,
    range: I,
    unroll: bool,
    body: &mut dyn FnMut(&Scope, I::Item),
) {
    if unroll || range.const_len() == Some(1) {
        range.expand_unroll(scope, body);
    } else {
        range.expand(scope, body);
    }
}

pub fn if_expand(scope: &Scope, condition: NativeExpand<bool>, block: impl FnOnce(&Scope)) {
    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
    match comptime_cond {
        Some(cond) => {
            if cond {
                block(scope);
            }
        }
        None => {
            let cond = condition.read_value(scope);
            let if_op = IfOp::new(scope.ctx_mut(), cond);

            let then_block = if_op.then_block(scope.ctx());
            let then_child = scope.branch_child(OpInserter::new_at_block_end(then_block));
            block(&then_child);
            then_child.terminate_yield();

            let else_block = if_op.else_block(scope.ctx());
            let else_child = scope.branch_child(OpInserter::new_at_block_end(else_block));
            else_child.terminate_yield();

            scope.register(&if_op);
            scope.update_flags_after_branch(&[then_child, else_child]);
        }
    }
}

#[allow(clippy::large_enum_variant)]
pub enum IfElseExpand {
    ComptimeThen,
    ComptimeElse,
    Runtime {
        runtime_cond: NativeExpand<bool>,
        if_op: IfOp,
        then_child: Scope,
    },
}

impl IfElseExpand {
    pub fn or_else(self, scope: &Scope, else_block: impl FnOnce(&Scope)) {
        match self {
            Self::Runtime {
                if_op, then_child, ..
            } => {
                let else_body = if_op.else_block(scope.ctx());
                let else_child = scope.branch_child(OpInserter::new_at_block_end(else_body));
                else_block(&else_child);
                else_child.terminate_yield();

                scope.register(&if_op);
                scope.update_flags_after_branch(&[then_child, else_child]);
            }
            Self::ComptimeElse => else_block(scope),
            Self::ComptimeThen => (),
        }
    }
}

pub fn if_else_expand(
    scope: &Scope,
    condition: NativeExpand<bool>,
    then_block: impl FnOnce(&Scope),
) -> IfElseExpand {
    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
    match comptime_cond {
        Some(true) => {
            then_block(scope);
            IfElseExpand::ComptimeThen
        }
        Some(false) => IfElseExpand::ComptimeElse,
        None => {
            let cond = condition.read_value(scope);
            let if_op = IfOp::new(scope.ctx_mut(), cond);
            let if_block = if_op.then_block(scope.ctx());
            let then_child = scope.branch_child(OpInserter::new_at_block_end(if_block));
            then_block(&then_child);
            then_child.terminate_yield();

            IfElseExpand::Runtime {
                runtime_cond: condition,
                if_op,
                then_child,
            }
        }
    }
}

#[allow(clippy::large_enum_variant)]
pub enum IfElseExprExpand<C: Assign> {
    ComptimeThen(C),
    ComptimeElse,
    Runtime {
        runtime_cond: NativeExpand<bool>,
        out: C,
        if_op: IfOp,
        then_child: Scope,
    },
}

impl<C: Assign> IfElseExprExpand<C> {
    pub fn or_else<R: RuntimeAssign<Expand = C>>(
        self,
        scope: &Scope,
        else_block: impl FnOnce(&Scope) -> R,
    ) -> C {
        match self {
            Self::Runtime {
                mut out,
                if_op,
                then_child,
                ..
            } => {
                let else_body = if_op.else_block(scope.ctx());
                let else_child = scope.branch_child(OpInserter::new_at_block_end(else_body));
                let ret = else_block(&else_child);
                out.__expand_assign_method(&else_child, ret.into_expand(scope));
                else_child.terminate_yield();

                scope.register(&if_op);
                scope.update_flags_after_branch(&[then_child, else_child]);
                out
            }
            Self::ComptimeElse => else_block(scope).into_expand(scope),
            Self::ComptimeThen(ret) => ret,
        }
    }
}

pub fn if_else_expr_expand<C: RuntimeAssign>(
    scope: &Scope,
    condition: NativeExpand<bool>,
    then_block: impl FnOnce(&Scope) -> C,
) -> IfElseExprExpand<C::Expand> {
    let comptime_cond = condition.expand.as_const().map(|it| it.as_bool());
    match comptime_cond {
        Some(true) => {
            let ret = then_block(scope);
            IfElseExprExpand::ComptimeThen(ret.into_expand(scope))
        }
        Some(false) => IfElseExprExpand::ComptimeElse,
        None => {
            let cond = condition.read_value(scope);
            let if_op = IfOp::new(scope.ctx_mut(), cond);
            let then_body = if_op.then_block(scope.ctx());
            let then_child = scope.branch_child(OpInserter::new_at_block_end(then_body));
            let ret = then_block(&then_child);
            let mut out = ret.init_mut(scope);
            out.__expand_assign_method(&then_child, ret.into_expand(scope));
            then_child.terminate_yield();

            IfElseExprExpand::Runtime {
                runtime_cond: condition,
                out,
                if_op,
                then_child,
            }
        }
    }
}

pub struct SwitchExpand<I: Int> {
    switch_op: SwitchOp,
    cases: Vec<I>,
    children: Vec<Scope>,
}

impl<I: Int> SwitchExpand<I> {
    pub fn case(mut self, scope: &Scope, value: impl Int, block: impl FnOnce(&Scope)) -> Self {
        let value = I::from(value).unwrap();
        self.cases.push(value);
        let body = self.switch_op.append_case_block(scope.ctx_mut());
        let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
        block(&case_child);
        case_child.terminate_yield();
        self.children.push(case_child);
        self
    }

    pub fn finish(self, scope: &Scope) {
        let cases = self.cases.into_iter().map(|case| {
            let ty = I::__expand_as_type(scope);
            let ty = TypedHandle::<IntegerType>::from_handle(ty, scope.ctx()).unwrap();
            let width = bw(ty.deref(scope.ctx()).width() as usize);
            IntegerAttr::new(ty, APInt::from_i128(case.to_i128().unwrap(), width))
        });
        self.switch_op.set_attr_cases(scope.ctx(), cases);
        scope.register(&self.switch_op);
        scope.update_flags_after_branch(&self.children);
    }
}

pub fn switch_expand<I: Int>(
    scope: &Scope,
    value: NativeExpand<I>,
    default_block: impl FnOnce(&Scope),
) -> SwitchExpand<I> {
    let value = value.read_value(scope);
    let switch_op = SwitchOp::new(scope.ctx_mut(), value);

    let default_body = switch_op.default_block(scope.ctx());
    let default_child = scope.branch_child(OpInserter::new_at_block_end(default_body));
    default_block(&default_child);
    default_child.terminate_yield();

    SwitchExpand {
        switch_op,
        cases: Vec::new(),
        children: alloc::vec![default_child],
    }
}

pub struct SwitchExpandExpr<I: Int, C: Assign> {
    switch_op: SwitchOp,
    cases: Vec<I>,
    children: Vec<Scope>,
    out: C,
}

impl<I: Int, C: Assign> SwitchExpandExpr<I, C> {
    pub fn case<T: RuntimeAssign<Expand = C>>(
        mut self,
        scope: &Scope,
        value: impl Int,
        block: impl FnOnce(&Scope) -> T,
    ) -> Self {
        let value = I::from(value).unwrap();
        self.cases.push(value);
        let body = self.switch_op.append_case_block(scope.ctx_mut());
        let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
        let ret = block(&case_child);
        self.out
            .__expand_assign_method(&case_child, ret.into_expand(scope));
        case_child.terminate_yield();
        self.children.push(case_child);
        self
    }

    pub fn finish(self, scope: &Scope) -> C {
        let cases = self.cases.into_iter().map(|case| {
            let ty = I::__expand_as_type(scope);
            let ty = TypedHandle::<IntegerType>::from_handle(ty, scope.ctx()).unwrap();
            let width = bw(ty.deref(scope.ctx()).width() as usize);
            IntegerAttr::new(ty, APInt::from_i128(case.to_i128().unwrap(), width))
        });
        self.switch_op.set_attr_cases(scope.ctx(), cases);
        scope.register(&self.switch_op);
        scope.update_flags_after_branch(&self.children);
        self.out
    }
}

pub fn switch_expand_expr<I: Int, C: RuntimeAssign>(
    scope: &Scope,
    value: NativeExpand<I>,
    default_block: impl FnOnce(&Scope) -> C,
) -> SwitchExpandExpr<I, C::Expand> {
    let value = value.read_value(scope);
    let switch_op = SwitchOp::new(scope.ctx_mut(), value);

    let default_body = switch_op.default_block(scope.ctx());
    let default_child = scope.branch_child(OpInserter::new_at_block_end(default_body));
    let default = default_block(&default_child);
    let mut out = default.init_mut(scope);
    out.__expand_assign_method(&default_child, default.into_expand(scope));
    default_child.terminate_yield();

    SwitchExpandExpr {
        switch_op,
        cases: Vec::new(),
        children: alloc::vec![default_child],
        out,
    }
}

#[allow(clippy::large_enum_variant)]
pub enum MatchExpand<T: CubeEnum> {
    ComptimeVariant {
        variant: i32,
        runtime_value: T::RuntimeValue,
        matched: bool,
    },
    RuntimeVariant {
        switch_op: SwitchOp,
        cases: Vec<i32>,
        children: Vec<Scope>,
        has_default: bool,
        runtime_value: T::RuntimeValue,
    },
}

impl<T: CubeEnum> MatchExpand<T> {
    pub fn case(
        mut self,
        scope: &Scope,
        value: i32,
        block: impl FnOnce(&Scope, T::RuntimeValue),
    ) -> Self {
        match &mut self {
            Self::RuntimeVariant {
                switch_op,
                cases,
                children,
                runtime_value,
                ..
            } => {
                cases.push(value);
                let body = switch_op.append_case_block(scope.ctx_mut());
                let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
                block(&case_child, (*runtime_value).clone_unchecked());
                case_child.terminate_yield();
                children.push(case_child);
            }
            Self::ComptimeVariant {
                variant,
                runtime_value,
                matched,
            } => {
                if value == *variant {
                    block(scope, (*runtime_value).clone_unchecked());
                    *matched = true;
                }
            }
        }
        self
    }

    pub fn default(mut self, scope: &Scope, block: impl FnOnce(&Scope, T::RuntimeValue)) -> Self {
        match &mut self {
            Self::RuntimeVariant {
                switch_op,
                children,
                runtime_value,
                has_default,
                ..
            } => {
                let body = switch_op.default_block(scope.ctx());
                let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
                block(&case_child, (*runtime_value).clone_unchecked());
                case_child.terminate_yield();
                children.push(case_child);
                *has_default = true;
            }
            Self::ComptimeVariant {
                runtime_value,
                matched,
                ..
            } => {
                if !*matched {
                    block(scope, (*runtime_value).clone_unchecked());
                    *matched = true;
                }
            }
        }
        self
    }

    pub fn finish(self, scope: &Scope) {
        match self {
            MatchExpand::ComptimeVariant { .. } => {}
            MatchExpand::RuntimeVariant {
                switch_op,
                cases,
                children,
                has_default,
                ..
            } => {
                if !has_default {
                    let default_body = switch_op.default_block(scope.ctx());
                    let mut inserter = OpInserter::new_at_block_end(default_body);
                    let unreachable = UnreachableOp::new(scope.ctx_mut());
                    inserter.append_op(scope.ctx(), &unreachable);
                }

                let cases = cases.into_iter().map(|case| {
                    let ty = IntegerType::get(scope.ctx(), 32, Signedness::Unsigned);
                    IntegerAttr::new(ty, APInt::from_i32(case, bw(32)))
                });
                switch_op.set_attr_cases(scope.ctx(), cases);
                scope.register(&switch_op);
                scope.update_flags_after_branch(&children);
            }
        }
    }
}

pub fn match_expand<T: CubeEnum>(
    scope: &Scope,
    value: T,
    discriminant0: i32,
    arm0: impl FnOnce(&Scope, T::RuntimeValue),
) -> MatchExpand<T> {
    let discriminant = value.discriminant();
    match discriminant.constant() {
        Some(const_variant) if const_variant.as_i32() == discriminant0 => {
            let runtime_value = value.runtime_value();
            arm0(scope, runtime_value.clone_unchecked());
            MatchExpand::ComptimeVariant {
                variant: const_variant.as_i32(),
                runtime_value,
                matched: true,
            }
        }
        Some(const_variant) => MatchExpand::ComptimeVariant {
            variant: const_variant.as_i32(),
            runtime_value: value.runtime_value(),
            matched: false,
        },
        None => {
            let discriminant = discriminant.read_value(scope);
            let runtime_value = value.runtime_value();

            let switch_op = SwitchOp::new(scope.ctx_mut(), discriminant);
            let body = switch_op.append_case_block(scope.ctx_mut());
            let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
            arm0(&case_child, runtime_value.clone_unchecked());
            case_child.terminate_yield();

            MatchExpand::RuntimeVariant {
                switch_op,
                cases: alloc::vec![discriminant0],
                children: alloc::vec![case_child],
                has_default: false,
                runtime_value,
            }
        }
    }
}

#[allow(clippy::large_enum_variant)]
pub enum MatchExpandExpr<T: CubeEnum, C: Assign> {
    ComptimeVariant {
        variant: i32,
        runtime_value: T::RuntimeValue,
        out: Option<C>,
        matched: bool,
    },
    RuntimeVariant {
        switch_op: SwitchOp,
        cases: Vec<i32>,
        children: Vec<Scope>,
        has_default: bool,
        out: C,
        runtime_value: T::RuntimeValue,
    },
}

impl<T: CubeEnum, C: Assign> MatchExpandExpr<T, C> {
    pub fn case<R: RuntimeAssign<Expand = C>>(
        mut self,
        scope: &Scope,
        value: i32,
        block: impl FnOnce(&Scope, T::RuntimeValue) -> R,
    ) -> Self {
        match &mut self {
            Self::RuntimeVariant {
                switch_op,
                cases,
                children,
                out,
                runtime_value,
                ..
            } => {
                cases.push(value);
                let body = switch_op.append_case_block(scope.ctx_mut());
                let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
                let ret_val = block(&case_child, (*runtime_value).clone_unchecked());
                out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
                case_child.terminate_yield();
                children.push(case_child);
            }
            Self::ComptimeVariant {
                variant,
                runtime_value,
                out,
                matched,
            } => {
                if value == *variant {
                    *out =
                        Some(block(scope, (*runtime_value).clone_unchecked()).into_expand(scope));
                    *matched = true;
                }
            }
        }
        self
    }

    pub fn default<R: RuntimeAssign<Expand = C>>(
        mut self,
        scope: &Scope,
        block: impl FnOnce(&Scope, T::RuntimeValue) -> R,
    ) -> Self {
        match &mut self {
            Self::RuntimeVariant {
                switch_op,
                children,
                runtime_value,
                out,
                has_default,
                ..
            } => {
                let body = switch_op.default_block(scope.ctx());
                let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
                let ret_val = block(&case_child, (*runtime_value).clone_unchecked());
                out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
                case_child.terminate_yield();
                children.push(case_child);
                *has_default = true;
            }
            Self::ComptimeVariant {
                runtime_value,
                out,
                matched,
                ..
            } => {
                if !*matched {
                    *out =
                        Some(block(scope, (*runtime_value).clone_unchecked()).into_expand(scope));
                    *matched = true;
                }
            }
        }
        self
    }

    pub fn finish(self, scope: &Scope) -> C {
        match self {
            MatchExpandExpr::ComptimeVariant { out, .. } => {
                out.expect("At least one variant should be matched")
            }
            MatchExpandExpr::RuntimeVariant {
                switch_op,
                cases,
                children,
                has_default,
                out,
                ..
            } => {
                if !has_default {
                    let default_body = switch_op.default_block(scope.ctx());
                    let mut inserter = OpInserter::new_at_block_end(default_body);
                    let unreachable = UnreachableOp::new(scope.ctx_mut());
                    inserter.append_op(scope.ctx(), &unreachable);
                }

                let cases = cases.into_iter().map(|case| {
                    let ty = IntegerType::get(scope.ctx(), 32, Signedness::Unsigned);
                    IntegerAttr::new(ty, APInt::from_i32(case, bw(32)))
                });
                switch_op.set_attr_cases(scope.ctx(), cases);
                scope.register(&switch_op);
                scope.update_flags_after_branch(&children);

                out
            }
        }
    }
}

pub fn match_expand_expr<T: CubeEnum, C: RuntimeAssign>(
    scope: &Scope,
    value: T,
    discriminant0: i32,
    arm0: impl FnOnce(&Scope, T::RuntimeValue) -> C,
) -> MatchExpandExpr<T, C::Expand> {
    let discriminant = value.discriminant();
    match discriminant.constant() {
        Some(const_variant) if const_variant.as_i32() == discriminant0 => {
            let runtime_value = value.runtime_value();
            let out = arm0(scope, runtime_value.clone_unchecked());
            MatchExpandExpr::ComptimeVariant {
                variant: const_variant.as_i32(),
                out: Some(out.into_expand(scope)),
                runtime_value,
                matched: true,
            }
        }
        Some(const_variant) => MatchExpandExpr::ComptimeVariant {
            variant: const_variant.as_i32(),
            out: None,
            runtime_value: value.runtime_value(),
            matched: false,
        },
        None => {
            let discriminant = discriminant.read_value(scope);
            let runtime_value = value.runtime_value();

            let switch_op = SwitchOp::new(scope.ctx_mut(), discriminant);
            let body = switch_op.append_case_block(scope.ctx_mut());
            let case_child = scope.branch_child(OpInserter::new_at_block_end(body));
            let ret_val = arm0(&case_child, runtime_value.clone_unchecked());

            let mut out = ret_val.init_mut(scope);
            out.__expand_assign_method(&case_child, ret_val.into_expand(scope));
            case_child.terminate_yield();

            MatchExpandExpr::RuntimeVariant {
                switch_op,
                out,
                cases: alloc::vec![discriminant0],
                children: alloc::vec![case_child],
                runtime_value,
                has_default: false,
            }
        }
    }
}

pub fn continue_expand(scope: &Scope) {
    let inv_continue_flag = scope
        .expand_state()
        .inv_continue_flag
        .expect("Should be in loop");
    let false_ = false.__expand_runtime_method(scope).expand;
    assign::expand_element(scope, false_, inv_continue_flag.into());
    scope.expand_state_mut().may_continue = true;
}

pub fn break_expand(scope: &Scope) {
    let inv_break_flag = scope
        .expand_state()
        .inv_break_flag
        .expect("Should be in loop");
    let false_ = false.__expand_runtime_method(scope).expand;
    assign::expand_element(scope, false_, inv_break_flag.into());
    scope.expand_state_mut().may_break = true;
}

pub fn return_expand(scope: &Scope) {
    let inv_return_flag = scope.expand_state().inv_return_flag;
    if let Some(inv_return_flag) = inv_return_flag {
        let false_ = false.__expand_runtime_method(scope).expand;
        assign::expand_element(scope, false_, inv_return_flag.into());
        scope.expand_state_mut().may_return = true;
    } else {
        panic!("Can't return when no return flag has been set up")
    }
}

/// This is a separate function because otherwise Rust can't infer the generic
pub fn return_with_value_expand<T: CubePrimitive>(
    scope: &Scope,
    value: NativeExpand<T>,
) -> NativeExpand<T> {
    let inv_return_flag = scope.expand_state().inv_return_flag;
    if let Some(inv_return_flag) = inv_return_flag {
        let return_value = scope.expand_state().return_value.expect("Should exist");
        let false_ = false.__expand_runtime_method(scope).expand;

        assign::expand_element(scope, false_, inv_return_flag.into());
        assign::expand_element(scope, value.expand, return_value.into());
        scope.expand_state_mut().may_return = true;
    } else {
        panic!("Can't return when no return flag has been set up")
    }
    value
}

pub fn terminate_expand(scope: &Scope) {
    let inv_terminate_flag = scope.expand_state().inv_terminate_flag;
    if let Some(inv_terminate_flag) = inv_terminate_flag {
        let false_ = false.__expand_runtime_method(scope).expand;
        assign::expand_element(scope, false_, inv_terminate_flag.into());
        scope.expand_state_mut().may_terminate = true;
    } else {
        panic!("Can't terminate when no terminate flag has been set up")
    }
}

#[diagnostic::on_unimplemented(
    message = "Only runtime primitives can be returned. {Self} is not a primitive."
)]
pub trait EarlyReturnable {
    fn init_value(scope: &Scope) -> Option<Value>;
    fn from_value(scope: &Scope, value: Option<Value>) -> Self;
}

impl<T: CubePrimitive> EarlyReturnable for NativeExpand<T> {
    fn init_value(scope: &Scope) -> Option<Value> {
        Some(init_mut_of_type(scope, T::__expand_as_type(scope)).value(scope))
    }
    fn from_value(scope: &Scope, value: Option<Value>) -> Self {
        read_value(scope, value.unwrap()).into()
    }
}
impl EarlyReturnable for () {
    fn init_value(_: &Scope) -> Option<Value> {
        None
    }
    fn from_value(_: &Scope, _: Option<Value>) -> Self {}
}

pub fn setup_early_return<T: EarlyReturnable>(scope: &Scope, body: impl FnOnce(&Scope) -> T) -> T {
    let value = T::init_value(scope);
    let inserter = OpInserter::new(scope.inserter().get_insertion_point());
    let child = scope.inlined_func_child(inserter, value);
    body(&child);

    let current_block = scope.inserter().get_insertion_block(scope.ctx()).unwrap();
    let after_point = op_insertion_point_in_block(
        scope.ctx(),
        child.inserter().get_insertion_point(),
        current_block,
    );
    if child.inserter().get_insertion_block(scope.ctx()) != Some(current_block) {
        child.terminate_yield();
    }

    scope.inserter().set_insertion_point(after_point);
    scope.update_flags_after_early_return(&[child]);
    T::from_value(scope, value)
}

pub mod unreachable_unchecked {
    use super::*;

    pub fn expand(scope: &Scope) {
        scope.register(&UnreachableOp::new(scope.ctx_mut()));
    }
}

/// `WhileOp` builder. Needs to work like this for borrow checker reasons.
/// Rust requires all closures passed to a function to be live at once, so we use a builder pattern
/// and separate the cond and body into separate functions to allow the borrows to work.
pub struct WhileBuilder {
    while_op: WhileOp,
    cond: ExpandValue,
    cond_scope: Scope,
}

impl WhileBuilder {
    pub fn new(scope: &Scope, mut cond: impl FnMut(&Scope) -> NativeExpand<bool>) -> Self {
        let while_op = WhileOp::new(scope.ctx_mut());
        let cond_scope = scope.branch_child(OpInserter::new_at_block_start(
            while_op.before_block(scope.ctx()),
        ));

        WhileBuilder {
            while_op,
            cond: cond(&cond_scope).expand,
            cond_scope,
        }
    }

    pub fn with_body(self, scope: &Scope, mut block: impl FnMut(&Scope)) {
        let Self {
            while_op,
            mut cond,
            cond_scope,
        } = self;

        let body = while_op.after_block(scope.ctx());
        let body = scope.loop_child(OpInserter::new_at_block_end(body));
        block(&body);
        body.terminate_yield();

        let expand_state = *body.expand_state();
        let break_flag = expand_state.inv_break_flag.unwrap().into();
        let return_flag = expand_state.inv_return_flag.map(Into::into);
        let terminate_flag = expand_state.inv_terminate_flag.map(Into::into);

        if expand_state.may_break {
            cond = binary_expand(&cond_scope, cond, break_flag, BoolAndOp::new);
        }
        if expand_state.may_return {
            let return_flag = return_flag.unwrap();
            cond = binary_expand(&cond_scope, cond, return_flag, BoolAndOp::new);
        }
        if expand_state.may_terminate {
            let terminate_flag = terminate_flag.unwrap();
            cond = binary_expand(&cond_scope, cond, terminate_flag, BoolAndOp::new);
        }

        let cond = cond.read_value(&cond_scope);
        cond_scope.register(&ConditionOp::new(scope.ctx_mut(), cond));

        scope.register(&while_op);
        scope.update_flags_after_loop(&[cond_scope, body]);
    }
}

/// register a range loop if it contains no break or return, destructure to while if it does
pub(crate) fn register_range_loop<I: Int>(scope: &Scope, for_op: &RangeLoopOp, body: &Scope) {
    let ctx = scope.ctx_mut();
    let ExpandState {
        may_terminate,
        may_return,
        may_break,
        may_continue,
        inv_terminate_flag,
        inv_return_flag,
        inv_break_flag,
        inv_continue_flag: _,
        return_value: _,
    } = *body.expand_state();
    if !may_break && !may_return && !may_terminate {
        body.terminate_yield();
        scope.register(for_op);
        return;
    }

    let mut rewriter = IRRewriter::<DummyListener>::default();

    let start = for_op.start(ctx);
    let end = for_op.end(ctx);
    let step = for_op.step(ctx);
    let iter_var_old = for_op.iter_var(ctx);

    let iter_var = ExpandValue::from(scope.create_local_mut(iter_var_old.get_type(ctx), None));
    assign::expand_element(scope, start.into(), iter_var);

    // Normally it needs to be predicated on `break;`/`return;`, but not `continue;`. However, this
    // would require additional branching and Rust actually doesn't allow loop counters to escape
    // the loop. So just not predicating it is cheaper, easier, and semantically identical. Yay Rust!
    if may_continue {
        set_inserter_before_terminator(body.inserter(), ctx, for_op.loop_body(ctx));
    }
    assign_binop_expand::<I>(
        body,
        &mut iter_var.into(),
        step.into(),
        I::__expand_native_add,
    );
    body.terminate_yield();

    body.inserter()
        .set_insertion_point_to_block_start(for_op.loop_body(ctx));
    let iter_val = iter_var.read_value(body);
    rewriter.replace_value_uses_with(ctx, iter_var_old, iter_val);
    BasicBlock::remove_argument(for_op.loop_body(ctx), ctx, 0);

    let while_op = WhileOp::new(ctx);

    let cond_scope = scope.branch_child(OpInserter::new_at_block_start(
        while_op.before_block(scope.ctx()),
    ));

    let mut cond = I::__expand_native_lt(&cond_scope, iter_var, end.into());
    if may_break {
        let inv_break_flag = inv_break_flag.unwrap().into();
        cond = binary_expand(&cond_scope, cond, inv_break_flag, BoolAndOp::new);
    }
    if may_return {
        let inv_return_flag = inv_return_flag.unwrap().into();
        cond = binary_expand(&cond_scope, cond, inv_return_flag, BoolAndOp::new);
    }
    if may_terminate {
        let inv_terminate_flag = inv_terminate_flag.unwrap().into();
        cond = binary_expand(&cond_scope, cond, inv_terminate_flag, BoolAndOp::new);
    }

    let cond = cond.read_value(&cond_scope);
    cond_scope.register(&ConditionOp::new(scope.ctx_mut(), cond));

    rewriter.erase_region(ctx, while_op.after_region(ctx));
    Region::move_to_op(for_op.get_region(ctx), while_op.get_operation(), ctx);
    rewriter.erase_operation(ctx, for_op.get_operation());

    scope.register(&while_op);
}

// Don't make this `FnOnce`, it must be executable multiple times
pub fn loop_expand(scope: &Scope, block: impl FnMut(&Scope)) {
    WhileBuilder::new(scope, |scope| true.__expand_runtime_method(scope)).with_body(scope, block);
}