From 330ee2abbf54319f766dbe3fd5876b3ecbd813f5 Mon Sep 17 00:00:00 2001
From: peefy <xpf6677@163.com>
Date: Mon, 5 Aug 2024 13:48:42 +0800
Subject: [PATCH] feat: add function mode for the evaluator

Signed-off-by: peefy <xpf6677@163.com>
---
 kclvm/evaluator/src/lib.rs                    | 19 ++++++++++--
 kclvm/evaluator/src/module.rs                 |  8 +++--
 ...lvm_evaluator__tests__function_stmt_0.snap |  5 +++
 kclvm/evaluator/src/tests.rs                  | 31 +++++++++++++++++++
 4 files changed, 58 insertions(+), 5 deletions(-)
 create mode 100644 kclvm/evaluator/src/snapshots/kclvm_evaluator__tests__function_stmt_0.snap

diff --git a/kclvm/evaluator/src/lib.rs b/kclvm/evaluator/src/lib.rs
index 44a589bd3..de5a1256b 100644
--- a/kclvm/evaluator/src/lib.rs
+++ b/kclvm/evaluator/src/lib.rs
@@ -132,13 +132,26 @@ impl<'ctx> Evaluator<'ctx> {
     pub fn run(self: &Evaluator<'ctx>) -> Result<(String, String)> {
         if let Some(modules) = self.program.pkgs.get(kclvm_ast::MAIN_PKG) {
             self.init_scope(kclvm_ast::MAIN_PKG);
-            self.compile_ast_modules(modules)
+            self.compile_ast_modules(modules);
         }
         Ok(self.plan_globals_to_string())
     }
 
+    /// Evaluate the program with the function mode and return the JSON and YAML result,
+    /// which means treating the files in the entire main package as a function run to
+    /// return the result of the function run, rather than a dictionary composed of each
+    /// configuration attribute.
+    pub fn run_as_function(self: &Evaluator<'ctx>) -> ValueRef {
+        if let Some(modules) = self.program.pkgs.get(kclvm_ast::MAIN_PKG) {
+            self.init_scope(kclvm_ast::MAIN_PKG);
+            self.compile_ast_modules(modules)
+        } else {
+            ValueRef::undefined()
+        }
+    }
+
     /// Plan globals to a planed json and yaml string.
-    pub fn plan_globals_to_string(&self) -> (String, String) {
+    pub(crate) fn plan_globals_to_string(&self) -> (String, String) {
         let current_pkgpath = self.current_pkgpath();
         let pkg_scopes = &self.pkg_scopes.borrow();
         let scopes = pkg_scopes
@@ -177,7 +190,7 @@ impl<'ctx> Evaluator<'ctx> {
         Ok(self.undefined_value())
     }
 
-    fn plan_value(&self, value: &ValueRef) -> (String, String) {
+    pub fn plan_value(&self, value: &ValueRef) -> (String, String) {
         let mut ctx = self.runtime_ctx.borrow_mut();
         let value = match ctx.buffer.custom_manifests_output.clone() {
             Some(output) => ValueRef::from_yaml_stream(&mut ctx, &output).unwrap(),
diff --git a/kclvm/evaluator/src/module.rs b/kclvm/evaluator/src/module.rs
index 613c78524..d79671521 100644
--- a/kclvm/evaluator/src/module.rs
+++ b/kclvm/evaluator/src/module.rs
@@ -2,6 +2,7 @@
 
 use kclvm_ast::ast;
 use kclvm_ast::walker::TypedResultWalker;
+use kclvm_runtime::ValueRef;
 
 use super::Evaluator;
 use crate::error as kcl_error;
@@ -80,7 +81,7 @@ impl<'ctx> Evaluator<'ctx> {
     /// 1. scan all possible global variables and allocate undefined values to global pointers.
     /// 2. build all user-defined schema/rule types.
     /// 3. evaluate all codes for the third time.
-    pub(crate) fn compile_ast_modules(&self, modules: &'ctx [ast::Module]) {
+    pub(crate) fn compile_ast_modules(&self, modules: &'ctx [ast::Module]) -> ValueRef {
         // Scan global variables
         for ast_module in modules {
             // Pre define global variables with undefined values
@@ -90,10 +91,13 @@ impl<'ctx> Evaluator<'ctx> {
         for ast_module in modules {
             self.compile_module_import_and_types(ast_module);
         }
+        let mut result = ValueRef::undefined();
         // Compile the ast module in the pkgpath.
         for ast_module in modules {
-            self.walk_module(ast_module)
+            result = self
+                .walk_module(ast_module)
                 .expect(kcl_error::RUNTIME_ERROR_MSG);
         }
+        result
     }
 }
diff --git a/kclvm/evaluator/src/snapshots/kclvm_evaluator__tests__function_stmt_0.snap b/kclvm/evaluator/src/snapshots/kclvm_evaluator__tests__function_stmt_0.snap
new file mode 100644
index 000000000..3d3de056a
--- /dev/null
+++ b/kclvm/evaluator/src/snapshots/kclvm_evaluator__tests__function_stmt_0.snap
@@ -0,0 +1,5 @@
+---
+source: evaluator/src/tests.rs
+expression: "format!(\"{}\", evaluator.run_as_function().to_string())"
+---
+"bar"
diff --git a/kclvm/evaluator/src/tests.rs b/kclvm/evaluator/src/tests.rs
index 17630d48a..25c3c09a8 100644
--- a/kclvm/evaluator/src/tests.rs
+++ b/kclvm/evaluator/src/tests.rs
@@ -24,6 +24,37 @@ macro_rules! evaluator_snapshot {
     };
 }
 
+#[macro_export]
+macro_rules! evaluator_function_snapshot {
+    ($name:ident, $src:expr) => {
+        #[test]
+        fn $name() {
+            let p = load_packages(&LoadPackageOptions {
+                paths: vec!["test.k".to_string()],
+                load_opts: Some(LoadProgramOptions {
+                    k_code_list: vec![$src.to_string()],
+                    ..Default::default()
+                }),
+                load_builtin: false,
+                ..Default::default()
+            })
+            .unwrap();
+            let evaluator = Evaluator::new(&p.program);
+            insta::assert_snapshot!(format!("{}", evaluator.run_as_function().to_string()));
+        }
+    };
+}
+
+evaluator_function_snapshot! {function_stmt_0, r#"
+import json
+
+config = {
+  foo: "bar"
+}
+
+json.encode("${config.foo}")
+"#}
+
 evaluator_snapshot! {expr_stmt_0, "1"}
 evaluator_snapshot! {expr_stmt_1, "2.0"}
 evaluator_snapshot! {expr_stmt_2, "True"}