Add Op registration into symbol table inside OpBuilder #546

hir/src/dialects/builtin/builders/component.rs

@@ -32,16 +32,6 @@ impl ComponentBuilder {
 
     pub fn define_module(&mut self, name: Ident) -> Result<ModuleRef, Report> {
         let module_ref = self.builder.create_module(name)?;
-        let is_new = self
-            .component
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(module_ref, crate::ProgramPoint::Invalid);
-        assert!(
-            is_new,
-            "module with the name {name} already exists in component {}",
-            self.component.borrow().name()
-        );
         Ok(module_ref)
     }
 
@@ -67,12 +57,6 @@ impl ComponentBuilder {
         signature: Signature,
     ) -> Result<FunctionRef, Report> {
         let function_ref = self.builder.create_function(name, signature)?;
-        let is_new = self
-            .component
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(function_ref, crate::ProgramPoint::Invalid);
-        assert!(is_new, "function with the name {name} already exists");
         Ok(function_ref)
     }
 }

hir/src/dialects/builtin/builders/module.rs

@@ -52,12 +52,6 @@ impl ModuleBuilder {
         signature: Signature,
     ) -> Result<FunctionRef, Report> {
         let function_ref = self.builder.create_function(name, signature)?;
-        let is_new = self
-            .module
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(function_ref, crate::ProgramPoint::Invalid);
-        assert!(is_new, "function with the name {name} already exists");
         Ok(function_ref)
     }
 
@@ -72,12 +66,6 @@ impl ModuleBuilder {
         ty: Type,
     ) -> Result<UnsafeIntrusiveEntityRef<GlobalVariable>, Report> {
         let global_var_ref = self.builder.create_global_variable(name, visibility, ty)?;
-        let is_new = self
-            .module
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(global_var_ref, crate::ProgramPoint::Invalid);
-        assert!(is_new, "global variable with the name {name} already exists");
         Ok(global_var_ref)
     }
 
@@ -119,12 +107,6 @@ impl ModuleBuilder {
     pub fn declare_module(&mut self, name: Ident) -> Result<ModuleRef, Report> {
         let builder = PrimModuleBuilder::new(&mut self.builder, name.span());
         let module_ref = builder(name)?;
-        let is_new = self
-            .module
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(module_ref, crate::ProgramPoint::Invalid);
-        assert!(is_new, "module with the name {name} already exists in world",);
         Ok(module_ref)
     }
 

hir/src/dialects/builtin/builders/world.rs

@@ -43,20 +43,6 @@ impl WorldBuilder {
     ) -> Result<ComponentRef, Report> {
         let builder = PrimComponentBuilder::new(&mut self.builder, name.span());
         let component_ref = builder(ns, name, ver.clone())?;
-        let is_new = self
-            .world
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(component_ref, crate::ProgramPoint::Invalid);
-        assert!(
-            is_new,
-            "component {} already exists in world",
-            ComponentId {
-                namespace: ns.name,
-                name: name.name,
-                version: ver
-            }
-        );
         Ok(component_ref)
     }
 
@@ -73,12 +59,6 @@ impl WorldBuilder {
     pub fn declare_module(&mut self, name: Ident) -> Result<ModuleRef, Report> {
         let builder = PrimModuleBuilder::new(&mut self.builder, name.span());
         let module_ref = builder(name)?;
-        let is_new = self
-            .world
-            .borrow_mut()
-            .symbol_manager_mut()
-            .insert_new(module_ref, crate::ProgramPoint::Invalid);
-        assert!(is_new, "module with the name {name} already exists in world",);
         Ok(module_ref)
     }
 

hir/src/ir/builder.rs

@@ -132,17 +132,17 @@ pub trait Builder: Listener {
                 block,
                 position: point,
             } => match point {
-                crate::Position::Before => op.borrow_mut().insert_at_start(block),
-                crate::Position::After => op.borrow_mut().insert_at_end(block),
+                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.borrow_mut().insert_before(other_op),
+                crate::Position::Before => op.insert_before(other_op),
                 crate::Position::After => {
-                    op.borrow_mut().insert_after(other_op);
+                    op.insert_after(other_op);
                     self.set_insertion_point(ProgramPoint::after(op));
                 }
             },

hir/src/ir/operation.rs

@@ -180,7 +180,30 @@ impl EntityWithParent for Operation {
     type Parent = Block;
 }
 impl EntityListItem for Operation {
-    fn on_inserted(this: UnsafeIntrusiveEntityRef<Self>, _cursor: &mut EntityCursorMut<'_, Self>) {
+    fn on_inserted(this: OperationRef, _cursor: &mut EntityCursorMut<'_, Self>) {
+        let parent = this.nearest_symbol_table();
+        if let Some(mut parent) = parent {
+            // NOTE: We use OperationName, instead of the Operation itself, to avoid borrowing.
+            if this.name().implements::<dyn Symbol>()
+                && parent.name().implements::<dyn SymbolTable>()
+            {
+                let mut symbol_table = parent.borrow_mut();
+                let sym_manager = symbol_table.as_trait_mut::<dyn SymbolTable>().expect(
+                    "Could not cast parent operation {parent.name()} as SymbolTable, even though \
+                     it implements said trait",
+                );
+                let mut sym_manager = sym_manager.symbol_manager_mut();
+
+                let symbol_ref = this.borrow().as_symbol_ref().expect(
+                    "Could not cast operation {this.name()} as Symbol, even though it implements \
+                     said trait",
+                );
+
+                let is_new = sym_manager.insert_new(symbol_ref, ProgramPoint::Invalid);
+                assert!(is_new, "{} already exists in {}", this.name(), parent.name());
+            };
+        }
+
         let order_offset = core::mem::offset_of!(Operation, order);
         unsafe {
             let ptr = UnsafeIntrusiveEntityRef::as_ptr(&this);
@@ -189,15 +212,35 @@ impl EntityListItem for Operation {
         }
     }
 
-    fn on_transfer(
-        _this: UnsafeIntrusiveEntityRef<Self>,
-        _from: &mut EntityList<Self>,
-        to: &mut EntityList<Self>,
-    ) {
+    fn on_transfer(_this: OperationRef, _from: &mut EntityList<Self>, to: &mut EntityList<Self>) {
         // Invalidate the ordering of the new parent block
         let mut to = to.parent();
         to.borrow_mut().invalidate_op_order();
     }
+
+    fn on_removed(this: OperationRef, _list: &mut EntityCursorMut<'_, Self>) {
+        let parent = this.nearest_symbol_table();
+        if let Some(mut parent) = parent {
+            // NOTE: We use OperationName, instead of the Operation itself, to avoid borrowing.
+            if this.name().implements::<dyn Symbol>()
+                && parent.name().implements::<dyn SymbolTable>()
+            {
+                let mut symbol_table = parent.borrow_mut();
+                let sym_manager = symbol_table.as_trait_mut::<dyn SymbolTable>().expect(
+                    "Could not cast parent operation {parent.name()} as SymbolTable, even though \
+                     it implements said trait",
+                );
+                let mut sym_manager = sym_manager.symbol_manager_mut();
+
+                let symbol_ref = this.borrow().as_symbol_ref().expect(
+                    "Could not cast operation {this.name()} as Symbol, even though it implements \
+                     said trait",
+                );
+
+                sym_manager.remove(symbol_ref);
+            };
+        }
+    }
 }
 
 impl EntityParent<Region> for Operation {
@@ -227,6 +270,59 @@ impl Operation {
     }
 }
 
+/// Insertion
+impl OperationRef {
+    pub fn insert_at_start(&self, mut block: BlockRef) {
+        assert!(
+            self.parent().is_none(),
+            "cannot insert operation that is already attached to another block"
+        );
+        {
+            let mut block = block.borrow_mut();
+            block.body_mut().push_front(*self);
+        }
+    }
+
+    pub fn insert_at_end(&self, mut block: BlockRef) {
+        assert!(
+            self.parent().is_none(),
+            "cannot insert operation that is already attached to another block"
+        );
+        {
+            let mut block = block.borrow_mut();
+            block.body_mut().push_back(*self);
+        }
+    }
+
+    pub fn insert_before(&mut self, before: OperationRef) {
+        assert!(
+            self.parent().is_none(),
+            "cannot insert operation that is already attached to another block"
+        );
+        let mut block = before.parent().expect("'before' block is not attached to a block");
+        {
+            let mut block = block.borrow_mut();
+            let block_body = block.body_mut();
+            let mut cursor = unsafe { block_body.cursor_mut_from_ptr(before) };
+            cursor.insert_before(*self);
+        }
+    }
+
+    pub fn insert_after(&mut self, after: OperationRef) {
+        assert!(
+            self.parent().is_none(),
+            "cannot insert operation that is already attached to another block"
+        );
+        let mut block = after.parent().expect("'after' block is not attached to a block");
+        {
+            let mut block = block.borrow_mut();
+            let block_body = block.body_mut();
+            let mut cursor = unsafe { block_body.cursor_mut_from_ptr(after) };
+            cursor.insert_after(*self);
+        }
+    }
+}
+
 /// Read-only Metadata
 impl OperationRef {
     pub fn name(&self) -> OperationName {
@@ -241,6 +337,31 @@ impl OperationRef {
             OperationName::clone(&*name_ptr)
         }
     }
+
+    /// Returns a handle to the nearest containing [Operation] of this operation, if it is attached
+    /// to one
+    pub fn parent_op(&self) -> Option<OperationRef> {
+        self.parent_region().and_then(|region| region.parent())
+    }
+
+    /// Returns a handle to the containing [Region] of this operation, if it is attached to one
+    pub fn parent_region(&self) -> Option<RegionRef> {
+        self.parent().and_then(|block| block.parent())
+    }
+
+    /// Returns the nearest [SymbolTable] from this operation.
+    ///
+    /// Returns `None` if no parent of this operation is a valid symbol table.
+    pub fn nearest_symbol_table(&self) -> Option<OperationRef> {
+        let mut parent = self.parent_op();
+        while let Some(parent_op) = parent.take() {
+            if parent_op.name().implements::<dyn SymbolTable>() {
+                return Some(parent_op);
+            }
+            parent = parent_op.parent_op();
+        }
+        None
+    }
 }
 
 /// Metadata
@@ -512,13 +633,13 @@ impl Operation {
 
     /// Returns a handle to the containing [Region] of this operation, if it is attached to one
     pub fn parent_region(&self) -> Option<RegionRef> {
-        self.parent().and_then(|block| block.parent())
+        self.as_operation_ref().parent_region()
     }
 
     /// Returns a handle to the nearest containing [Operation] of this operation, if it is attached
     /// to one
     pub fn parent_op(&self) -> Option<OperationRef> {
-        self.parent_region().and_then(|region| region.parent())
+        self.as_operation_ref().parent_op()
     }
 
     /// Returns a handle to the nearest containing [Operation] of type `T` for this operation, if it
@@ -1033,54 +1154,20 @@ impl Operation {
 
 /// Insertion
 impl Operation {
-    pub fn insert_at_start(&mut self, mut block: BlockRef) {
-        assert!(
-            self.parent().is_none(),
-            "cannot insert operation that is already attached to another block"
-        );
-        {
-            let mut block = block.borrow_mut();
-            block.body_mut().push_front(self.as_operation_ref());
-        }
+    pub fn insert_at_start(&mut self, block: BlockRef) {
+        self.as_operation_ref().insert_at_start(block);
     }
 
-    pub fn insert_at_end(&mut self, mut block: BlockRef) {
-        assert!(
-            self.parent().is_none(),
-            "cannot insert operation that is already attached to another block"
-        );
-        {
-            let mut block = block.borrow_mut();
-            block.body_mut().push_back(self.as_operation_ref());
-        }
+    pub fn insert_at_end(&mut self, block: BlockRef) {
+        self.as_operation_ref().insert_at_end(block);
     }
 
     pub fn insert_before(&mut self, before: OperationRef) {
-        assert!(
-            self.parent().is_none(),
-            "cannot insert operation that is already attached to another block"
-        );
-        let mut block = before.parent().expect("'before' block is not attached to a block");
-        {
-            let mut block = block.borrow_mut();
-            let block_body = block.body_mut();
-            let mut cursor = unsafe { block_body.cursor_mut_from_ptr(before) };
-            cursor.insert_before(self.as_operation_ref());
-        }
+        self.as_operation_ref().insert_before(before);
     }
 
     pub fn insert_after(&mut self, after: OperationRef) {
-        assert!(
-            self.parent().is_none(),
-            "cannot insert operation that is already attached to another block"
-        );
-        let mut block = after.parent().expect("'after' block is not attached to a block");
-        {
-            let mut block = block.borrow_mut();
-            let block_body = block.body_mut();
-            let mut cursor = unsafe { block_body.cursor_mut_from_ptr(after) };
-            cursor.insert_after(self.as_operation_ref());
-        }
+        self.as_operation_ref().insert_after(after);
     }
 }
 

hir/src/ir/symbols.rs

@@ -151,16 +151,7 @@ impl Operation {
     ///
     /// Returns `None` if no parent of this operation is a valid symbol table.
     pub fn nearest_symbol_table(&self) -> Option<OperationRef> {
-        let mut parent = self.parent_op();
-        while let Some(parent_op) = parent.take() {
-            let op = parent_op.borrow();
-            if op.implements::<dyn SymbolTable>() {
-                drop(op);
-                return Some(parent_op);
-            }
-            parent = op.parent_op();
-        }
-        None
+        self.as_operation_ref().nearest_symbol_table()
     }
 
     /// Returns the operation registered with the given symbol name within the closest symbol table

hir/src/patterns/rewriter.rs

@@ -356,9 +356,9 @@ pub trait Rewriter: Builder + RewriterListener {
     }
 
     /// Insert an unlinked operation at the end of `ip`
-    fn insert_op_at_end(&mut self, mut op: OperationRef, ip: BlockRef) {
+    fn insert_op_at_end(&mut self, op: OperationRef, ip: BlockRef) {
         let prev = ProgramPoint::before(op);
-        op.borrow_mut().insert_at_end(ip);
+        op.insert_at_end(ip);
         self.notify_operation_inserted(op, prev);
     }