Skip to content

Commit

Permalink
feat: kcl rust plugin impl and example
Browse files Browse the repository at this point in the history
Signed-off-by: peefy <[email protected]>
  • Loading branch information
Peefy committed Aug 24, 2024
1 parent 32f3739 commit 3885840
Show file tree
Hide file tree
Showing 7 changed files with 78 additions and 4 deletions.
1 change: 1 addition & 0 deletions kclvm/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
source: evaluator/src/tests.rs
expression: "format!(\"{}\", evaluator.run().unwrap().1)"
---
sum: 2
45 changes: 45 additions & 0 deletions kclvm/evaluator/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::Evaluator;
use kclvm_ast::MAIN_PKG;
use kclvm_loader::{load_packages, LoadPackageOptions};
use kclvm_parser::LoadProgramOptions;
use kclvm_runtime::{Context, ValueRef};

#[macro_export]
macro_rules! evaluator_snapshot {
Expand Down Expand Up @@ -504,6 +505,8 @@ fn test_if_stmt_setters() {
assert_eq!(var_setters.len(), 3);
}

use std::cell::RefCell;
use std::rc::Rc;
use std::sync::Arc;
use std::thread;

Expand Down Expand Up @@ -555,3 +558,45 @@ fn run_code(source: &str) -> (String, String) {
let evaluator = Evaluator::new(&p.program);
evaluator.run().unwrap()
}

fn testing_sum(_: &Context, args: &ValueRef, _: &ValueRef) -> anyhow::Result<ValueRef> {
let a = args
.arg_i_int(0, Some(0))
.ok_or(anyhow::anyhow!("expect int value for the first param"))?;
let b = args
.arg_i_int(1, Some(0))
.ok_or(anyhow::anyhow!("expect int value for the second param"))?;
Ok((a + b).into())
}

fn context_with_plugin() -> Rc<RefCell<Context>> {
let mut plugin_functions: kclvm_runtime::IndexMap<String, kclvm_runtime::PluginFunction> =
Default::default();
let func = Arc::new(testing_sum);
plugin_functions.insert("testing.add".to_string(), func);
let mut ctx = Context::new();
ctx.plugin_functions = plugin_functions;
Rc::new(RefCell::new(ctx))
}

#[test]
fn test_exec_with_plugin() {
let src = r#"
import kcl_plugin.testing
sum = testing.add(1, 1)
"#;
let p = load_packages(&LoadPackageOptions {
paths: vec!["test.k".to_string()],
load_opts: Some(LoadProgramOptions {
load_plugins: true,
k_code_list: vec![src.to_string()],
..Default::default()
}),
load_builtin: false,
..Default::default()
})
.unwrap();
let evaluator = Evaluator::new_with_runtime_ctx(&p.program, context_with_plugin());
insta::assert_snapshot!(format!("{}", evaluator.run().unwrap().1));
}
1 change: 1 addition & 0 deletions kclvm/runtime/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ glob = "0.3.0"
uuid = { version = "1.7.0", features = ["serde", "v4"] }
handlebars = "5.1.2"
walkdir = "2.5.0"
anyhow = "1"

[[bin]]
name = "gen-api-spec"
Expand Down
13 changes: 10 additions & 3 deletions kclvm/runtime/src/api/kclvm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::panic::{RefUnwindSafe, UnwindSafe};
use std::rc::Rc;
use std::sync::Arc;
use std::{
cell::RefCell,
cmp::Ordering,
Expand Down Expand Up @@ -329,15 +330,19 @@ impl Default for ContextBuffer {
}
}

#[derive(PartialEq, Clone, Default, Debug)]
/// Plugin functions
pub type PluginFunction =
Arc<dyn Fn(&Context, &ValueRef, &ValueRef) -> anyhow::Result<ValueRef> + Send + Sync>;

#[derive(Clone, Default)]
pub struct Context {
/// Runtime evaluation config.
pub cfg: ContextConfig,

/// kcl.mod path or the pwd path
pub module_path: String,
/// Program work directory
pub workdir: String,
/// Runtime backtrace frame for the debugger.
pub backtrace: Vec<BacktraceFrame>,
/// Imported package path to check the cyclic import process.
pub imported_pkgpath: HashSet<String>,
Expand All @@ -352,7 +357,7 @@ pub struct Context {
pub import_names: IndexMap<String, IndexMap<String, String>>,
/// A buffer to store plugin or hooks function calling results.
pub buffer: ContextBuffer,
/// Objects is to store all KCL object pointers at runtime.
/// Objects is to store all KCL object pointers at runtime for GC.
pub objects: IndexSet<usize>,
/// Log message used to store print results.
pub log_message: String,
Expand All @@ -364,6 +369,8 @@ pub struct Context {
pub panic_info: PanicInfo,
/// Planning options
pub plan_opts: PlanOptions,
/// Builtin plugin functions, the key of the map is the form <module_name>.<module_func> e.g., `hello.say_hello`
pub plugin_functions: IndexMap<String, PluginFunction>,
}

impl UnwindSafe for Context {}
Expand Down
2 changes: 1 addition & 1 deletion kclvm/runtime/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -128,4 +128,4 @@ pub use self::_kclvm::*;
pub mod _kclvm_addr;
pub use self::_kclvm_addr::*;

type IndexMap<K, V> = indexmap::IndexMap<K, V, ahash::RandomState>;
pub type IndexMap<K, V> = indexmap::IndexMap<K, V, ahash::RandomState>;
15 changes: 15 additions & 0 deletions kclvm/runtime/src/stdlib/plugin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ lazy_static! {
> = Mutex::new(None);
}

/// KCL plugin module prefix
pub const PLUGIN_MODULE_PREFIX: &str = "kcl_plugin.";

#[no_mangle]
#[runtime_fn]
pub extern "C" fn kclvm_plugin_init(
Expand All @@ -46,6 +49,18 @@ pub unsafe extern "C" fn kclvm_plugin_invoke(
args: *const kclvm_value_ref_t,
kwargs: *const kclvm_value_ref_t,
) -> *const kclvm_value_ref_t {
let ctx_ref = mut_ptr_as_ref(ctx);
let method_ref = c2str(method);
let plugin_short_method = match method_ref.strip_prefix(PLUGIN_MODULE_PREFIX) {
Some(s) => s,
None => method_ref,
};
if let Some(func) = ctx_ref.plugin_functions.get(plugin_short_method) {
let args = ptr_as_ref(args);
let kwargs = ptr_as_ref(kwargs);
let result = func(ctx_ref, args, kwargs);
return result.unwrap().into_raw(ctx_ref);
}
let args_s = kclvm_value_to_json_value_with_null(ctx, args);
let kwargs_s = kclvm_value_to_json_value_with_null(ctx, kwargs);

Expand Down

0 comments on commit 3885840

Please sign in to comment.