fix: fix schema index signature check and schema symbol init in advan…

…ced_resolver

Signed-off-by: he1pa <[email protected]>

kclvm/sema/src/advanced_resolver/mod.rs

@@ -49,8 +49,8 @@ use crate::{
     resolver::scope::{NodeKey, NodeTyMap},
 };
 
-use kclvm_ast::ast::AstIndex;
 use kclvm_ast::ast::Program;
+use kclvm_ast::ast::{AstIndex, Stmt};
 use kclvm_ast::walker::MutSelfTypedResultWalker;
 mod node;
 
@@ -113,6 +113,7 @@ impl<'ctx> AdvancedResolver<'ctx> {
                 maybe_def: false,
             },
         };
+        // Scan all scehma symbol
         for (name, modules) in advanced_resolver.ctx.program.pkgs.iter() {
             if !advanced_resolver.gs.new_or_invalidate_pkgs.contains(name) {
                 continue;
@@ -130,6 +131,38 @@ impl<'ctx> AdvancedResolver<'ctx> {
                     pkg_info.pkg_filepath.clone(),
                     pkg_info.kfile_paths.clone(),
                 );
+                for module in modules.iter() {
+                    advanced_resolver.ctx.current_filename = Some(module.filename.clone());
+                    for stmt in &module.body {
+                        if matches!(stmt.node, Stmt::Schema(_)) {
+                            advanced_resolver.stmt(stmt)?;
+                        }
+                    }
+                }
+                advanced_resolver.leave_scope()
+            }
+        }
+
+        for (name, modules) in advanced_resolver.ctx.program.pkgs.iter() {
+            if !advanced_resolver.gs.new_or_invalidate_pkgs.contains(name) {
+                continue;
+            }
+            advanced_resolver.ctx.current_pkgpath = Some(name.clone());
+            if let Some(_) = advanced_resolver.gs.get_packages().get_package_info(name) {
+                if modules.is_empty() {
+                    continue;
+                }
+                if !advanced_resolver.ctx.scopes.is_empty() {
+                    advanced_resolver.ctx.scopes.clear();
+                }
+
+                let scope_ref = advanced_resolver
+                    .gs
+                    .get_scopes_mut()
+                    .get_root_scope(name.to_string())
+                    .unwrap();
+
+                advanced_resolver.ctx.scopes.push(scope_ref);
                 for module in modules.iter() {
                     advanced_resolver.ctx.current_filename = Some(module.filename.clone());
                     advanced_resolver.walk_module(module)?;

kclvm/sema/src/resolver/ty.rs

@@ -352,23 +352,39 @@ impl<'ctx> Resolver<'ctx> {
             // empty dict {}
             TypeKind::Any => true,
             // single key: {key1: value1}
-            TypeKind::StrLit(s) => schema_ty.attrs.len() == 1 && schema_ty.attrs.contains_key(s),
+            TypeKind::StrLit(s) => schema_ty.attrs.len() >= 1 && schema_ty.attrs.contains_key(s),
             // multi key: {
             // key1: value1
             // key2: value2
             // ...
             // }
             TypeKind::Union(types) => {
-                schema_ty.attrs.len() == types.len()
-                    && types.iter().all(|ty| match &ty.kind {
-                        TypeKind::StrLit(s) => schema_ty.attrs.contains_key(s),
+                let (attrs, has_index_signature) = Self::get_schema_attrs(schema_ty);
+                match (attrs.len() >= types.len(), has_index_signature) {
+                    (true, _) => types.iter().all(|ty| match &ty.kind {
+                        TypeKind::StrLit(s) => attrs.contains(s),
                         _ => false,
-                    })
+                    }),
+                    // Todo: do more index_signature check
+                    (false, true) => true,
+                    (false, false) => false,
+                }
             }
             _ => false,
         }
     }
 
+    fn get_schema_attrs(schema_ty: &SchemaType) -> (Vec<String>, bool) {
+        let mut attrs: Vec<String> = schema_ty.attrs.keys().map(|attr| attr.clone()).collect();
+        let mut has_index_signature = schema_ty.index_signature.is_some();
+        if let Some(base) = &schema_ty.base {
+            let (base_attrs, index_signature) = Self::get_schema_attrs(base);
+            attrs.extend(base_attrs);
+            has_index_signature &= index_signature;
+        }
+        (attrs, has_index_signature)
+    }
+
     fn upgrade_named_ty_with_scope(
         &mut self,
         ty: TypeRef,

kclvm/tools/src/LSP/src/goto_def.rs

@@ -441,4 +441,18 @@ mod tests {
         22,
         22
     );
+
+    goto_def_test_snapshot!(
+        goto_base_schema_attr_test,
+        "src/test_data/goto_def_test/goto_base_schema_attr_test/goto_base_schema_attr_test.k",
+        8,
+        12
+    );
+
+    goto_def_test_snapshot!(
+        goto_base_schema_attr_1_test,
+        "src/test_data/goto_def_test/goto_base_schema_attr_1_test/goto_base_schema_attr_1_test.k",
+        4,
+        12
+    );
 }

kclvm/tools/src/LSP/src/snapshots/kcl_language_server__goto_def__tests__goto_base_schema_attr_1_test.snap

@@ -0,0 +1,5 @@
+---
+source: tools/src/LSP/src/goto_def.rs
+expression: "format!(\"{:?}\", { fmt_resp(& res) })"
+---
+"path: \"src/test_data/goto_def_test/goto_base_schema_attr_1_test/types/host.k\", range: Range { start: Position { line: 1, character: 4 }, end: Position { line: 1, character: 8 } }"

kclvm/tools/src/LSP/src/snapshots/kcl_language_server__goto_def__tests__goto_base_schema_attr_test.snap

@@ -0,0 +1,5 @@
+---
+source: tools/src/LSP/src/goto_def.rs
+expression: "format!(\"{:?}\", { fmt_resp(& res) })"
+---
+"path: \"src/test_data/goto_def_test/goto_base_schema_attr_test/goto_base_schema_attr_test.k\", range: Range { start: Position { line: 1, character: 4 }, end: Position { line: 1, character: 8 } }"

kclvm/tools/src/LSP/src/snapshots/kcl_language_server__goto_def__tests__goto_nested_schema_attr_test.snap.new → kclvm/tools/src/LSP/src/snapshots/kcl_language_server__goto_def__tests__goto_nested_schema_attr_test.snap

@@ -1,6 +1,5 @@
 ---
 source: tools/src/LSP/src/goto_def.rs
-assertion_line: 438
 expression: "format!(\"{:?}\", { fmt_resp(& res) })"
 ---
 "path: \"src/test_data/goto_def_test/goto_nested_schema_attr_test/goto_nested_schema_attr_test.k\", range: Range { start: Position { line: 14, character: 4 }, end: Position { line: 14, character: 8 } }"

kclvm/tools/src/LSP/src/test_data/goto_def_test/goto_base_schema_attr_1_test/goto_base_schema_attr_1_test.k

@@ -0,0 +1,6 @@
+import types
+
+hosts: {str: types.HostPort} = {
+    foo: { host: "foo.example.net", port: 80}
+}
+

kclvm/tools/src/LSP/src/test_data/goto_def_test/goto_base_schema_attr_1_test/types/host.k

@@ -0,0 +1,5 @@
+schema Host:
+    host: str
+
+schema HostPort(Host):
+    port: int

kclvm/tools/src/LSP/src/test_data/goto_def_test/goto_base_schema_attr_test/goto_base_schema_attr_test.k

@@ -0,0 +1,9 @@
+schema Host:
+    host: str
+
+schema HostPort(Host):
+    port: int
+
+hosts: {str:HostPort} = {
+    foo: {host: "foo.example.net", port: 80}
+}