test(jit): prove provider and await continuations

Author: olwang

crates/rsscript-sdk/tests/native_jit_differential.rs

@@ -1,10 +1,21 @@
+use std::collections::BTreeMap;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+
 use rsscript_sdk::{
     artifact::ArtifactVerifier,
     compile::Compiler,
     operation::CancellationToken,
-    provider_api::ProviderRegistry,
+    provider_api::{
+        BlockingBehavior, CancellationBehavior, DataEffect, ExternalSymbol, FunctionSignature,
+        ParameterSignature, ProviderCallMode, ProviderDescriptor, ProviderError,
+        ProviderErrorMapping, ProviderFunction, ProviderFunctionDescriptor, ProviderRegistry,
+        RUNTIME_ABI_VERSION, ResourceCleanupContract, WireInterpreterFn, WireValue,
+    },
     report::ExecutionEngineTelemetry,
-    runtime::{ExecutionRequest, NativeCostModel, NativeJitOptions, RunLimits, Runtime},
+    runtime::{
+        ExecutionRequest, NativeCostModel, NativeJitOptions, RunLimits, Runtime, TracePolicy,
+    },
 };
 
 const CASES: &[(&str, &str)] = &[
@@ -32,6 +43,10 @@ const CASES: &[(&str, &str)] = &[
         "aggregate-continuation.rss",
         "struct AggregateBox { value: Int } fn main() -> Int { let boxed = AggregateBox(value: 13); let extracted = boxed.value; let a = extracted * 3; let b = a + 11; let c = b * 5; return c }",
     ),
+    (
+        "await-continuation.rss",
+        "async fn boundary(value: Int) -> Int { return value + 4 } async fn main() -> Int { let a = 7; let b = a * 3; let c = b + 11; task_group { async let pending = boundary(value: c); let d = await pending; let e = d * 5; let f = e - 9; let g = f + 2; return g } }",
+    ),
     (
         "native-list-write.rss",
         include_str!("../../../benchmarks/vm-jit/kernels/native_list_write_loop.rss"),
@@ -165,6 +180,18 @@ fn native_engine_matches_the_verified_interpreter_corpus() {
                 "scalar work after aggregate materialization must re-enter native code; entries={continuation_entries}, yields={continuation_yields}, barriers={native_barrier_counts:?}, missed={interpreted_native_work}"
             );
         }
+        if *file == "await-continuation.rss" {
+            assert!(
+                continuation_entries >= 2 && continuation_yields >= 2,
+                "scalar work around await must use native continuations; entries={continuation_entries}, yields={continuation_yields}, barriers={native_barrier_counts:?}"
+            );
+            assert!(
+                native_barrier_counts
+                    .get("await")
+                    .is_some_and(|count| *count >= 1),
+                "await must remain a VM-owned barrier"
+            );
+        }
         cases_with_native_entry +=
             usize::from(native_calls > 0 || osr_entries > 0 || continuation_entries > 0);
     }
@@ -252,3 +279,108 @@ fn bounded_step_accounting_matches_across_call_continuations() {
     };
     assert!(continuation_entries >= 2);
 }
+
+#[test]
+fn provider_barrier_executes_once_and_reenters_native() {
+    const SOURCE: &str = "module app\nuse host.math.*\nfn main() -> Int { let a = 7; let b = a * 3; let c = b + 11; let d = adjust(value: read c); let e = d * 5; let f = e - 9; let g = f + 2; return g }";
+    const INTERFACE: &str = "module host.math\npub fn adjust(value: read Int) -> Int\n";
+
+    let symbol = ExternalSymbol::new("host.math.adjust").expect("test symbol is valid");
+    let signature = FunctionSignature {
+        parameters: vec![ParameterSignature {
+            name: "value".into(),
+            effect: DataEffect::Read,
+            ty: "Int".into(),
+            retained: false,
+        }],
+        result: "Int".into(),
+        asynchronous: false,
+    };
+    let descriptor = ProviderDescriptor {
+        provider_id: "jit.test.math".into(),
+        provider_version: "1".into(),
+        supported_abi: vec![RUNTIME_ABI_VERSION],
+        record_layouts: Vec::new(),
+        variant_layouts: Vec::new(),
+        functions: vec![ProviderFunctionDescriptor {
+            symbol: symbol.clone(),
+            signature: signature.clone(),
+            entry: "adjust".into(),
+            call_mode: ProviderCallMode::Sync,
+            blocking: BlockingBehavior::NonBlocking,
+            cancellation: CancellationBehavior::NotApplicable,
+            thread_safe: true,
+            reentrant: true,
+            resource_cleanup: ResourceCleanupContract::None,
+            error_mapping: ProviderErrorMapping::StructuredV1,
+        }],
+    };
+    let calls = Arc::new(AtomicU64::new(0));
+    let provider_calls = Arc::clone(&calls);
+    let mut providers = ProviderRegistry::default();
+    providers
+        .register(
+            &descriptor,
+            BTreeMap::from([(
+                symbol,
+                ProviderFunction {
+                    signature,
+                    callable: WireInterpreterFn::new(move |args| match args.as_slice() {
+                        [WireValue::Int { value }] => {
+                            provider_calls.fetch_add(1, Ordering::SeqCst);
+                            Ok(WireValue::Int { value: value + 4 })
+                        }
+                        _ => Err(ProviderError::invalid_argument(
+                            "adjust expects one Int argument",
+                        )),
+                    }),
+                },
+            )]),
+        )
+        .expect("test Provider matches its descriptor");
+
+    let built = Compiler
+        .compile_with_interfaces(&[("main.rss", SOURCE)], &[("math.rssi", INTERFACE)])
+        .expect("provider continuation source compiles");
+    let admitted = ArtifactVerifier
+        .verify(built)
+        .expect("provider continuation artifact verifies")
+        .admit_trusted_input();
+    let linked = Runtime::new(providers)
+        .link(&admitted)
+        .expect("test Provider links");
+
+    let interpreter = linked.execute(ExecutionRequest::default().trace(TracePolicy::MetadataOnly));
+    assert_eq!(calls.load(Ordering::SeqCst), 1);
+    let native = linked.execute(
+        ExecutionRequest::default()
+            .trace(TracePolicy::MetadataOnly)
+            .native_jit(NativeJitOptions {
+                cost_model: NativeCostModel::Off,
+                ..NativeJitOptions::default()
+            }),
+    );
+    assert_eq!(calls.load(Ordering::SeqCst), 2);
+    assert_eq!(native.outcome(), interpreter.outcome());
+    assert_eq!(
+        native.usage.steps_consumed,
+        interpreter.usage.steps_consumed
+    );
+    assert_eq!(native.provider_call_traces.len(), 1);
+    let ExecutionEngineTelemetry::Native {
+        continuation_entries,
+        continuation_yields,
+        native_barrier_counts,
+        ..
+    } = native.telemetry.engine
+    else {
+        panic!("Provider continuation must report native telemetry");
+    };
+    assert!(continuation_entries >= 2);
+    assert!(continuation_yields >= 2);
+    assert!(
+        native_barrier_counts
+            .get("external_call")
+            .is_some_and(|count| *count >= 1)
+    );
+}

docs/spec/native-jit-contract.md

@@ -52,6 +52,11 @@ currently keep every continuation in the interpreter.
   coexist in the VM frame: continuation marshalling validates only the exact
   register footprint of the selected scalar region, so scalar work after an
   interpreter-materialized aggregate can re-enter native code safely.
+- Provider calls and `await` are exercised as normal mixed-mode boundaries by
+  interpreter/native differential tests. The VM executes each boundary exactly
+  once, preserves Provider traces and scheduler semantics, then probes the next
+  scalar continuation. Generated code never re-enters the interpreter or spans a
+  suspension.
 - Structural compilation work is bounded by `JitLimits` before Cranelift code
   generation. Instruction, register, CFG-edge, operand, analysis-word, deopt,
   memo-scope, callee, and recursive-group counts have deterministic limits.