fix: add type annotation check during type checking in assignment

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

kclvm/sema/src/resolver/node.rs

@@ -7,7 +7,7 @@ use std::rc::Rc;
 
 use crate::info::is_private_field;
 use crate::ty::{
-    sup, DictType, FunctionType, Parameter, Type, TypeInferMethods, TypeKind,
+    is_upper_bound, sup, DictType, FunctionType, Parameter, Type, TypeInferMethods, TypeKind,
     RESERVED_TYPE_IDENTIFIERS,
 };
 
@@ -185,6 +185,51 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Resolver<'ctx> {
                             self.ctx.ty_ctx.infer_to_variable_type(value_ty.clone()),
                         );
                     }
+                    // If the assign statement has a type annotation like: `a: str = "test_str"`,
+                } else if assign_stmt.type_annotation.is_some() {
+                    // Get the type annotation.
+                    if let Some(ty_annotation) = &assign_stmt.ty {
+                        let annotation_ty = self
+                            .parse_ty_with_scope(&ty_annotation.node, ty_annotation.get_span_pos());
+                        // Set the type annotation as the target type.
+                        self.set_type_to_scope(name, annotation_ty.clone(), target.get_span_pos());
+                        // Check the type of target whether is upper bound of the type annotation.
+                        if let Some(obj) = self.scope.borrow().elems.get(name) {
+                            let obj = obj.borrow();
+                            if !is_upper_bound(obj.ty.clone(), annotation_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,
+                                        },
+                                    ],
+                                );
+                            }
+                        }
+                        // Check the type of value and the type annotation of target
+                        self.must_assignable_to(
+                            value_ty.clone(),
+                            annotation_ty,
+                            target.get_span_pos(),
+                            None,
+                        )
+                    }
                 }
             } else {
                 self.lookup_type_from_scope(name, target.get_span_pos());

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]
+}

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]");
+}

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]
+}

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"));