fix: add module ty and typealias to symbol

Signed-off-by: never <[email protected]>
kcl-lang · Dec 19, 2023 · 1f44ec9 · 1f44ec9
1 parent fc2c67a
commit 1f44ec9
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 8 deletions.
diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
@@ -331,11 +331,8 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for AdvancedResolver<'ctx> {
     }
 
     fn walk_quant_expr(&mut self, quant_expr: &'ctx ast::QuantExpr) -> Self::Result {
+        let (start, end) = (self.ctx.start_pos.clone(), self.ctx.end_pos.clone());
         self.expr(&quant_expr.target);
-        let (start, mut end) = quant_expr.test.get_span_pos();
-        if let Some(if_cond) = &quant_expr.if_cond {
-            end = if_cond.get_end_pos();
-        }
         self.enter_local_scope(
             &self.ctx.current_filename.as_ref().unwrap().clone(),
             start,
@@ -704,6 +701,7 @@ impl<'ctx> AdvancedResolver<'ctx> {
                 | ast::Expr::Config(_)
                 | ast::Expr::Schema(_)
                 | ast::Expr::ConfigIfEntry(_)
+                | ast::Expr::Quant(_)
         ) {
             let (start, end) = expr.get_span_pos();
             self.ctx.start_pos = start;
@@ -741,7 +739,25 @@ impl<'ctx> AdvancedResolver<'ctx> {
                 first_symbol = self
                     .gs
                     .get_symbols()
-                    .get_symbol_by_fully_qualified_name(&import_info.unwrap().fully_qualified_name)
+                    .get_symbol_by_fully_qualified_name(&import_info.unwrap().fully_qualified_name);
+            }
+
+            if let Some(first_symbol) = first_symbol {
+                if self
+                    .gs
+                    .get_symbols()
+                    .get_symbol(first_symbol)
+                    .unwrap()
+                    .get_sema_info()
+                    .ty
+                    .is_none()
+                {
+                    if let Some(ty) = self.ctx.node_ty_map.get(&first_name.id) {
+                        self.gs
+                            .get_symbols_mut()
+                            .set_symbol_type(first_symbol, ty.clone());
+                    }
+                }
             }
         }
         match first_symbol {

diff --git a/kclvm/sema/src/core/symbol.rs b/kclvm/sema/src/core/symbol.rs
@@ -6,7 +6,7 @@ use indexmap::IndexMap;
 use kclvm_error::{diagnostic::Range, Position};
 
 use super::package::ModuleInfo;
-use crate::ty::{Type, TypeKind};
+use crate::ty::{Type, TypeKind, TypeRef};
 use kclvm_ast::ast::AstIndex;
 
 pub trait Symbol {
@@ -156,6 +156,53 @@ impl KCLSymbolData {
         }
     }
 
+    pub fn set_symbol_type(&mut self, id: SymbolRef, ty: TypeRef) {
+        match id.get_kind() {
+            SymbolKind::Schema => {
+                self.schemas.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::Attribute => {
+                self.attributes.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::Value => {
+                self.values.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::Package => {
+                self.packages.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::TypeAlias => {
+                self.type_aliases.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::Unresolved => {
+                self.unresolved.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+            SymbolKind::Rule => {
+                self.rules.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
+        }
+    }
+
     pub fn get_type_symbol(
         &self,
         ty: &Type,
@@ -1151,7 +1198,7 @@ impl Symbol for TypeAliasSymbol {
     }
 
     fn get_owner(&self) -> Option<SymbolRef> {
-        None
+        Some(self.owner)
     }
 
     fn get_definition(&self) -> Option<SymbolRef> {
@@ -1275,7 +1322,7 @@ impl Symbol for RuleSymbol {
     }
 
     fn get_owner(&self) -> Option<SymbolRef> {
-        None
+        Some(self.owner)
     }
 
     fn get_definition(&self) -> Option<SymbolRef> {