Add WASI p2 support

examples/wasi-p2.rb

@@ -0,0 +1,17 @@
+require "wasmtime"
+
+engine = Wasmtime::Engine.new
+component = Wasmtime::Component::Component.from_file(engine, "spec/fixtures/wasi-debug-p2.wasm")
+
+linker = Wasmtime::Component::Linker.new(engine, wasi: true)
+
+wasi_config = Wasmtime::WasiConfig.new
+  .use_p2
+  .set_stdin_string("hi!")
+  .inherit_stdout
+  .inherit_stderr
+  .set_argv(ARGV)
+  .set_env(ENV)
+store = Wasmtime::Store.new(engine, wasi_config: wasi_config)
+
+Wasmtime::Component::WasiCommand.new(store, component, linker).call_run(store)

ext/src/ruby_api/component.rs

@@ -2,6 +2,7 @@ mod convert;
 mod func;
 mod instance;
 mod linker;
+mod wasi_command;
 
 use super::root;
 use magnus::{
@@ -13,6 +14,7 @@ use wasmtime::component::Component as ComponentImpl;
 
 pub use func::Func;
 pub use instance::Instance;
+pub use wasi_command::WasiCommand;
 
 pub fn component_namespace(ruby: &Ruby) -> RModule {
     static COMPONENT_NAMESPACE: Lazy<RModule> =
@@ -162,6 +164,7 @@ pub fn init(ruby: &Ruby) -> Result<(), Error> {
     instance::init(ruby, &namespace)?;
     func::init(ruby, &namespace)?;
     convert::init(ruby)?;
+    wasi_command::init(ruby, &namespace)?;
 
     Ok(())
 }

ext/src/ruby_api/component/linker.rs

@@ -1,19 +1,30 @@
 use super::{Component, Instance};
 use crate::{
-    err,
+    define_rb_intern, err,
     ruby_api::{
         store::{StoreContextValue, StoreData},
         Engine, Module, Store,
     },
 };
-use std::{borrow::BorrowMut, cell::RefCell};
+use std::{
+    borrow::BorrowMut,
+    cell::{RefCell, RefMut},
+};
 
 use crate::error;
 use magnus::{
     class, function, gc::Marker, method, r_string::RString, scan_args, typed_data::Obj,
     DataTypeFunctions, Error, Module as _, Object, RModule, Ruby, TryConvert, TypedData, Value,
 };
 use wasmtime::component::{Linker as LinkerImpl, LinkerInstance as LinkerInstanceImpl};
+use wasmtime_wasi::{
+    p2::{IoView, WasiCtx, WasiView},
+    ResourceTable,
+};
+
+define_rb_intern!(
+    WASI => "wasi",
+);
 
 /// @yard
 /// @rename Wasmtime::Component::Linker
@@ -23,6 +34,7 @@ use wasmtime::component::{Linker as LinkerImpl, LinkerInstance as LinkerInstance
 pub struct Linker {
     inner: RefCell<LinkerImpl<StoreData>>,
     refs: RefCell<Vec<Value>>,
+    has_wasi: bool,
 }
 unsafe impl Send for Linker {}
 
@@ -34,17 +46,36 @@ impl DataTypeFunctions for Linker {
 
 impl Linker {
     /// @yard
-    /// @def new(engine)
+    /// @def new(engine, wasi: false)
     /// @param engine [Engine]
+    /// @param wasi [Boolean] Whether WASI should be defined in this Linker. Defaults to false.
     /// @return [Linker]
-    pub fn new(engine: &Engine) -> Result<Self, Error> {
-        let linker = LinkerImpl::new(engine.get());
+    pub fn new(args: &[Value]) -> Result<Self, Error> {
+        let args = scan_args::scan_args::<(&Engine,), (), (), (), _, ()>(args)?;
+        let kw = scan_args::get_kwargs::<_, (), (Option<bool>,), ()>(args.keywords, &[], &[*WASI])?;
+        let (engine,) = args.required;
+        let wasi = kw.optional.0.unwrap_or(false);
+
+        let mut linker: LinkerImpl<StoreData> = LinkerImpl::new(engine.get());
+        if wasi {
+            wasmtime_wasi::p2::add_to_linker_sync(&mut linker).map_err(|e| error!("{}", e))?
+        }
+
         Ok(Linker {
             inner: RefCell::new(linker),
             refs: RefCell::new(Vec::new()),
+            has_wasi: wasi,
         })
     }
 
+    pub(crate) fn inner_mut(&self) -> RefMut<LinkerImpl<StoreData>> {
+        self.inner.borrow_mut()
+    }
+
+    pub(crate) fn has_wasi(&self) -> bool {
+        self.has_wasi
+    }
+
     /// @yard
     /// @def root
     /// Define items in the root of this {Linker}.
@@ -105,6 +136,16 @@ impl Linker {
         store: Obj<Store>,
         component: &Component,
     ) -> Result<Instance, Error> {
+        if rb_self.has_wasi && !store.context().data().has_wasi_p2_ctx() {
+            return err!(
+                "Store is missing WASI p2 configuration.\n\n\
+                When using `wasi: true`, the Store given to\n\
+                `Linker#instantiate` must have a WASI p2 configuration.\n\
+                To fix this, provide the `wasi_config` when creating the Store:\n\
+                    Wasmtime::Store.new(engine, wasi_config: WasiConfig.new.use_p2)"
+            );
+        }
+
         let inner = rb_self.inner.borrow();
         inner
             .instantiate(store.context_mut(), component.get())
@@ -230,7 +271,7 @@ impl<'a> LinkerInstance<'a> {
 
 pub fn init(_ruby: &Ruby, namespace: &RModule) -> Result<(), Error> {
     let linker = namespace.define_class("Linker", class::object())?;
-    linker.define_singleton_method("new", function!(Linker::new, 1))?;
+    linker.define_singleton_method("new", function!(Linker::new, -1))?;
     linker.define_method("root", method!(Linker::root, 0))?;
     linker.define_method("instance", method!(Linker::instance, 1))?;
     linker.define_method("instantiate", method!(Linker::instantiate, 2))?;

ext/src/ruby_api/component/wasi_command.rs

@@ -0,0 +1,61 @@
+use magnus::{
+    class, function, method, module::Module, typed_data::Obj, DataTypeFunctions, Error, Object,
+    RModule, Ruby,
+};
+use wasmtime_wasi::p2::bindings::sync::Command;
+
+use crate::{
+    err, error,
+    ruby_api::component::{linker::Linker, Component},
+    Store,
+};
+
+#[magnus::wrap(class = "Wasmtime::Component::WasiCommand", size, free_immediately)]
+pub struct WasiCommand {
+    command: Command,
+}
+
+impl WasiCommand {
+    /// @yard
+    /// @def new(store, component, linker)
+    /// @param store [Store]
+    /// @param component [Component]
+    /// @param linker [Linker]
+    /// @return [WasiCommand]
+    pub fn new(store: &Store, component: &Component, linker: &Linker) -> Result<Self, Error> {
+        if linker.has_wasi() && !store.context().data().has_wasi_p2_ctx() {
+            return err!(
+                "Store is missing WASI p2 configuration.\n\n\
+                When using `wasi: true`, the Store given to\n\
+                `Linker#instantiate` must have a WASI p2 configuration.\n\
+                To fix this, provide the `wasi_config` when creating the Store:\n\
+                    Wasmtime::Store.new(engine, wasi_config: WasiConfig.new.use_p2)"
+            );
+        }
+        let command =
+            Command::instantiate(store.context_mut(), component.get(), &linker.inner_mut())
+                .map_err(|e| error!("{e}"))?;
+        Ok(Self { command })
+    }
+
+    /// @yard
+    /// @def call_run(store)
+    /// @param store [Store]
+    /// @return [nil]
+    pub fn call_run(_ruby: &Ruby, rb_self: Obj<Self>, store: &Store) -> Result<(), Error> {
+        rb_self
+            .command
+            .wasi_cli_run()
+            .call_run(store.context_mut())
+            .map_err(|err| error!("{err}"))?
+            .map_err(|_| error!("Error running `run`"))
+    }
+}
+
+pub fn init(_ruby: &Ruby, namespace: &RModule) -> Result<(), Error> {
+    let linker = namespace.define_class("WasiCommand", class::object())?;
+    linker.define_singleton_method("new", function!(WasiCommand::new, 3))?;
+    linker.define_method("call_run", method!(WasiCommand::call_run, 1))?;
+
+    Ok(())
+}

ext/src/ruby_api/linker.rs

@@ -54,7 +54,7 @@ impl Linker {
 
         let mut inner: LinkerImpl<StoreData> = LinkerImpl::new(engine.get());
         if wasi {
-            wasmtime_wasi::preview1::add_to_linker_sync(&mut inner, |s| s.wasi_ctx_mut())
+            wasmtime_wasi::preview1::add_to_linker_sync(&mut inner, |s| s.wasi_p1_ctx_mut())
                 .map_err(|e| error!("{}", e))?
         }
         Ok(Self {
@@ -280,7 +280,7 @@ impl Linker {
     /// @param mod [Module]
     /// @return [Instance]
     pub fn instantiate(&self, store: Obj<Store>, module: &Module) -> Result<Instance, Error> {
-        if self.has_wasi && !store.context().data().has_wasi_ctx() {
+        if self.has_wasi && !store.context().data().has_wasi_p1_ctx() {
             return err!(
                 "Store is missing WASI configuration.\n\n\
                 When using `wasi: true`, the Store given to\n\

ext/src/ruby_api/store.rs

@@ -1,5 +1,6 @@
 use super::errors::wasi_exit_error;
 use super::{caller::Caller, engine::Engine, root, trap::Trap};
+use crate::ruby_api::wasi_config::WasiVersion;
 use crate::{define_rb_intern, error, WasiConfig};
 use magnus::value::StaticSymbol;
 use magnus::{
@@ -19,8 +20,9 @@ use wasmtime::{
     AsContext, AsContextMut, ResourceLimiter, Store as StoreImpl, StoreContext, StoreContextMut,
     StoreLimits, StoreLimitsBuilder,
 };
+use wasmtime_wasi::p2::{IoView, WasiCtx, WasiView};
 use wasmtime_wasi::preview1::WasiP1Ctx;
-use wasmtime_wasi::I32Exit;
+use wasmtime_wasi::{I32Exit, ResourceTable};
 
 define_rb_intern!(
     WASI_CONFIG => "wasi_config",
@@ -29,23 +31,31 @@ define_rb_intern!(
 
 pub struct StoreData {
     user_data: Value,
-    wasi: Option<WasiP1Ctx>,
+    wasi_p1: Option<WasiP1Ctx>,
+    wasi_p2: Option<WasiCtx>,
     refs: Vec<Value>,
     last_error: Option<Error>,
     store_limits: TrackingResourceLimiter,
+    resource_table: ResourceTable,
 }
 
 impl StoreData {
     pub fn user_data(&self) -> Value {
         self.user_data
     }
 
-    pub fn has_wasi_ctx(&self) -> bool {
-        self.wasi.is_some()
+    pub fn has_wasi_p1_ctx(&self) -> bool {
+        self.wasi_p1.is_some()
     }
 
-    pub fn wasi_ctx_mut(&mut self) -> &mut WasiP1Ctx {
-        self.wasi.as_mut().expect("Store must have a WASI context")
+    pub fn has_wasi_p2_ctx(&self) -> bool {
+        self.wasi_p2.is_some()
+    }
+
+    pub fn wasi_p1_ctx_mut(&mut self) -> &mut WasiP1Ctx {
+        self.wasi_p1
+            .as_mut()
+            .expect("Store must have a WASI context")
     }
 
     pub fn retain(&mut self, value: Value) {
@@ -148,7 +158,14 @@ impl Store {
         let (user_data,) = args.optional;
         let user_data = user_data.unwrap_or_else(|| ().into_value());
         let wasi_config = kw.optional.0;
-        let wasi = wasi_config.map(|config| config.build(&ruby)).transpose()?;
+
+        let (wasi_p1, wasi_p2) = match wasi_config {
+            Some(config) => match config.wasi_version() {
+                WasiVersion::P1 => (Some(config.build_p1(&ruby)?), None),
+                WasiVersion::P2 => (None, Some(config.build_p2(&ruby)?)),
+            },
+            None => (None, None),
+        };
 
         let limiter = match kw.optional.1 {
             None => StoreLimitsBuilder::new(),
@@ -160,10 +177,12 @@ impl Store {
         let eng = engine.get();
         let store_data = StoreData {
             user_data,
-            wasi,
+            wasi_p1,
+            wasi_p2,
             refs: Default::default(),
             last_error: Default::default(),
             store_limits: limiter,
+            resource_table: Default::default(),
         };
         let store = Self {
             inner: UnsafeCell::new(StoreImpl::new(eng, store_data)),
@@ -485,3 +504,17 @@ impl ResourceLimiter for TrackingResourceLimiter {
         self.inner.memories()
     }
 }
+
+impl WasiView for StoreData {
+    fn ctx(&mut self) -> &mut WasiCtx {
+        self.wasi_p2
+            .as_mut()
+            .expect("Should have WASI p2 context defined if using WASI p2")
+    }
+}
+
+impl IoView for StoreData {
+    fn table(&mut self) -> &mut ResourceTable {
+        &mut self.resource_table
+    }
+}