midenc-hir 0.8.1

use alloc::{boxed::Box, format, rc::Rc, vec};

use super::OperationState;
use crate::{
    BlockArgument, BlockRef, BuildableOp, Context, OperationRef, ProgramPoint, RegionRef,
    SourceSpan, Type, Value,
    diagnostics::{LabeledSpan, Report, Severity, miette::diagnostic},
};

/// The [Builder] trait encompasses all of the functionality needed to construct and insert blocks
/// and operations into the IR.
pub trait Builder: Listener {
    fn context(&self) -> &Context;
    fn context_rc(&self) -> Rc<Context>;
    /// Returns the current insertion point of the builder
    fn insertion_point(&self) -> &ProgramPoint;
    /// Clears the current insertion point
    fn clear_insertion_point(&mut self) -> ProgramPoint;
    /// Restores the current insertion point to `ip`
    fn restore_insertion_point(&mut self, ip: ProgramPoint);
    /// Sets the current insertion point to `ip`
    fn set_insertion_point(&mut self, ip: ProgramPoint);

    /// Sets the insertion point before `op`, causing subsequent insertions to be placed before it.
    #[inline]
    fn set_insertion_point_before(&mut self, op: OperationRef) {
        self.set_insertion_point(ProgramPoint::before(op));
    }

    /// Sets the insertion point after `op`, causing subsequent insertions to be placed after it.
    #[inline]
    fn set_insertion_point_after(&mut self, op: OperationRef) {
        self.set_insertion_point(ProgramPoint::after(op));
    }

    /// Sets the insertion point to the node after the specified value is defined.
    ///
    /// If value has a defining operation, this sets the insertion point after that operation, so
    /// that all insertions are placed following the definition.
    ///
    /// Otherwise, a value must be a block argument, so the insertion point is placed at the start
    /// of the block, causing insertions to be placed starting at the front of the block.
    fn set_insertion_point_after_value(&mut self, value: &dyn Value) {
        let pp = if let Some(op) = value.get_defining_op() {
            ProgramPoint::after(op)
        } else {
            let block_argument = value.downcast_ref::<BlockArgument>().unwrap();
            ProgramPoint::at_start_of(block_argument.owner())
        };
        self.set_insertion_point(pp);
    }

    /// Sets the current insertion point to the start of `block`.
    ///
    /// Operations inserted will be placed starting at the beginning of the block.
    #[inline]
    fn set_insertion_point_to_start(&mut self, block: BlockRef) {
        self.set_insertion_point(ProgramPoint::at_start_of(block));
    }

    /// Sets the current insertion point to the end of `block`.
    ///
    /// Operations inserted will be placed starting at the end of the block.
    #[inline]
    fn set_insertion_point_to_end(&mut self, block: BlockRef) {
        self.set_insertion_point(ProgramPoint::at_end_of(block));
    }

    /// Return the block the current insertion point belongs to.
    ///
    /// NOTE: The insertion point is not necessarily at the end of the block.
    ///
    /// Returns `None` if the insertion point is unset, or is pointing at an operation which is
    /// detached from a block.
    fn insertion_block(&self) -> Option<BlockRef> {
        self.insertion_point().block()
    }

    /// Add a new block with `args` arguments, and set the insertion point to the end of it.
    ///
    /// The block is inserted after the provided insertion point `ip`, or at the end of `parent` if
    /// not.
    ///
    /// Panics if `ip` is in a different region than `parent`, or if the position it refers to is no
    /// longer valid.
    fn create_block(&mut self, parent: RegionRef, ip: Option<BlockRef>, args: &[Type]) -> BlockRef {
        let mut block = self.context_rc().create_block_with_params(args.iter().cloned());
        if let Some(at) = ip {
            let region = at.parent().unwrap();
            assert!(
                RegionRef::ptr_eq(&parent, &region),
                "insertion point region differs from 'parent'"
            );

            block.borrow_mut().insert_after(at);
        } else {
            block.borrow_mut().insert_at_end(parent);
        }

        self.notify_block_inserted(block, None, None);

        self.set_insertion_point_to_end(block);

        block
    }

    /// Add a new block with `args` arguments, and set the insertion point to the end of it.
    ///
    /// The block is inserted before `before`.
    fn create_block_before(&mut self, before: BlockRef, args: &[Type]) -> BlockRef {
        let mut block = self.context_rc().create_block_with_params(args.iter().cloned());
        block.borrow_mut().insert_before(before);
        self.notify_block_inserted(block, None, None);
        self.set_insertion_point_to_end(block);
        block
    }

    /// Insert `op` at the current insertion point.
    ///
    /// If the insertion point is inserting after the current operation, then after calling this
    /// function, the insertion point will have been moved to the newly inserted operation. This
    /// ensures that subsequent calls to `insert` will place operations in the block in the same
    /// sequence as they were inserted. The other insertion point placements already have more or
    /// less intuitive behavior, e.g. inserting _before_ the current operation multiple times will
    /// result in operations being placed in the same sequence they were inserted, just before the
    /// current op.
    ///
    /// This function will panic if no insertion point is set.
    fn insert(&mut self, op: OperationRef) {
        let ip = self.insertion_point();

        match *ip {
            ProgramPoint::Block {
                block,
                position: point,
            } => match point {
                crate::Position::Before => op.insert_at_start(block),
                crate::Position::After => op.insert_at_end(block),
            },
            ProgramPoint::Op {
                op: other_op,
                position: point,
                ..
            } => match point {
                crate::Position::Before => op.insert_before(other_op),
                crate::Position::After => {
                    op.insert_after(other_op);
                    self.set_insertion_point(ProgramPoint::after(op));
                }
            },
            ProgramPoint::Invalid => panic!("insertion point is invalid/unset"),
        }
        self.notify_operation_inserted(op, *self.insertion_point());
    }

    /// Create an [super::Operation] from the provided [OperationState]
    fn create_operation(&mut self, state: &mut OperationState) -> Result<OperationRef, Report> {
        let mut op = state.name.alloc_default(self.context_rc());
        op.borrow_mut().set_span(state.span);

        let mut builder = crate::GenericOperationBuilder::new(self, op);

        for prop in op.name().properties() {
            if !state.attrs.iter().any(|p| p.name == prop.name) {
                return Err(Report::from(diagnostic!(
                    severity = Severity::Error,
                    labels = vec![LabeledSpan::at(
                        state.span,
                        format!("missing required property '{}'", prop.name)
                    )],
                    "invalid operation"
                )));
            }
        }

        for attr in state.attrs.drain(..) {
            if state.name.has_property(attr.name) {
                builder.with_property_boxed(attr.name, attr.value)?;
            } else {
                builder.with_attr_boxed(attr.name, attr.value);
            }
        }

        for (i, group) in state.operands.drain(..).enumerate() {
            builder.with_operands_in_group(i, group);
        }

        for successor in state.successors.drain(..) {
            builder.with_pending_successor(successor);
        }

        if !state.regions.is_empty() {
            let mut op = op.borrow_mut();
            let regions = op.regions_mut();
            for region in state.regions.drain(..) {
                regions.push_back(region);
            }
        }

        if !state.results.is_empty() {
            builder.with_results(state.results.drain(..));
        }

        builder.build()
    }
}

pub trait BuilderExt: Builder {
    /// Returns a specialized builder for a concrete [super::Op], `T`, which can be called like a
    /// closure with the arguments required to create an instance of the specified operation.
    ///
    /// # How it works
    ///
    /// The set of arguments which are valid for the specialized builder returned by `create`, are
    /// determined by what implementations of the [BuildableOp] trait exist for `T`. The specific
    /// impl that is chosen will depend on the types of the arguments given to it. Typically, there
    /// should only be one implementation, or if there are multiple, they should not overlap in
    /// ways that may confuse type inference, or you will be forced to specify the full type of the
    /// argument pack.
    ///
    /// This mechanism for constructing ops using arbitrary arguments is essentially a workaround
    /// for the lack of variadic generics in Rust, and isn't quite as nice as what you can acheive
    /// in C++ with varidadic templates and `std::forward` and such, but is close enough so that
    /// the ergonomics are still a significant improvement over the alternative approaches.
    ///
    /// The nice thing about this is that we can generate all of the boilerplate, and hide all of
    /// the sensitive/unsafe parts of initializing operations. Alternative approaches require
    /// exposing more unsafe APIs for use by builders, whereas this approach can conceal those
    /// details within this crate.
    ///
    /// ## Example
    ///
    /// ```text,ignore
    /// // Get an OpBuilder
    /// let builder = context.builder();
    /// // Obtain a builder for AddOp
    /// let add_builder = builder.create::<AddOp, _>(span);
    /// // Consume the builder by creating the op with the given arguments
    /// let add = add_builder(lhs, rhs, Overflow::Wrapping).expect("invalid add op");
    /// ```
    ///
    /// Or, simplified/collapsed:
    ///
    /// ```text,ignore
    /// let builder = context.builder();
    /// let add = builder.create::<AddOp, _>(span)(lhs, rhs, Overflow::Wrapping)
    ///     .expect("invalid add op");
    /// ```
    #[inline(always)]
    fn create<T, Args>(&mut self, span: SourceSpan) -> <T as BuildableOp<Args>>::Builder<'_, Self>
    where
        Args: core::marker::Tuple,
        T: BuildableOp<Args>,
    {
        <T as BuildableOp<Args>>::builder(self, span)
    }
}

impl<B: ?Sized + Builder> BuilderExt for B {}

pub struct OpBuilder<L = NoopBuilderListener> {
    context: Rc<Context>,
    listener: Option<L>,
    ip: ProgramPoint,
}

impl OpBuilder {
    pub fn new(context: Rc<Context>) -> Self {
        Self {
            context,
            listener: None,
            ip: ProgramPoint::default(),
        }
    }
}

impl<L: Listener> OpBuilder<L> {
    /// Sets the listener of this builder to `listener`
    pub fn with_listener<L2>(self, listener: L2) -> OpBuilder<L2>
    where
        L2: Listener,
    {
        OpBuilder {
            context: self.context,
            listener: Some(listener),
            ip: self.ip,
        }
    }

    pub fn listener(&self) -> Option<&L> {
        self.listener.as_ref()
    }

    #[inline]
    pub fn into_parts(self) -> (Rc<Context>, Option<L>, ProgramPoint) {
        (self.context, self.listener, self.ip)
    }
}

impl<L: Listener> Listener for OpBuilder<L> {
    fn kind(&self) -> ListenerType {
        self.listener.as_ref().map(|l| l.kind()).unwrap_or(ListenerType::Builder)
    }

    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {
        if let Some(listener) = self.listener.as_ref() {
            listener.notify_operation_inserted(op, prev);
        }
    }

    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
        if let Some(listener) = self.listener.as_ref() {
            listener.notify_block_inserted(block, prev, ip);
        }
    }
}

impl<L: Listener> Builder for OpBuilder<L> {
    #[inline(always)]
    fn context(&self) -> &Context {
        self.context.as_ref()
    }

    #[inline(always)]
    fn context_rc(&self) -> Rc<Context> {
        self.context.clone()
    }

    #[inline(always)]
    fn insertion_point(&self) -> &ProgramPoint {
        &self.ip
    }

    #[inline]
    fn clear_insertion_point(&mut self) -> ProgramPoint {
        let ip = self.ip;
        self.ip = ProgramPoint::Invalid;
        ip
    }

    #[inline]
    fn restore_insertion_point(&mut self, ip: ProgramPoint) {
        self.ip = ip;
    }

    #[inline(always)]
    fn set_insertion_point(&mut self, ip: ProgramPoint) {
        self.ip = ip;
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ListenerType {
    Builder,
    Rewriter,
}

#[allow(unused_variables)]
pub trait Listener {
    fn kind(&self) -> ListenerType;
    /// Notify the listener that the specified operation was inserted.
    ///
    /// * If the operation was moved, then `prev` is the previous location of the op
    /// * If the operation was unlinked before it was inserted, then `prev` is `Invalid`
    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {}
    /// Notify the listener that the specified block was inserted.
    ///
    /// * If the block was created, but not inserted, then `prev` and `ip` are `None`
    /// * If the block was unlinked before it was inserted, then `prev` is `None` and `ip` is the
    ///   location where it was inserted.
    /// * If the block was moved, then `prev` is the previous region it was inserted into, and `ip`
    ///   is the location where it was inserted
    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
    }
}

impl<L: Listener> Listener for Option<L> {
    fn kind(&self) -> ListenerType {
        ListenerType::Builder
    }

    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
        if let Some(listener) = self.as_ref() {
            listener.notify_block_inserted(block, prev, ip);
        }
    }

    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {
        if let Some(listener) = self.as_ref() {
            listener.notify_operation_inserted(op, prev);
        }
    }
}

impl<L: ?Sized + Listener> Listener for Box<L> {
    #[inline]
    fn kind(&self) -> ListenerType {
        (**self).kind()
    }

    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {
        (**self).notify_operation_inserted(op, prev)
    }

    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
        (**self).notify_block_inserted(block, prev, ip)
    }
}

impl<L: ?Sized + Listener> Listener for Rc<L> {
    #[inline]
    fn kind(&self) -> ListenerType {
        (**self).kind()
    }

    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {
        (**self).notify_operation_inserted(op, prev)
    }

    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
        (**self).notify_block_inserted(block, prev, ip)
    }
}

/// A listener of kind `Builder` that does nothing
pub struct NoopBuilderListener;
impl Listener for NoopBuilderListener {
    #[inline]
    fn kind(&self) -> ListenerType {
        ListenerType::Builder
    }
}

/// This is used to allow [InsertionGuard] to be agnostic about the type of builder/rewriter it
/// wraps, while still performing the necessary insertion point restoration on drop. Without this,
/// we would be required to specify a `B: Builder` bound on the definition of [InsertionGuard].
#[doc(hidden)]
#[allow(unused_variables)]
trait RestoreInsertionPointOnDrop {
    fn restore_insertion_point_on_drop(&mut self, ip: ProgramPoint);
}
impl<B: ?Sized> RestoreInsertionPointOnDrop for InsertionGuard<'_, B> {
    #[inline(always)]
    default fn restore_insertion_point_on_drop(&mut self, _ip: ProgramPoint) {}
}
impl<B: ?Sized + Builder> RestoreInsertionPointOnDrop for InsertionGuard<'_, B> {
    fn restore_insertion_point_on_drop(&mut self, ip: ProgramPoint) {
        self.builder.restore_insertion_point(ip);
    }
}

pub struct InsertionGuard<'a, B: ?Sized> {
    builder: &'a mut B,
    ip: ProgramPoint,
}
impl<'a, B> InsertionGuard<'a, B>
where
    B: ?Sized + Builder,
{
    #[allow(unused)]
    pub fn new(builder: &'a mut B) -> Self {
        let ip = *builder.insertion_point();
        Self { builder, ip }
    }
}
impl<B: ?Sized> core::ops::Deref for InsertionGuard<'_, B> {
    type Target = B;

    #[inline(always)]
    fn deref(&self) -> &Self::Target {
        self.builder
    }
}
impl<B: ?Sized> core::ops::DerefMut for InsertionGuard<'_, B> {
    #[inline(always)]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.builder
    }
}
impl<B: ?Sized> Drop for InsertionGuard<'_, B> {
    fn drop(&mut self) {
        let ip = self.ip;
        self.restore_insertion_point_on_drop(ip);
    }
}
impl<B: ?Sized + Listener> Listener for InsertionGuard<'_, B> {
    fn kind(&self) -> ListenerType {
        self.builder.kind()
    }

    fn notify_block_inserted(
        &self,
        block: BlockRef,
        prev: Option<RegionRef>,
        ip: Option<BlockRef>,
    ) {
        self.builder.notify_block_inserted(block, prev, ip);
    }

    fn notify_operation_inserted(&self, op: OperationRef, prev: ProgramPoint) {
        self.builder.notify_operation_inserted(op, prev);
    }
}
impl<B: ?Sized + Builder> Builder for InsertionGuard<'_, B> {
    fn context(&self) -> &Context {
        self.builder.context()
    }

    fn context_rc(&self) -> Rc<Context> {
        self.builder.context_rc()
    }

    fn insertion_point(&self) -> &ProgramPoint {
        self.builder.insertion_point()
    }

    fn clear_insertion_point(&mut self) -> ProgramPoint {
        self.builder.clear_insertion_point()
    }

    fn restore_insertion_point(&mut self, ip: ProgramPoint) {
        self.builder.restore_insertion_point(ip);
    }

    fn set_insertion_point(&mut self, ip: ProgramPoint) {
        self.builder.set_insertion_point(ip);
    }

    fn set_insertion_point_before(&mut self, op: OperationRef) {
        self.builder.set_insertion_point_before(op);
    }

    fn set_insertion_point_after(&mut self, op: OperationRef) {
        self.builder.set_insertion_point_after(op);
    }

    fn set_insertion_point_after_value(&mut self, value: &dyn Value) {
        self.builder.set_insertion_point_after_value(value);
    }

    fn set_insertion_point_to_start(&mut self, block: BlockRef) {
        self.builder.set_insertion_point_to_start(block);
    }

    fn set_insertion_point_to_end(&mut self, block: BlockRef) {
        self.builder.set_insertion_point_to_end(block);
    }

    fn insertion_block(&self) -> Option<BlockRef> {
        self.builder.insertion_block()
    }

    fn create_block(&mut self, parent: RegionRef, ip: Option<BlockRef>, args: &[Type]) -> BlockRef {
        self.builder.create_block(parent, ip, args)
    }

    fn create_block_before(&mut self, before: BlockRef, args: &[Type]) -> BlockRef {
        self.builder.create_block_before(before, args)
    }

    fn insert(&mut self, op: OperationRef) {
        self.builder.insert(op);
    }
}