midenc-hir 0.8.0

use alloc::{rc::Rc, sync::Arc, vec::Vec};
use core::{
    cell::{Cell, RefCell},
    mem::MaybeUninit,
};

use blink_alloc::Blink;
use midenc_session::Session;
use traits::BranchOpInterface;

use super::{traits::BuildableTypeConstraint, *};
use crate::{
    AttributeRef, AttributeRegistration, FxHashMap,
    attributes::{AttributeName, DerivableTypeAttribute, Marker},
    constants::{ConstantData, ConstantId, ConstantPool},
};

/// Represents the shared state of the IR, used during a compilation session.
///
/// The primary purpose(s) of the context are:
///
/// * Provide storage/memory for all allocated IR entities for the lifetime of the session.
/// * Provide unique value and block identifiers for printing the IR
/// * Provide a uniqued constant pool
/// * Provide configuration used during compilation
///
/// # Safety
///
/// The [Context] _must_ live as long as any reference to an IR entity may be dereferenced.
pub struct Context {
    session: Rc<Session>,
    allocator: Rc<Blink>,
    registered_dialects: RefCell<FxHashMap<interner::Symbol, Rc<dyn Dialect>>>,
    constants: RefCell<ConstantPool>,
    type_cache: RefCell<FxHashMap<core::any::TypeId, Arc<Type>>>,
    uniqued_attr_cache: RefCell<FxHashMap<AttributeName, AttributeRef>>,
    next_block_id: Cell<u32>,
    next_value_id: Cell<u32>,
}

impl Default for Context {
    fn default() -> Self {
        use alloc::sync::Arc;

        use midenc_session::diagnostics::DefaultSourceManager;

        let target_dir = std::env::current_dir().unwrap();
        let options = midenc_session::Options::default();
        let source_manager = Arc::new(DefaultSourceManager::default());
        let session =
            Rc::new(Session::new([], None, None, target_dir, options, None, source_manager));
        Self::new(session)
    }
}

impl Context {
    /// Create a new [Context] for the given [Session]
    pub fn new(session: Rc<Session>) -> Self {
        use hashbrown::hash_map::Entry;

        let mut registered_dialects = FxHashMap::default();
        for dialect in inventory::iter::<DialectRegistrationInfo>() {
            let namespace = interner::Symbol::intern(dialect.namespace());
            match registered_dialects.entry(namespace) {
                Entry::Vacant(entry) => {
                    entry.insert(dialect.create());
                }
                Entry::Occupied(prev) => {
                    panic!(
                        "conflicting dialect registrations for namespace '{namespace}': {} \
                         conflicts with previous registration of {}",
                        dialect.type_name(),
                        prev.key()
                    );
                }
            }
        }

        let allocator = Rc::new(Blink::new());
        Self {
            session,
            allocator,
            registered_dialects: RefCell::new(registered_dialects),
            constants: Default::default(),
            type_cache: Default::default(),
            uniqued_attr_cache: Default::default(),
            next_block_id: Cell::new(0),
            next_value_id: Cell::new(0),
        }
    }

    #[inline]
    pub fn session(&self) -> &Session {
        &self.session
    }

    #[inline]
    pub fn session_rc(&self) -> Rc<Session> {
        self.session.clone()
    }

    #[inline]
    pub fn diagnostics(&self) -> &::midenc_session::DiagnosticsHandler {
        &self.session.diagnostics
    }

    pub fn registered_dialects(
        &self,
    ) -> core::cell::Ref<'_, FxHashMap<interner::Symbol, Rc<dyn Dialect>>> {
        self.registered_dialects.borrow()
    }

    pub fn get_registered_dialect(&self, dialect: impl Into<interner::Symbol>) -> Rc<dyn Dialect> {
        let dialect = dialect.into();
        self.registered_dialects.borrow()[&dialect].clone()
    }

    pub fn get_registered_name<T>(&self) -> OperationName
    where
        T: OpRegistration,
    {
        self.get_registered_dialect(T::dialect_name()).expect_registered_name::<T>()
    }

    pub fn get_registered_attribute_name<T>(&self) -> AttributeName
    where
        T: AttributeRegistration,
    {
        self.get_registered_dialect(T::dialect_name())
            .expect_registered_attribute_name::<T>()
    }

    pub fn get_or_register_dialect<T>(&self) -> Rc<dyn Dialect>
    where
        T: DialectRegistration,
    {
        let dialect_name = <T as DialectRegistration>::NAMESPACE.into();
        if let Some(dialect) = self.registered_dialects.borrow().get(&dialect_name).cloned() {
            return dialect;
        }

        let dialect = super::dialect::dialect_init::<T>();
        self.registered_dialects.borrow_mut().insert(dialect_name, Rc::clone(&dialect));
        dialect
    }

    pub fn get_marker<T>(self: &Rc<Context>) -> UnsafeIntrusiveEntityRef<T>
    where
        T: AttributeRegistration<Value = ()>,
    {
        <T as UniquedAttribute>::get_or_create_uniqued_attribute(self, ())
    }

    pub fn create_attribute<T, V>(self: &Rc<Context>, value: V) -> UnsafeIntrusiveEntityRef<T>
    where
        T: AttributeRegistration,
        <T as AttributeRegistration>::Value: From<V>,
    {
        <T as DerivableTypeAttribute>::create(self, value)
    }

    pub fn create_attribute_with_type<T, V>(
        self: &Rc<Context>,
        value: V,
        ty: Type,
    ) -> UnsafeIntrusiveEntityRef<T>
    where
        T: AttributeRegistration,
        <T as AttributeRegistration>::Value: From<V>,
    {
        <T as AttributeRegistration>::create(self, value, ty)
    }

    pub fn create_constant(&self, data: impl Into<ConstantData>) -> ConstantId {
        let mut constants = self.constants.borrow_mut();
        constants.insert(data.into())
    }

    pub fn get_constant(&self, id: ConstantId) -> Arc<ConstantData> {
        self.constants.borrow().get(id)
    }

    pub fn get_constant_size_in_bytes(&self, id: ConstantId) -> usize {
        self.constants.borrow().get_by_ref(id).len()
    }

    pub fn get_cached_type<T: BuildableTypeConstraint>(&self) -> Option<Arc<Type>> {
        self.type_cache.borrow().get(&core::any::TypeId::of::<T>()).cloned()
    }

    pub fn get_or_insert_type<T: BuildableTypeConstraint>(&self) -> Arc<Type> {
        match self.get_cached_type::<T>() {
            Some(ty) => ty,
            None => {
                let ty = Arc::new(<T as BuildableTypeConstraint>::build(self));
                let mut types = self.type_cache.borrow_mut();
                types.insert(core::any::TypeId::of::<T>(), Arc::clone(&ty));
                ty
            }
        }
    }

    /// Get a new [OpBuilder] for this context
    pub fn builder(self: Rc<Self>) -> OpBuilder {
        OpBuilder::new(Rc::clone(&self))
    }

    /// Create a new, detached and empty [Region]
    pub fn create_region(&self) -> RegionRef {
        self.alloc_tracked(Region::default())
    }

    /// Create a new, detached and empty [Block] with no parameters
    pub fn create_block(self: &Rc<Self>) -> BlockRef {
        let id = self.alloc_block_id();
        let block = Block::new(Rc::clone(self), id);
        self.alloc_tracked(block)
    }

    /// Create a new, detached and empty [Block], with parameters corresponding to the given types
    pub fn create_block_with_params<I>(self: &Rc<Self>, tys: I) -> BlockRef
    where
        I: IntoIterator<Item = Type>,
    {
        let mut block = Rc::clone(self).create_block();
        let owner = block;
        let args = tys.into_iter().enumerate().map(|(index, ty)| {
            let id = self.alloc_value_id();
            let arg = BlockArgument::new(
                SourceSpan::default(),
                id,
                ty,
                owner,
                index.try_into().expect("too many block arguments"),
            );
            self.alloc(arg)
        });
        block.borrow_mut().arguments_mut().extend(args);
        block
    }

    /// Append a new [BlockArgument] to `block`, with the given type and source location
    ///
    /// Returns the block argument as a `dyn Value` reference
    pub fn append_block_argument(
        &self,
        mut block: BlockRef,
        ty: Type,
        span: SourceSpan,
    ) -> ValueRef {
        let next_index = block.borrow().num_arguments();
        let id = self.alloc_value_id();
        let arg = BlockArgument::new(
            span,
            id,
            ty,
            block,
            next_index.try_into().expect("too many block arguments"),
        );
        let arg = self.alloc(arg);
        block.borrow_mut().arguments_mut().push(arg);
        arg.upcast()
    }

    /// Create a new [OpOperand] with the given value, owner, and index.
    ///
    /// NOTE: This inserts the operand as a user of `value`, but does _not_ add the operand to
    /// `owner`'s operand storage, the caller is expected to do that. This makes this function a
    /// more useful primitive.
    pub fn make_operand(&self, mut value: ValueRef, owner: OperationRef, index: u8) -> OpOperand {
        let op_operand = self.alloc_tracked(OpOperandImpl::new(value, owner, index));
        let mut value = value.borrow_mut();
        value.insert_use(op_operand);
        op_operand
    }

    /// Create a new [BlockOperand] with the given block, owner, and index.
    ///
    /// NOTE: This inserts the block operand as a user of `block`, but does _not_ add the block
    /// operand to `owner`'s successor storage, the caller is expected to do that. This makes this
    /// function a more useful primitive.
    pub fn make_block_operand(
        &self,
        mut block: BlockRef,
        owner: OperationRef,
        index: u8,
    ) -> BlockOperandRef {
        let block_operand = self.alloc_tracked(BlockOperand::new(owner, index));
        let mut block = block.borrow_mut();
        block.insert_use(block_operand);
        block_operand
    }

    /// Create a new [OpResult] with the given type, owner, and index
    ///
    /// NOTE: This does not attach the result to the operation, it is expected that the caller will
    /// do so.
    pub fn make_result(
        &self,
        span: SourceSpan,
        ty: Type,
        owner: OperationRef,
        index: u8,
    ) -> OpResultRef {
        let id = self.alloc_value_id();
        self.alloc(OpResult::new(span, id, ty, owner, index))
    }

    /// Create a new [OpResult] named `name`, with the given type, owner, and index
    ///
    /// NOTE: This does not attach the result to the operation, it is expected that the caller will
    /// do so.
    pub fn make_named_result(
        &self,
        span: SourceSpan,
        ty: Type,
        owner: OperationRef,
        name: interner::Symbol,
        index: u8,
    ) -> OpResultRef {
        let id = ValueId::from_symbol(name).with_result_index(index);
        self.alloc(OpResult::new(span, id, ty, owner, index))
    }

    /// Appends `value` as an argument to the `branch_inst` instruction arguments list if the
    /// destination block of the `branch_inst` is `dest`.
    ///
    /// NOTE: Panics if `branch_inst` is not a branch instruction.
    pub fn append_branch_destination_argument(
        &self,
        mut branch_inst: OperationRef,
        dest: BlockRef,
        value: ValueRef,
    ) {
        let mut borrow = branch_inst.borrow_mut();
        let op = borrow.as_mut().as_operation_mut();
        assert!(
            op.as_trait::<dyn BranchOpInterface>().is_some(),
            "expected branch instruction, got {branch_inst:?}"
        );
        let dest_operand_groups: Vec<usize> = op
            .successors()
            .iter()
            .filter(|succ| succ.block.borrow().successor() == dest)
            .map(|succ| succ.operand_group as usize)
            .collect();
        for dest_group in dest_operand_groups {
            let current_dest_operands_len = op.operands.group(dest_group).len();
            let operand = self.make_operand(
                value,
                op.as_operation_ref(),
                (current_dest_operands_len + 1) as u8,
            );
            op.operands_mut().extend_group(dest_group, [operand]);
        }
    }

    /// Allocate a new uninitialized entity of type `T`
    ///
    /// In general, you can probably prefer [Context::alloc] instead, but for use cases where you
    /// need to allocate the space for `T` first, and then perform initialization, this can be
    /// used.
    pub fn alloc_uninit<T: 'static>(&self) -> UnsafeEntityRef<MaybeUninit<T>> {
        UnsafeEntityRef::new_uninit(&self.allocator)
    }

    /// Allocate a new uninitialized entity of type `T`, which needs to be tracked in an intrusive
    /// doubly-linked list.
    ///
    /// In general, you can probably prefer [Context::alloc_tracked] instead, but for use cases
    /// where you need to allocate the space for `T` first, and then perform initialization,
    /// this can be used.
    pub fn alloc_uninit_tracked<T: 'static>(&self) -> UnsafeIntrusiveEntityRef<MaybeUninit<T>> {
        UnsafeIntrusiveEntityRef::<T>::new_uninit_with_metadata(Default::default(), &self.allocator)
    }

    /// Allocate a new `UnsafeEntityRef<T>`.
    ///
    /// [UnsafeEntityRef] is a smart-pointer type for IR entities, which behaves like a ref-counted
    /// pointer with dynamically-checked borrow checking rules. It is designed to play well with
    /// entities allocated from a [Context], and with the somewhat cyclical nature of the IR.
    pub fn alloc<T: 'static>(&self, value: T) -> UnsafeEntityRef<T> {
        UnsafeEntityRef::new(value, &self.allocator)
    }

    /// Allocate a new `UnsafeIntrusiveEntityRef<T>`.
    ///
    /// [UnsafeIntrusiveEntityRef] is like [UnsafeEntityRef], except that it is specially designed
    /// for entities which are meant to be tracked in intrusive linked lists. For example, the
    /// blocks in a region, or the ops in a block. It does this without requiring the entity to know
    /// about the link at all, while still making it possible to access the link from the entity.
    pub fn alloc_tracked<T: 'static>(&self, value: T) -> UnsafeIntrusiveEntityRef<T> {
        UnsafeIntrusiveEntityRef::new_with_metadata(value, Default::default(), &self.allocator)
    }

    /// Allocate a new `UnsafeIntrusiveMapEntityRef<T>`.
    ///
    /// [UnsafeIntrusiveMapEntityRef] is like [UnsafeEntityRef], except that it is specially
    /// designed for entities which are meant to be tracked in intrusive red/black trees. For
    /// example, the attributes of an op. It does this without requiring the entity to know about
    /// the link at all, while still making it possible to access the link from the entity.
    pub fn alloc_map_item<T: EntityMapItem + 'static>(
        &self,
        value: T,
    ) -> UnsafeIntrusiveMapEntityRef<T> {
        UnsafeIntrusiveMapEntityRef::new_with_metadata(value, Default::default(), &self.allocator)
    }

    fn alloc_block_id(&self) -> BlockId {
        let id = self.next_block_id.get();
        self.next_block_id.set(id + 1);
        BlockId::from_u32(id)
    }

    fn alloc_value_id(&self) -> ValueId {
        let id = self.next_value_id.get();
        self.next_value_id.set(id + 1);
        ValueId::from_u32(id)
    }
}

trait UniquedAttribute: AttributeRegistration {
    fn get_or_create_uniqued_attribute(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
    ) -> UnsafeIntrusiveEntityRef<Self>;
    #[allow(unused)]
    fn get_or_create_uniqued_attribute_with_type(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
        ty: Type,
    ) -> UnsafeIntrusiveEntityRef<Self>;
}

impl<T: AttributeRegistration> UniquedAttribute for T {
    default fn get_or_create_uniqued_attribute(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
    ) -> UnsafeIntrusiveEntityRef<Self> {
        context.create_attribute::<T, _>(value)
    }

    default fn get_or_create_uniqued_attribute_with_type(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
        ty: Type,
    ) -> UnsafeIntrusiveEntityRef<Self> {
        context.create_attribute_with_type::<T, _>(value, ty)
    }
}

impl<T: AttributeRegistration + Marker> UniquedAttribute for T {
    fn get_or_create_uniqued_attribute(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
    ) -> UnsafeIntrusiveEntityRef<Self> {
        use hashbrown::hash_map::Entry;
        let name = context.get_registered_attribute_name::<T>();
        let mut cache = context.uniqued_attr_cache.borrow_mut();
        match cache.entry(name) {
            Entry::Vacant(entry) => {
                let attr = context.create_attribute::<T, _>(value);
                entry.insert(attr as AttributeRef);
                attr
            }
            Entry::Occupied(entry) => entry.get().try_downcast_attr().unwrap(),
        }
    }

    fn get_or_create_uniqued_attribute_with_type(
        context: &Rc<Context>,
        value: <Self as AttributeRegistration>::Value,
        ty: Type,
    ) -> UnsafeIntrusiveEntityRef<Self> {
        use hashbrown::hash_map::Entry;
        let name = context.get_registered_attribute_name::<T>();
        let mut cache = context.uniqued_attr_cache.borrow_mut();
        match cache.entry(name) {
            Entry::Vacant(entry) => {
                let attr = context.create_attribute_with_type::<T, _>(value, ty);
                entry.insert(attr as AttributeRef);
                attr
            }
            Entry::Occupied(entry) => entry.get().try_downcast_attr().unwrap(),
        }
    }
}