feat: added getters for function ctx

Author: Zffu

compiler/astoir_hir/src/ctx.rs

@@ -13,7 +13,7 @@ use diagnostics::{
 use crate::{nodes::HIRNode, scope::HIRGlobalScopeStorage};
 
 pub type HIRFunction = (Option<Type>, Vec<(u64, Type)>, String);
-pub type HIRFunctionImpl = (HIRBranchedContext, Box<HIRNode>);
+pub type HIRFunctionImpl = Box<HIRNode>;
 
 /// The function HIR context. Contains a mapping from element name hash to element index and other variable information.
 /// Uses a branch based system to contain variables.

compiler/astoir_hir/src/scope.rs

@@ -16,6 +16,7 @@ pub struct HIRGlobalScopeStorage {
     pub scope: TypedGlobalScope,
     pub descriptors: Vec<HIRFunction>,
     pub implementations: Vec<HIRFunctionImpl>,
+    pub contexts: Vec<HIRBranchedContext>,
 }
 
 impl HIRGlobalScopeStorage {
@@ -24,6 +25,7 @@ impl HIRGlobalScopeStorage {
             scope: TypedGlobalScope::new(),
             descriptors: vec![],
             implementations: vec![],
+            contexts: vec![],
         }
     }
 
@@ -46,7 +48,8 @@ impl HIRGlobalScopeStorage {
         origin: &K,
     ) -> DiagnosticResult<usize> {
         self.descriptors.push(descriptor);
-        self.implementations.push((brctx, implementation));
+        self.implementations.push(implementation);
+        self.contexts.push(brctx);
 
         self.scope.append(
             name,
@@ -83,7 +86,8 @@ impl HIRGlobalScopeStorage {
         origin: &K,
     ) -> DiagnosticResult<usize> {
         self.descriptors.push(descriptor);
-        self.implementations.push((brctx, implementation));
+        self.implementations.push(implementation);
+        self.contexts.push(brctx);
 
         let ind = self.scope.value_to_ind[&TypedGlobalScopeEntry::Type(struct_type)];
 
@@ -164,14 +168,27 @@ impl HIRGlobalScopeStorage {
         return Ok(self.descriptors[ind].clone());
     }
 
+    pub fn get_function_ctx<K: DiagnosticSpanOrigin>(
+        &self,
+        name: EntryKey,
+        origin: &K,
+    ) -> DiagnosticResult<HIRBranchedContext> {
+        let ind = self.scope.get_function_ctx(name, origin)?;
+
+        return Ok(self.contexts[ind].clone());
+    }
+
     pub fn get_function_impl<K: DiagnosticSpanOrigin>(
         &self,
         name: EntryKey,
         origin: &K,
-    ) -> DiagnosticResult<HIRFunctionImpl> {
+    ) -> DiagnosticResult<(HIRFunctionImpl, HIRBranchedContext)> {
         let ind = self.scope.get_function_impl(name, origin)?;
 
-        return Ok(self.implementations[ind].clone());
+        return Ok((
+            self.implementations[ind].clone(),
+            self.contexts[ind].clone(),
+        ));
     }
 
     pub fn get_implless_function<K: DiagnosticSpanOrigin>(
@@ -201,12 +218,13 @@ impl HIRGlobalScopeStorage {
         &self,
         name: EntryKey,
         origin: &K,
-    ) -> DiagnosticResult<(HIRFunction, HIRFunctionImpl, RawType)> {
+    ) -> DiagnosticResult<(HIRFunction, HIRFunctionImpl, HIRBranchedContext, RawType)> {
         let res = self.scope.get_exact_struct_function(name, origin)?;
 
         return Ok((
             self.descriptors[res.0].clone(),
             self.implementations[res.1].clone(),
+            self.contexts[res.1].clone(),
             res.2,
         ));
     }

compiler/astoir_mir_lowering/src/lib.rs

@@ -6,7 +6,8 @@ use astoir_hir::{
 };
 use astoir_mir::{ctx::MIRContext, funcs::MIRFunction};
 use compiler_typing::{
-    SizedType, TypedGlobalScopeEntry, raw::RawType, structs::LoweredStructTypeContainer, tree::Type,
+    SizedType, TypedGlobalScope, TypedGlobalScopeEntry, raw::RawType,
+    structs::LoweredStructTypeContainer, tree::Type,
 };
 use compiler_utils::utils::indexed::IndexStorage;
 use diagnostics::{DiagnosticResult, unsure_panic};
@@ -68,6 +69,9 @@ pub fn lower_hir(ctx: HIRContext) -> DiagnosticResult<MIRContext> {
                 lower_hir_top_level(node, &mut lowering_ctx)?;
             }
 
+            TypedGlobalScopeEntry::TypeAlias(_) => continue,
+            TypedGlobalScopeEntry::Type(_) => continue,
+
             TypedGlobalScopeEntry::ImplLessFunction(descriptor_ind) => {
                 let descriptor =
                     lowering_ctx.hir_ctx.global_scope.descriptors[descriptor_ind].clone();
@@ -96,7 +100,7 @@ pub fn lower_hir(ctx: HIRContext) -> DiagnosticResult<MIRContext> {
                 lowering_ctx.mir_ctx.append_function(func);
             }
 
-            _ => todo!("Add support for remaining nodes"),
+            _ => todo!("Add support for remaining nodes {:#?}", entry),
         };
     }
 

compiler/astoir_mir_lowering/src/vars.rs

@@ -46,7 +46,7 @@ pub fn lower_hir_variable_declaration(
             }
         };
 
-        let local_ctx = ctx.hir_ctx.global_scope.implementations[*fns_ind].0.clone();
+        let local_ctx = ctx.hir_ctx.global_scope.contexts[*fns_ind].0.clone();
 
         if local_ctx.is_eligible_for_ssa(variable) {
             if default_val.is_some() {

compiler/compiler_global_scope/src/entry.rs

@@ -13,6 +13,7 @@ pub enum GlobalStorageEntryType<T: Hash, R: Hash> {
         impl_ind: usize,
     },
     ImplLessFunction(usize),
+	HalfImplFunction { descriptor_ind: usize, branch_ctx: usize }
     StructFunction {
         descriptor_ind: usize,
         impl_ind: usize,

compiler/compiler_global_scope/src/lib.rs

@@ -57,6 +57,20 @@ impl<T: Clone + Hash + Eq, R: Clone + Hash + Eq> GlobalScopeStorage<T, R> {
             return Err(make_already_in_scope(origin, &name.name_hash).into());
         }
 
+        if let GlobalStorageEntryType::Function { .. } = entry {
+            self.descriptor_counter += 1;
+            self.impl_counter += 1;
+        }
+
+        if let GlobalStorageEntryType::ImplLessFunction(_) = entry {
+            self.descriptor_counter += 1;
+        }
+
+        if let GlobalStorageEntryType::StructFunction { .. } = entry {
+            self.descriptor_counter += 1;
+            self.impl_counter += 1;
+        }
+
         let parent_index = self.entries.len();
 
         self.value_to_ind.insert(entry.clone(), parent_index);
@@ -125,6 +139,10 @@ impl<T: Clone + Hash + Eq, R: Clone + Hash + Eq> GlobalScopeStorage<T, R> {
                 impl_ind: _,
             } => Ok(descriptor_ind),
             GlobalStorageEntryType::ImplLessFunction(descriptor_ind) => Ok(descriptor_ind),
+            GlobalStorageEntryType::HalfImplFunction {
+                descriptor_ind,
+                branch_ctx: _,
+            } => Ok(descriptor_ind),
             GlobalStorageEntryType::StructFunction {
                 descriptor_ind,
                 impl_ind: _,
@@ -135,6 +153,36 @@ impl<T: Clone + Hash + Eq, R: Clone + Hash + Eq> GlobalScopeStorage<T, R> {
         };
     }
 
+    pub fn get_function_ctx<K: DiagnosticSpanOrigin>(
+        &self,
+        name: EntryKey,
+        origin: &K,
+    ) -> DiagnosticResult<usize> {
+        let base = self.get_base(name, origin)?;
+
+        return match base {
+            GlobalStorageEntryType::Function {
+                descriptor_ind: _,
+                impl_ind,
+            } => Ok(impl_ind),
+
+            GlobalStorageEntryType::StructFunction {
+                descriptor_ind: _,
+                impl_ind,
+                struct_type: _,
+            } => Ok(impl_ind),
+
+            GlobalStorageEntryType::HalfImplFunction {
+                descriptor_ind: _,
+                branch_ctx,
+            } => Ok(branch_ctx),
+
+            _ => Err(
+                make_expected_simple_error(origin, &"function with implementation", &base).into(),
+            ),
+        };
+    }
+
     pub fn get_function_impl<K: DiagnosticSpanOrigin>(
         &self,
         name: EntryKey,