Zed 扩展的生命周期：Rust、WIT、Wasm

$ mkdir output
$ zed-extension --source-dir ./metals-zed --output ./output --scratch-dir $(mktemp -d)
info: downloading component 'rust-std' for 'wasm32-wasip1'
info: installing component 'rust-std' for 'wasm32-wasip1'
 
$ ls -1 ./output
archive.tar.gz
manifest.json

$ cat ./output/manifest.json | jq .
{
  "name": "Scala",
  "version": "0.1.0",
  "description": "Scala support.",
  "authors": [
    "Igal Tabachnik <[email protected]>",
    "Jamie Thompson <[email protected]>",
 
  ],
  "repository": "https://github.com/scalameta/metals-zed",
  "schema_version": 1,
  "wasm_api_version": "0.0.6"
}

$ tar -ztf ./output/archive.tar.gz
./
./extension.toml
./extension.wasm
./languages/
./grammars/
./grammars/scala.wasm
./languages/scala/
./languages/scala/outline.scm
./languages/scala/indents.scm
./languages/scala/highlights.scm
./languages/scala/config.toml
./languages/scala/overrides.scm
./languages/scala/injections.scm
./languages/scala/runnables.scm
./languages/scala/brackets.scm

use zed_extension_api::{self as zed, Result};
 
struct MyExtension;
 
impl MyExtension {
    const SERVER_BINARY_NAME: &'static str = "my-language-server";
}
 
impl zed::Extension for MyExtension {
    fn new() -> Self {
        Self
    }
 
    fn language_server_command(
        &mut self,
        _language_server_id: &zed::LanguageServerId,
        worktree: &zed::Worktree,
    ) -> Result<zed::Command> {
        let path = worktree
            .which(Self::SERVER_BINARY_NAME)
            .ok_or_else(|| format!("Could not find {} binary", Self::SERVER_BINARY_NAME))?;
 
        Ok(zed::Command {
            command: path,
            args: vec!["--use-printf-debugging".to_string()],
            env: worktree.shell_env(),
        })
    }
}
 
zed::register_extension!(MyExtension);

[package]
name = "test-extension"
version = "0.1.0"
edition = "2021"
 
[lib]
crate-type = ["cdylib"]
 
[dependencies]
zed_extension_api = "0.0.6"

// Excerpt from `extension_api.rs`
 
/// A Zed extension.
pub trait Extension: Send + Sync {
    /// Returns a new instance of the extension.
    fn new() -> Self
    where
        Self: Sized;
 
    /// Returns the command used to start the language server for the specified
    /// language.
    fn language_server_command(
        &mut self,
        _language_server_id: &LanguageServerId,
        _worktree: &Worktree,
    ) -> Result<Command> {
        Err("`language_server_command` not implemented".to_string())
    }
 
    /// Returns the initialization options to pass to the specified language server.
    fn language_server_initialization_options(
        &mut self,
        _language_server_id: &LanguageServerId,
        _worktree: &Worktree,
    ) -> Result<Option<serde_json::Value>> {
        Ok(None)
    }
 
    // [...]
}

/// The ID of a language server.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
pub struct LanguageServerId(String);

// Excerpt from extension_api.rs
 
mod wit {
    wit_bindgen::generate!({
        skip: ["init-extension"],
        path: "./wit/since_v0.2.0",
    });
}

wit_bindgen::generate!({
    skip: ["init-extension"],
    path: "./wit/since_v0.2.0",
});

package zed:extension;
 
world extension {
    /// A command.
    record command {
        /// The command to execute.
        command: string,
        /// The arguments to pass to the command.
        args: list<string>,
        /// The environment variables to set for the command.
        env: env-vars,
    }
 
    /// A Zed worktree.
    resource worktree {
        /// Returns the ID of the worktree.
        id: func() -> u64;
        /// Returns the root path of the worktree.
        root-path: func() -> string;
        /// Returns the textual contents of the specified file in the worktree.
        read-text-file: func(path: string) -> result<string, string>;
        /// Returns the path to the given binary name, if one is present on the `$PATH`.
        which: func(binary-name: string) -> option<string>;
        /// Returns the current shell environment.
        shell-env: func() -> env-vars;
    }
 
    /// Returns the command used to start up the language server.
    export language-server-command: func(language-server-id: string, worktree: borrow<worktree>) -> result<command, string>;
 
    /// [... other definitions ...]
}

fn language_server_command(
    &mut self,
    language_server_id: &LanguageServerId,
    worktree: &Worktree,
) -> Result<Command> {
    // ...
}

pub struct LanguageServerId(String);

export language-server-command: func(language-server-id: string, worktree: borrow<worktree>) -> result<command, string>;

#[export_name = "language-server-command"]
unsafe extern "C" fn export_language_server_command(
    arg0: *mut u8,
    arg1: usize,
    arg2: i32,
) -> *mut u8 {
    self::_export_language_server_command_cabi::<Component>(arg0, arg1, arg2)
}
 
unsafe fn _export_language_server_command_cabi<T: Guest>(
    arg0: *mut u8,
    arg1: usize,
    arg2: i32,
) -> *mut u8 {
    let handle1;
    let len0 = arg1;
    let bytes0 = _rt::Vec::from_raw_parts(arg0.cast(), len0, len0);
    let result2 = T::language_server_command(
        _rt::string_lift(bytes0),
        {
            handle1 = Worktree::from_handle(arg2 as u32);
            &handle1
        },
    );
 
    // [... a lot more code that only deals with numbers ...]
}

// Extract from extension_store.rs
 
let mut path = root_dir.clone();
path.extend([extension.manifest.clone().id.as_ref(), "extension.wasm"]);
let mut wasm_file = fs
    .open_sync(&path)
    .await
    .context("failed to open wasm file")?;
 
let mut wasm_bytes = Vec::new();
wasm_file
    .read_to_end(&mut wasm_bytes)
    .context("failed to read wasm")?;
 
wasm_host
    .load_extension(
        wasm_bytes,
        extension.manifest.clone().clone(),
        cx.background_executor().clone(),
    )
    .await
    .with_context(|| {
        format!("failed to load wasm extension {}", extension.manifest.id)
    })

// Excerpt from `wasm_host.rs`
 
pub fn load_extension(
    self: &Arc<Self>,
    wasm_bytes: Vec<u8>,
    manifest: Arc<ExtensionManifest>,
    executor: BackgroundExecutor,
) -> Task<Result<WasmExtension>> {
    let this = self.clone();
    executor.clone().spawn(async move {
        let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
 
        let component = Component::from_binary(&this.engine, &wasm_bytes)
            .context("failed to compile wasm component")?;
 
        let mut store = wasmtime::Store::new(
            &this.engine,
            WasmState {
                ctx: this.build_wasi_ctx(&manifest).await?,
                manifest: manifest.clone(),
                table: ResourceTable::new(),
                host: this.clone(),
            },
        );
 
        let mut extension = Extension::instantiate_async(
            &mut store,
            this.release_channel,
            zed_api_version,
            &component,
        )
        .await?;
 
        extension
            .call_init_extension(&mut store)
            .await
            .context("failed to initialize wasm extension")?;
 
        let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
        executor
            .spawn(async move {
                while let Some(call) = rx.next().await {
                    (call)(&mut extension, &mut store).await;
                }
            })
            .detach();
 
        Ok(WasmExtension {
            manifest,
            tx,
            zed_api_version,
        })
    })
}

// Excerpt from wasm_host/wit.rs
 
impl Extension {
    pub async fn instantiate_async(
        store: &mut Store<WasmState>,
        release_channel: ReleaseChannel,
        version: SemanticVersion,
        component: &Component,
    ) -> Result<Self> {
        // Note: The release channel can be used to stage a new version of the extension API.
        let allow_latest_version = match release_channel {
            ReleaseChannel::Dev | ReleaseChannel::Nightly => true,
            ReleaseChannel::Stable | ReleaseChannel::Preview => false,
        };
 
        if allow_latest_version && version >= latest::MIN_VERSION {
            let extension =
                latest::Extension::instantiate_async(store, component, latest::linker())
                    .await
                    .context("failed to instantiate wasm extension")?;
            Ok(Self::V020(extension))
        } else {
            // [... code that handles other versions ...]
        }
    }
}

wasmtime::component::bindgen!({
    async: true,
    trappable_imports: true,
    path: "../extension_api/wit/since_v0.2.0",
    with: {
         "worktree": ExtensionWorktree,
         "key-value-store": ExtensionKeyValueStore,
         "zed:extension/http-client/http-response-stream": ExtensionHttpResponseStream
    },
});

// extension_lsp_adapter.rs, simplified
 
struct ExtensionLspAdapter {
    extension: WasmExtension,
    language_server_id: LanguageServerName,
    config: LanguageServerConfig,
    host: Arc<WasmHost>,
}
 
#[async_trait(?Send)]
impl LspAdapter for ExtensionLspAdapter {
    fn get_language_server_command<'a>(
        self: Arc<Self>,
        delegate: Arc<dyn LspAdapterDelegate>,
        _: LanguageServerBinaryOptions,
        _: futures::lock::MutexGuard<'a, Option<LanguageServerBinary>>,
        _: &'a mut AsyncAppContext,
    ) -> Pin<Box<dyn 'a + Future<Output = Result<LanguageServerBinary>>>> {
        async move {
            let command = self
                .extension
                .call({
                    let this = self.clone();
                    |extension, store| {
                        async move {
                            let resource = store.data_mut().table().push(delegate)?;
                            let command = extension
                                .call_language_server_command(
                                    store,
                                    &this.language_server_id,
                                    &this.config,
                                    resource,
                                )
                                .await?
                                .map_err(|e| anyhow!("{}", e))?;
                            anyhow::Ok(command)
                        }
                        .boxed()
                    }
                })
                .await?;
 
            let path = self
                .host
                .path_from_extension(&self.extension.manifest.id, command.command.as_ref());
 
            // ...
 
            Ok(LanguageServerBinary {
                path,
                arguments: command.args.into_iter().map(|arg| arg.into()).collect(),
                env: Some(command.env.into_iter().collect()),
            })
        }
        .boxed_local()
    }
// [...]
}

pub async fn call_language_server_command(
    &self,
    store: &mut Store<WasmState>,
    language_server_id: &LanguageServerName,
    config: &LanguageServerConfig,
    resource: Resource<Arc<dyn LspAdapterDelegate>>,
) -> Result<Result<Command, String>> {
    match self {
        Extension::V020(ext) => {
            ext.call_language_server_command(store, &language_server_id.0, resource)
                .await
        }
        // [...]
    }
}