fix: fix advanced resolver current schema symbol (#1110)

* fix: fix advanced resolver current schema symbol Signed-off-by: he1pa <[email protected]> * add ut Signed-off-by: he1pa <[email protected]> * fix ut Signed-off-by: he1pa <[email protected]> --------- Signed-off-by: he1pa <[email protected]>
kcl-lang · Mar 6, 2024 · 259ac7a · 259ac7a
1 parent abfe4e7
commit 259ac7a
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 5 deletions.
diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
@@ -1153,12 +1153,13 @@ impl<'ctx> AdvancedResolver<'ctx> {
             if let Some(key) = &entry.node.key {
                 self.ctx.maybe_def = true;
                 if let Some(key_symbol_ref) = self.expr(key) {
-                    self.set_config_scope_owner(key_symbol_ref);
+                    self.set_current_schema_symbol(key_symbol_ref);
                 }
                 self.ctx.maybe_def = false;
             }
 
             let (start, end) = entry.node.value.get_span_pos();
+
             self.enter_local_scope(
                 &self.ctx.current_filename.as_ref().unwrap().clone(),
                 start,
@@ -1172,19 +1173,33 @@ impl<'ctx> AdvancedResolver<'ctx> {
         self.leave_scope()
     }
 
-    pub(crate) fn set_config_scope_owner(&mut self, key_symbol_ref: SymbolRef) {
+    pub(crate) fn set_current_schema_symbol(&mut self, key_symbol_ref: SymbolRef) {
         let symbols = self.gs.get_symbols();
 
         if let Some(def_symbol_ref) = symbols.get_symbol(key_symbol_ref).unwrap().get_definition() {
             if let Some(node_key) = symbols.symbols_info.symbol_node_map.get(&def_symbol_ref) {
                 if let Some(def_ty) = self.ctx.node_ty_map.get(node_key) {
-                    if def_ty.is_schema() {
+                    if let Some(ty) = get_possible_schema_ty(def_ty.clone()) {
                         self.ctx.current_schema_symbol =
-                            self.gs.get_symbols().get_type_symbol(&def_ty, None);
+                            self.gs.get_symbols().get_type_symbol(&ty, None);
                     }
                 }
             }
         }
+        fn get_possible_schema_ty(ty: Arc<Type>) -> Option<Arc<Type>> {
+            match &ty.kind {
+                crate::ty::TypeKind::List(ty) => get_possible_schema_ty(ty.clone()),
+                crate::ty::TypeKind::Dict(dict_ty) => {
+                    get_possible_schema_ty(dict_ty.val_ty.clone())
+                }
+                crate::ty::TypeKind::Union(_) => {
+                    // Todo: fix union schema type
+                    None
+                }
+                crate::ty::TypeKind::Schema(_) => Some(ty.clone()),
+                _ => None,
+            }
+        }
     }
 
     pub(crate) fn resolve_decorator(&mut self, decorators: &'ctx [ast::NodeRef<ast::CallExpr>]) {

diff --git a/kclvm/tools/src/LSP/src/completion.rs b/kclvm/tools/src/LSP/src/completion.rs
@@ -1613,4 +1613,41 @@ mod tests {
             CompletionResponse::List(_) => panic!("test failed"),
         }
     }
+
+    #[test]
+    #[bench_test]
+    fn schema_type_attr_completion() {
+        let (file, program, _, _, gs) =
+            compile_test_file("src/test_data/completion_test/schema/schema.k");
+
+        let pos = KCLPos {
+            filename: file.to_owned(),
+            line: 18,
+            column: Some(15),
+        };
+
+        let mut got = completion(None, &program, &pos, &gs).unwrap();
+        match &mut got {
+            CompletionResponse::Array(arr) => {
+                let labels: Vec<String> = arr.iter().map(|item| item.label.clone()).collect();
+                assert!(labels.contains(&"name".to_string()));
+            }
+            CompletionResponse::List(_) => panic!("test failed"),
+        }
+
+        let pos = KCLPos {
+            filename: file.to_owned(),
+            line: 19,
+            column: Some(21),
+        };
+
+        let mut got = completion(None, &program, &pos, &gs).unwrap();
+        match &mut got {
+            CompletionResponse::Array(arr) => {
+                let labels: Vec<String> = arr.iter().map(|item| item.label.clone()).collect();
+                assert!(labels.contains(&"name".to_string()));
+            }
+            CompletionResponse::List(_) => panic!("test failed"),
+        }
+    }
 }
diff --git a/kclvm/tools/src/LSP/src/test_data/completion_test/schema/schema.k b/kclvm/tools/src/LSP/src/test_data/completion_test/schema/schema.k
@@ -4,4 +4,17 @@ schema Base:
 schema Person[b: int](Base):
     c: int
 
-p = 
+p = 
+
+schema Name:
+    name: str
+
+schema Config:
+    names: [Name]
+    names1: {str: Name}
+    names2: str | Name
+
+Config{
+    names: [{  }]
+    names1: {"a": {  }}
+}