From 4eac50d4ef3feb942b4a24a7b9c912d812a576ad Mon Sep 17 00:00:00 2001
From: zongz <68977949+zong-zhe@users.noreply.github.com>
Date: Thu, 26 Oct 2023 12:34:59 +0800
Subject: [PATCH] fix: fix the lambda argument type, when it is schema type,
 into schema type instance (#803)

* fix: fix the lambda argument type, when it is schema type, into schema type instance

Signed-off-by: zongz <zongzhe1024@163.com>

* fix: make fmt

Signed-off-by: zongz <zongzhe1024@163.com>

---------

Signed-off-by: zongz <zongzhe1024@163.com>
---
 kclvm/sema/src/resolver/node.rs            | 11 +++++++++++
 test/grammar/lambda/arg_type/main.k        |  9 +++++++++
 test/grammar/lambda/arg_type/stdout.golden |  2 ++
 test/grammar/lambda/arg_type/sub.k         |  2 ++
 4 files changed, 24 insertions(+)
 create mode 100644 test/grammar/lambda/arg_type/main.k
 create mode 100644 test/grammar/lambda/arg_type/stdout.golden
 create mode 100644 test/grammar/lambda/arg_type/sub.k

diff --git a/kclvm/sema/src/resolver/node.rs b/kclvm/sema/src/resolver/node.rs
index 190554eb5..c25fcbed3 100644
--- a/kclvm/sema/src/resolver/node.rs
+++ b/kclvm/sema/src/resolver/node.rs
@@ -943,6 +943,17 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for Resolver<'ctx> {
                 let name = arg.node.get_name();
                 let arg_ty = args.node.get_arg_type(i);
                 let ty = self.parse_ty_with_scope(&arg_ty, arg.get_span_pos());
+
+                // If the arguments type of a lambda is a schema type,
+                // It should be marked as an schema instance type.
+                let ty = if let TypeKind::Schema(sty) = &ty.kind {
+                    let mut arg_ty = sty.clone();
+                    arg_ty.is_instance = true;
+                    Rc::new(Type::schema(arg_ty))
+                } else {
+                    ty.clone()
+                };
+
                 params.push(Parameter {
                     name,
                     ty: ty.clone(),
diff --git a/test/grammar/lambda/arg_type/main.k b/test/grammar/lambda/arg_type/main.k
new file mode 100644
index 000000000..3cf3c8756
--- /dev/null
+++ b/test/grammar/lambda/arg_type/main.k
@@ -0,0 +1,9 @@
+import sub as s
+
+identity = lambda res: s.Test {
+  res {
+    name = "world" 
+  }
+}
+
+c = identity(s.Test{name="hello"})
diff --git a/test/grammar/lambda/arg_type/stdout.golden b/test/grammar/lambda/arg_type/stdout.golden
new file mode 100644
index 000000000..81be50f43
--- /dev/null
+++ b/test/grammar/lambda/arg_type/stdout.golden
@@ -0,0 +1,2 @@
+c:
+  name: world
\ No newline at end of file
diff --git a/test/grammar/lambda/arg_type/sub.k b/test/grammar/lambda/arg_type/sub.k
new file mode 100644
index 000000000..df60bab68
--- /dev/null
+++ b/test/grammar/lambda/arg_type/sub.k
@@ -0,0 +1,2 @@
+schema Test:
+    name: str
\ No newline at end of file