fix: add type annotation check during type checking in assignment (#757)

* fix: add type annotation check during type checking in assignment Signed-off-by: zongz <[email protected]> * fix: fix failed test case Signed-off-by: zongz <[email protected]> * fix: rm useless comments Signed-off-by: zongz <[email protected]> * fix: fix test case Signed-off-by: zongz <[email protected]> * fix: rm unnecessary changes Signed-off-by: zongz <[email protected]> * fix: rm an empty line Signed-off-by: zongz <[email protected]> --------- Signed-off-by: zongz <[email protected]>
kcl-lang · Oct 12, 2023 · e13ef0e · e13ef0e
1 parent f63c71b
commit e13ef0e
Show file tree

Hide file tree

Showing 8 changed files with 154 additions and 2 deletions.
diff --git a/kclvm/sema/src/resolver/node.rs b/kclvm/sema/src/resolver/node.rs
@@ -195,6 +195,9 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Resolver<'ctx> {
                 self.must_assignable_to(value_ty.clone(), expected_ty, target.get_span_pos(), None)
             }
         }
+        // Check type annotation if exists.
+        self.check_assignment_type_annotation(assign_stmt, value_ty.clone());
+
         value_ty
     }
 

diff --git a/kclvm/sema/src/resolver/test_data/assign_in_lambda.k b/kclvm/sema/src/resolver/test_data/assign_in_lambda.k
@@ -0,0 +1,6 @@
+lambda {
+    containers = []
+    if True:
+        containers = []
+    images: [str] = [c.image for c in containers]
+}
diff --git a/kclvm/sema/src/resolver/tests.rs b/kclvm/sema/src/resolver/tests.rs
@@ -521,3 +521,21 @@ fn test_resolve_program_import_suggest() {
         "name 's' is not defined, did you mean '[\"s1\"]'?"
     );
 }
+
+#[test]
+fn test_resolve_assignment_in_lambda() {
+    let sess = Arc::new(ParseSession::default());
+    let mut program = load_program(
+        sess.clone(),
+        &["./src/resolver/test_data/assign_in_lambda.k"],
+        None,
+    )
+    .unwrap();
+    let scope = resolve_program(&mut program);
+    let main_scope = scope.scope_map.get("__main__").unwrap().clone();
+    assert_eq!(main_scope.borrow().children.len(), 1);
+    let lambda_scope = main_scope.borrow().children[0].clone();
+    assert_eq!(lambda_scope.borrow().elems.len(), 2);
+    let images_scope_obj = lambda_scope.borrow().elems.get("images").unwrap().clone();
+    assert_eq!(images_scope_obj.borrow().ty.ty_str(), "[str]");
+}
diff --git a/kclvm/sema/src/resolver/ty.rs b/kclvm/sema/src/resolver/ty.rs
@@ -2,7 +2,7 @@ use std::rc::Rc;
 
 use crate::resolver::Resolver;
 use crate::ty::parser::parse_type_str;
-use crate::ty::{assignable_to, Attr, DictType, SchemaType, Type, TypeKind};
+use crate::ty::{assignable_to, is_upper_bound, Attr, DictType, SchemaType, Type, TypeKind};
 use indexmap::IndexMap;
 use kclvm_ast::ast;
 use kclvm_ast::pos::GetPos;
@@ -118,6 +118,70 @@ impl<'ctx> Resolver<'ctx> {
         }
     }
 
+    /// Check the type assignment statement between type annotation and target.
+    pub fn check_assignment_type_annotation(
+        &mut self,
+        assign_stmt: &kclvm_ast::ast::AssignStmt,
+        value_ty: Rc<Type>,
+    ) {
+        if assign_stmt.type_annotation.is_none() {
+            return;
+        }
+        for target in &assign_stmt.targets {
+            if target.node.names.is_empty() {
+                continue;
+            }
+            let name = &target.node.names[0].node;
+            // If the assignment statement has type annotation, check the type of value and the type annotation of target
+
+            if let Some(ty_annotation) = &assign_stmt.ty {
+                let annotation_ty =
+                    self.parse_ty_with_scope(&ty_annotation.node, ty_annotation.get_span_pos());
+                // If the target defined in the scope, check the type of value and the type annotation of target
+                let target_ty = if let Some(obj) = self.scope.borrow().elems.get(name) {
+                    let obj = obj.borrow();
+                    if obj.ty.is_any() {
+                        annotation_ty
+                    } else {
+                        if !is_upper_bound(annotation_ty.clone(), obj.ty.clone()) {
+                            self.handler.add_error(
+                                ErrorKind::TypeError,
+                                &[
+                                    Message {
+                                        range: target.get_span_pos(),
+                                        style: Style::LineAndColumn,
+                                        message: format!(
+                                            "can not change the type of '{}' to {}",
+                                            name,
+                                            annotation_ty.ty_str()
+                                        ),
+                                        note: None,
+                                        suggested_replacement: None,
+                                    },
+                                    Message {
+                                        range: obj.get_span_pos(),
+                                        style: Style::LineAndColumn,
+                                        message: format!("expected {}", obj.ty.ty_str()),
+                                        note: None,
+                                        suggested_replacement: None,
+                                    },
+                                ],
+                            );
+                        }
+                        obj.ty.clone()
+                    }
+                } else {
+                    annotation_ty
+                };
+
+                self.set_type_to_scope(name, target_ty.clone(), target.get_span_pos());
+
+                // Check the type of value and the type annotation of target
+                self.must_assignable_to(value_ty.clone(), target_ty, target.get_span_pos(), None)
+            }
+        }
+    }
+
     /// The check type main function, returns a boolean result.
     #[inline]
     pub fn check_type(&mut self, ty: Rc<Type>, expected_ty: Rc<Type>, range: &Range) -> bool {

diff --git a/kclvm/sema/src/ty/unify.rs b/kclvm/sema/src/ty/unify.rs
@@ -64,6 +64,8 @@ pub fn subsume(ty_lhs: Rc<Type>, ty_rhs: Rc<Type>, check_left_any: bool) -> bool
         let (ty_rhs_key, ty_rhs_val) = ty_rhs.dict_entry_ty();
         subsume(ty_lhs_key, ty_rhs_key, check_left_any)
             && subsume(ty_lhs_val, ty_rhs_val, check_left_any)
+    } else if ty_lhs.is_str() && ty_rhs.is_literal() && ty_rhs.is_str() {
+        return true;
     } else {
         equal(ty_lhs, ty_rhs)
     }

diff --git a/kclvm/tools/src/LSP/src/test_data/goto_def_test/goto_def.k b/kclvm/tools/src/LSP/src/test_data/goto_def_test/goto_def.k
@@ -35,7 +35,7 @@ schema Reviewer:
     teams?: [int]
     users?: [int]
 
-reviewers: [Reviewer] = [{team: [1]}]
+reviewers: [Reviewer] = [Reviewer {teams: [1]}]
 
 schema Fib:
     n1 = n - 1

diff --git a/kclvm/tools/src/LSP/src/test_data/hover_test/assign_in_lambda.k b/kclvm/tools/src/LSP/src/test_data/hover_test/assign_in_lambda.k
@@ -0,0 +1,6 @@
+lambda {
+    containers = []
+    if True:
+        containers = []
+    images: [str] = [c.image for c in containers]
+}
diff --git a/kclvm/tools/src/LSP/src/tests.rs b/kclvm/tools/src/LSP/src/tests.rs
@@ -757,6 +757,59 @@ fn hover_test() {
     )
 }
 
+#[test]
+fn hover_assign_in_lambda_test() {
+    let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+    let mut path = root.clone();
+
+    path.push("src/test_data/hover_test/assign_in_lambda.k");
+
+    let path = path.to_str().unwrap();
+    let src = std::fs::read_to_string(path.clone()).unwrap();
+    let server = Project {}.server();
+
+    // Mock open file
+    server.notification::<lsp_types::notification::DidOpenTextDocument>(
+        lsp_types::DidOpenTextDocumentParams {
+            text_document: TextDocumentItem {
+                uri: Url::from_file_path(path).unwrap(),
+                language_id: "KCL".to_string(),
+                version: 0,
+                text: src,
+            },
+        },
+    );
+
+    let id = server.next_request_id.get();
+    server.next_request_id.set(id.wrapping_add(1));
+
+    let r: Request = Request::new(
+        id.into(),
+        "textDocument/hover".to_string(),
+        HoverParams {
+            text_document_position_params: TextDocumentPositionParams {
+                text_document: TextDocumentIdentifier {
+                    uri: Url::from_file_path(path).unwrap(),
+                },
+                position: Position::new(4, 7),
+            },
+            work_done_progress_params: Default::default(),
+        },
+    );
+
+    // Send request and wait for it's response
+    let res = server.send_and_receive(r);
+
+    assert_eq!(
+        res.result.unwrap(),
+        to_json(Hover {
+            contents: HoverContents::Scalar(MarkedString::String("images: [str]".to_string()),),
+            range: None
+        })
+        .unwrap()
+    )
+}
+
 #[test]
 fn formatting_test() {
     let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));