feat(air): embedded Starlark script support

Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "crates/air-lib/interpreter-interface",
     "crates/air-lib/interpreter-sede",
     "crates/air-lib/interpreter-signatures",
+    "crates/air-lib/interpreter-starlark",
     "crates/air-lib/interpreter-value",
     "crates/air-lib/lambda/ast",
     "crates/air-lib/lambda/parser",

air/Cargo.toml

@@ -24,6 +24,7 @@ air-interpreter-sede = { version = "0.1.0", path = "../crates/air-lib/interprete
 air-interpreter-signatures = { version = "0.1.7", path = "../crates/air-lib/interpreter-signatures", features = ["rkyv"] }
 air-interpreter-value = { version = "0.1.0", path = "../crates/air-lib/interpreter-value" }
 air-interpreter-interface = { version = "0.19.0", path = "../crates/air-lib/interpreter-interface", default-features = false }
+air-interpreter-starlark = { version = "0.1.0", path = "../crates/air-lib/interpreter-starlark", default-features = false }
 air-log-targets = { version = "0.1.0", path = "../crates/air-lib/log-targets" }
 air-lambda-ast = { version = "0.1.0", path = "../crates/air-lib/lambda/ast" }
 air-lambda-parser = { version = "0.1.0", path = "../crates/air-lib/lambda/parser" }

air/src/execution_step/errors/catchable_errors.rs

@@ -101,6 +101,9 @@ pub enum CatchableError {
     /// Stream map related errors.
     #[error(transparent)]
     StreamMapError(#[from] StreamMapError),
+
+    #[error("Starlark error: {0}")]
+    StarlarkError(air_interpreter_starlark::ExecutionError),
 }
 
 impl From<LambdaError> for Rc<CatchableError> {

air/src/execution_step/errors/uncatchable_errors.rs

@@ -143,6 +143,9 @@ pub enum UncatchableError {
 
     #[error("failed to serialize call arguments {0}")]
     CallArgumentsSerializationFailed(<CallArgumentsRepr as Representation>::SerializeError),
+
+    #[error("Starlark error: {0}")]
+    StarlarkError(air_interpreter_starlark::ExecutionError),
 }
 
 impl ToErrorCode for UncatchableError {

air/src/execution_step/instructions/embed.rs

@@ -0,0 +1,113 @@
+/*
+ * AquaVM Workflow Engine
+ *
+ * Copyright (C) 2024 Fluence DAO
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation version 3 of the
+ * License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+use std::rc::Rc;
+
+use air_interpreter_data::Provenance;
+use air_interpreter_starlark::execute as starlark_execute;
+use air_parser::ast::Embed;
+use air_parser::ast::ImmutableValue;
+use polyplets::SecurityTetraplet;
+
+use super::fail::fail_with_error_object;
+use super::ExecutableInstruction;
+use super::ExecutionCtx;
+use super::ExecutionResult;
+use super::TraceHandler;
+use crate::execution_step::errors::Joinable as _;
+use crate::execution_step::resolver::Resolvable as _;
+use crate::execution_step::LiteralAggregate;
+use crate::execution_step::RcSecurityTetraplets;
+use crate::execution_step::ValueAggregate;
+use crate::joinable;
+use crate::CatchableError;
+use crate::ExecutionError;
+use crate::JValue;
+use crate::UncatchableError;
+
+impl<'i> ExecutableInstruction<'i> for Embed<'i> {
+    fn execute(&self, exec_ctx: &mut ExecutionCtx<'i>, _trace_ctx: &mut TraceHandler) -> ExecutionResult<()> {
+        let args = joinable!(collect_args(&self.args, exec_ctx), exec_ctx, ())?;
+
+        let output_value = starlark_execute(self.script, args).map_err(classify_starlark_error)?;
+        match output_value {
+            Ok(value) => maybe_set_output_value(&self.output, value, exec_ctx),
+            Err((error_code, error_message)) => {
+                let error_object = crate::execution_step::execution_context::error_from_raw_fields_w_peerid(
+                    error_code.into(),
+                    &error_message,
+                    &self.to_string(),
+                    exec_ctx.run_parameters.init_peer_id.as_ref(),
+                );
+                let literal_tetraplet =
+                    SecurityTetraplet::literal_tetraplet(exec_ctx.run_parameters.init_peer_id.as_ref());
+                let literal_tetraplet = Rc::new(literal_tetraplet);
+                // in (fail x y), x and y are always literals
+                let provenance = Provenance::literal();
+
+                fail_with_error_object(exec_ctx, error_object, Some(literal_tetraplet), provenance)
+            }
+        }
+    }
+}
+
+fn collect_args(
+    args: &[ImmutableValue<'_>],
+    exec_ctx: &ExecutionCtx<'_>,
+) -> ExecutionResult<Vec<(JValue, RcSecurityTetraplets)>> {
+    let mut result = Vec::with_capacity(args.len());
+
+    for instruction_value in args {
+        let (arg, tetraplet, _) = instruction_value.resolve(exec_ctx)?;
+        result.push((arg, tetraplet));
+    }
+    Ok(result)
+}
+
+fn classify_starlark_error(err: air_interpreter_starlark::ExecutionError) -> ExecutionError {
+    use air_interpreter_starlark::ExecutionError::*;
+
+    match err {
+        // TODO perhaps, Other should be uncatchable
+        Value(_) | Function(_) | Other(_) => ExecutionError::Catchable(CatchableError::StarlarkError(err).into()),
+        Scope(_) | Lexer(_) | Internal(_) => ExecutionError::Uncatchable(UncatchableError::StarlarkError(err)),
+    }
+}
+
+fn maybe_set_output_value(
+    embed_output_value: &air_parser::ast::EmbedOutputValue<'_>,
+    result_value: air_interpreter_value::JValue,
+    exec_ctx: &mut ExecutionCtx<'_>,
+) -> Result<(), ExecutionError> {
+    match embed_output_value {
+        air_parser::ast::EmbedOutputValue::Scalar(scalar) => {
+            // TODO for now, we treat value produced by Starlark as a literal, as it has to be
+            // same on every peer.
+            let result = ValueAggregate::from_literal_result(LiteralAggregate::new(
+                result_value,
+                exec_ctx.run_parameters.init_peer_id.clone(),
+                0.into(), // TODO is it correct?
+            ));
+
+            exec_ctx.scalars.set_scalar_value(scalar.name, result)?;
+        }
+        air_parser::ast::EmbedOutputValue::None => {}
+    }
+    Ok(())
+}

air/src/execution_step/instructions/fail.rs

@@ -148,7 +148,7 @@ fn fail_with_error(exec_ctx: &mut ExecutionCtx<'_>) -> ExecutionResult<()> {
     result
 }
 
-fn fail_with_error_object(
+pub(crate) fn fail_with_error_object(
     exec_ctx: &mut ExecutionCtx<'_>,
     error: JValue,
     tetraplet: Option<RcSecurityTetraplet>,

air/src/execution_step/instructions/mod.rs

@@ -25,6 +25,7 @@ mod canon_map;
 mod canon_stream_map_scalar;
 mod canon_utils;
 mod compare_matchable;
+mod embed;
 mod fail;
 mod fold;
 mod fold_scalar;
@@ -95,6 +96,7 @@ impl<'i> ExecutableInstruction<'i> for Instruction<'i> {
             Instruction::Xor(xor) => execute!(self, xor, exec_ctx, trace_ctx),
             Instruction::Match(match_) => execute!(self, match_, exec_ctx, trace_ctx),
             Instruction::MisMatch(mismatch) => execute!(self, mismatch, exec_ctx, trace_ctx),
+            Instruction::Embed(embed) => execute!(self, embed, exec_ctx, trace_ctx),
 
             Instruction::Error => unreachable!("should not execute if parsing succeeded. QED."),
         }

air/tests/test_module/instructions/embed.rs

@@ -0,0 +1,217 @@
+/*
+ * AquaVM Workflow Engine
+ *
+ * Copyright (C) 2024 Fluence DAO
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation version 3 of the
+ * License.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+use air::{CatchableError, UncatchableError};
+use air_interpreter_starlark::ExecutionError as StarlarkExecutionError;
+use air_test_utils::prelude::*;
+
+#[tokio::test]
+async fn embed_basic() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+        (seq
+            (embed []
+#"
+"a string\nwith escape"
+"#
+                var)
+            (call %init_peer_id% ("" "") [var] result_name))"##;
+
+    let result = checked_call_vm!(vm, <_>::default(), script, "", "");
+    assert!(result.next_peer_pks.is_empty());
+
+    let expected_trace = vec![scalar!(
+        json!("a string\nwith escape"),
+        peer = "",
+        args = ["a string\nwith escape"]
+    )];
+
+    let trace = trace_from_result(&result);
+    assert_eq!(&*trace, expected_trace);
+}
+
+#[tokio::test]
+async fn embed_args() {
+    let init_peer_id = "my_id";
+    let mut vm = create_avm(echo_call_service(), init_peer_id).await;
+
+    let script = r##"
+        (seq
+           (call %init_peer_id% ("myservice" "myfunc") [42] arg)
+           (seq
+               (embed [arg]
+#"
+t = get_tetraplet(0)[0]
+"{}: {}/{}:{}".format(get_value(0), t.peer_pk, t.service_id, t.function_name)
+"#
+                      var)
+               (call %init_peer_id% ("" "") [var] result_name)))"##;
+
+    let run_params = TestRunParameters::from_init_peer_id(init_peer_id);
+    let result = checked_call_vm!(vm, run_params, script, "", "");
+    assert!(result.next_peer_pks.is_empty());
+
+    let expected_val = format!("42: {init_peer_id}/myservice:myfunc");
+    let expected_trace = vec![
+        scalar!(
+            json!(42),
+            peer = init_peer_id,
+            service = "myservice",
+            function = "myfunc",
+            args = [42]
+        ),
+        scalar!(json!(expected_val), peer = init_peer_id, args = [expected_val]),
+    ];
+
+    let trace = trace_from_result(&result);
+    assert_eq!(&*trace, expected_trace);
+}
+
+#[tokio::test]
+async fn embed_error_fail() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+        (xor
+            (embed []
+#"
+fail(42, "my message")
+"#
+                var)
+            (call %init_peer_id% ("" "") [%last_error%.$.error_code %last_error%.$.message] result_name))"##;
+
+    let result = checked_call_vm!(vm, <_>::default(), script, "", "");
+    assert!(result.next_peer_pks.is_empty());
+
+    let expected_trace = vec![scalar!(json!(42), peer = "", args = [json!(42), json!("my message")])];
+
+    let trace = trace_from_result(&result);
+    assert_eq!(&*trace, expected_trace);
+}
+
+#[tokio::test]
+async fn embed_error_value() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+       (embed []
+#"
+42 + "string"
+"#
+              var)"##;
+
+    let result = call_vm!(vm, <_>::default(), script, "", "");
+    let expected_error = CatchableError::StarlarkError(StarlarkExecutionError::Value(
+        "error: Operation `+` not supported for types `int` and `string`\n --> dummy.star:2:1\n  |\n2 | 42 + \"string\"\n  | ^^^^^^^^^^^^^\n  |\n".to_owned(),
+    ));
+    assert_error_eq!(&result, expected_error);
+}
+
+// TODO 42.length gives Other, and it is a problem
+#[tokio::test]
+async fn embed_error_lexer() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+       (embed []
+#"
+"an unterminated string
+"#
+              var)"##;
+
+    let result = call_vm!(vm, <_>::default(), script, "", "");
+    let expected_error = UncatchableError::StarlarkError(StarlarkExecutionError::Lexer(
+        "Parse error: unfinished string literal".to_owned(),
+    ));
+    assert_error_eq!(&result, expected_error);
+}
+
+#[tokio::test]
+async fn embed_with_join_behavior() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+        (par
+            (call "other_peer" ("" "") [] var)
+            (seq
+                (embed [var] #"var + var"# var2)
+                (call %init_peer_id% ("" "") [var2])))"##;
+
+    let result = call_vm!(vm, <_>::default(), script, "", "");
+
+    assert_eq!(result.error_message, "");
+    assert_eq!(result.ret_code, 0);
+
+    let trace = trace_from_result(&result);
+    assert_eq!(trace.len(), 2);
+}
+
+#[tokio::test]
+async fn embed_zip_reverse() {
+    let mut vm = create_avm(echo_call_service(), "").await;
+
+    let script = r##"
+        (seq
+            (seq
+               (seq
+                   (ap 1 $stream)
+                   (ap 2 $stream))
+            (canon %init_peer_id% $stream #canon))
+            (seq
+                (embed [#canon #canon]
+#"
+def main():
+    v1 = get_value(0)
+    v2 = get_value(1)
+
+    if get_tetraplet(0)[0].peer_pk != get_tetraplet(1)[0].peer_pk:
+        fail(42, 'tetraplet peer_pk mismatch')
+
+    return list(zip(v1, reversed(v2)))
+
+main()
+"#
+                       var2)
+                (call %init_peer_id% ("" "") [var2] var3)))"##;
+
+    let result = call_vm!(vm, <_>::default(), script, "", "");
+
+    assert_eq!(result.error_message, "", "{}", result.error_message);
+    assert_eq!(result.ret_code, 0);
+
+    let expected_trace = vec![
+        ap(0),
+        ap(0),
+        canon(json!({
+            "tetraplet":  {"function_name": "", "lens": "", "peer_pk": "", "service_id": ""},
+            "values": [{
+                "tetraplet": {"function_name": "", "lens": "", "peer_pk": "", "service_id": ""},
+                "result": 1
+            }, {
+                "tetraplet": {"function_name": "", "lens": "", "peer_pk": "", "service_id": ""},
+                "result": 2
+            }]
+        })),
+        scalar!(json!([(1, 2), (2, 1)]), args = [json!([(1, 2), (2, 1)])]),
+    ];
+    let data = data_from_result(&result);
+    let trace = trace_from_result(&result);
+    assert_eq!(&*trace, expected_trace, "{:#?}", data.cid_info);
+}