1use crate::{
  2    ExtensionLibraryKind, ExtensionManifest, GrammarManifestEntry, build_debug_adapter_schema_path,
  3    parse_wasm_extension_version,
  4};
  5use anyhow::{Context as _, Result, bail};
  6use async_compression::futures::bufread::GzipDecoder;
  7use async_tar::Archive;
  8use futures::{AsyncReadExt, io::Cursor};
  9use heck::ToSnakeCase;
 10use http_client::{self, AsyncBody, HttpClient};
 11use serde::Deserialize;
 12use std::{
 13    env, fs, mem,
 14    path::{Path, PathBuf},
 15    process::Stdio,
 16    str::FromStr,
 17    sync::Arc,
 18};
 19use wasm_encoder::{ComponentSectionId, Encode as _, RawSection, Section as _};
 20use wasmparser::Parser;
 21
 22/// Currently, we compile with Rust's `wasm32-wasip2` target, which works with WASI `preview2` and the component model.
 23const RUST_TARGET: &str = "wasm32-wasip2";
 24
 25/// Compiling Tree-sitter parsers from C to WASM requires Clang 17, and a WASM build of libc
 26/// and clang's runtime library. The `wasi-sdk` provides these binaries.
 27///
 28/// Once Clang 17 and its wasm target are available via system package managers, we won't need
 29/// to download this.
 30const WASI_SDK_URL: &str = "https://github.com/WebAssembly/wasi-sdk/releases/download/wasi-sdk-25/";
 31const WASI_SDK_ASSET_NAME: Option<&str> = if cfg!(all(target_os = "macos", target_arch = "x86_64"))
 32{
 33    Some("wasi-sdk-25.0-x86_64-macos.tar.gz")
 34} else if cfg!(all(target_os = "macos", target_arch = "aarch64")) {
 35    Some("wasi-sdk-25.0-arm64-macos.tar.gz")
 36} else if cfg!(all(target_os = "linux", target_arch = "x86_64")) {
 37    Some("wasi-sdk-25.0-x86_64-linux.tar.gz")
 38} else if cfg!(all(target_os = "linux", target_arch = "aarch64")) {
 39    Some("wasi-sdk-25.0-arm64-linux.tar.gz")
 40} else if cfg!(all(target_os = "freebsd", target_arch = "x86_64")) {
 41    Some("wasi-sdk-25.0-x86_64-linux.tar.gz")
 42} else if cfg!(all(target_os = "freebsd", target_arch = "aarch64")) {
 43    Some("wasi-sdk-25.0-arm64-linux.tar.gz")
 44} else if cfg!(all(target_os = "windows", target_arch = "x86_64")) {
 45    Some("wasi-sdk-25.0-x86_64-windows.tar.gz")
 46} else {
 47    None
 48};
 49
 50pub struct ExtensionBuilder {
 51    cache_dir: PathBuf,
 52    pub http: Arc<dyn HttpClient>,
 53}
 54
 55pub struct CompileExtensionOptions {
 56    pub release: bool,
 57}
 58
 59#[derive(Deserialize)]
 60struct CargoToml {
 61    package: CargoTomlPackage,
 62}
 63
 64#[derive(Deserialize)]
 65struct CargoTomlPackage {
 66    name: String,
 67}
 68
 69impl ExtensionBuilder {
 70    pub fn new(http_client: Arc<dyn HttpClient>, cache_dir: PathBuf) -> Self {
 71        Self {
 72            cache_dir,
 73            http: http_client,
 74        }
 75    }
 76
 77    pub async fn compile_extension(
 78        &self,
 79        extension_dir: &Path,
 80        extension_manifest: &mut ExtensionManifest,
 81        options: CompileExtensionOptions,
 82    ) -> Result<()> {
 83        populate_defaults(extension_manifest, extension_dir)?;
 84
 85        if extension_dir.is_relative() {
 86            bail!(
 87                "extension dir {} is not an absolute path",
 88                extension_dir.display()
 89            );
 90        }
 91
 92        fs::create_dir_all(&self.cache_dir).context("failed to create cache dir")?;
 93
 94        if extension_manifest.lib.kind == Some(ExtensionLibraryKind::Rust) {
 95            log::info!("compiling Rust extension {}", extension_dir.display());
 96            self.compile_rust_extension(extension_dir, extension_manifest, options)
 97                .await
 98                .context("failed to compile Rust extension")?;
 99            log::info!("compiled Rust extension {}", extension_dir.display());
100        }
101
102        for (debug_adapter_name, meta) in &mut extension_manifest.debug_adapters {
103            let debug_adapter_schema_path =
104                extension_dir.join(build_debug_adapter_schema_path(debug_adapter_name, meta));
105
106            let debug_adapter_schema = fs::read_to_string(&debug_adapter_schema_path)
107                .with_context(|| {
108                    format!("failed to read debug adapter schema for `{debug_adapter_name}` from `{debug_adapter_schema_path:?}`")
109                })?;
110            _ = serde_json::Value::from_str(&debug_adapter_schema).with_context(|| {
111                format!("Debug adapter schema for `{debug_adapter_name}` (path: `{debug_adapter_schema_path:?}`) is not a valid JSON")
112            })?;
113        }
114        for (grammar_name, grammar_metadata) in &extension_manifest.grammars {
115            let snake_cased_grammar_name = grammar_name.to_snake_case();
116            if grammar_name.as_ref() != snake_cased_grammar_name.as_str() {
117                bail!(
118                    "grammar name '{grammar_name}' must be written in snake_case: {snake_cased_grammar_name}"
119                );
120            }
121
122            log::info!(
123                "compiling grammar {grammar_name} for extension {}",
124                extension_dir.display()
125            );
126            self.compile_grammar(extension_dir, grammar_name.as_ref(), grammar_metadata)
127                .await
128                .with_context(|| format!("failed to compile grammar '{grammar_name}'"))?;
129            log::info!(
130                "compiled grammar {grammar_name} for extension {}",
131                extension_dir.display()
132            );
133        }
134
135        log::info!("finished compiling extension {}", extension_dir.display());
136        Ok(())
137    }
138
139    async fn compile_rust_extension(
140        &self,
141        extension_dir: &Path,
142        manifest: &mut ExtensionManifest,
143        options: CompileExtensionOptions,
144    ) -> anyhow::Result<()> {
145        self.install_rust_wasm_target_if_needed().await?;
146
147        let cargo_toml_content = fs::read_to_string(extension_dir.join("Cargo.toml"))?;
148        let cargo_toml: CargoToml = toml::from_str(&cargo_toml_content)?;
149
150        log::info!(
151            "compiling Rust crate for extension {}",
152            extension_dir.display()
153        );
154        let output = util::command::new_smol_command("cargo")
155            .args(["build", "--target", RUST_TARGET])
156            .args(options.release.then_some("--release"))
157            .arg("--target-dir")
158            .arg(extension_dir.join("target"))
159            // WASI builds do not work with sccache and just stuck, so disable it.
160            .env("RUSTC_WRAPPER", "")
161            .current_dir(extension_dir)
162            .output()
163            .await
164            .context("failed to run `cargo`")?;
165        if !output.status.success() {
166            bail!(
167                "failed to build extension {}",
168                String::from_utf8_lossy(&output.stderr)
169            );
170        }
171
172        log::info!(
173            "compiled Rust crate for extension {}",
174            extension_dir.display()
175        );
176
177        let mut wasm_path = PathBuf::from(extension_dir);
178        wasm_path.extend([
179            "target",
180            RUST_TARGET,
181            if options.release { "release" } else { "debug" },
182            &cargo_toml
183                .package
184                .name
185                // The wasm32-wasip2 target normalizes `-` in package names to `_` in the resulting `.wasm` file.
186                .replace('-', "_"),
187        ]);
188        wasm_path.set_extension("wasm");
189
190        log::info!(
191            "encoding wasm component for extension {}",
192            extension_dir.display()
193        );
194
195        let component_bytes = fs::read(&wasm_path)
196            .with_context(|| format!("failed to read output module `{}`", wasm_path.display()))?;
197
198        let component_bytes = self
199            .strip_custom_sections(&component_bytes)
200            .context("failed to strip debug sections from wasm component")?;
201
202        let wasm_extension_api_version =
203            parse_wasm_extension_version(&manifest.id, &component_bytes)
204                .context("compiled wasm did not contain a valid zed extension api version")?;
205        manifest.lib.version = Some(wasm_extension_api_version);
206
207        let extension_file = extension_dir.join("extension.wasm");
208        fs::write(extension_file.clone(), &component_bytes)
209            .context("failed to write extension.wasm")?;
210
211        log::info!(
212            "extension {} written to {}",
213            extension_dir.display(),
214            extension_file.display()
215        );
216
217        Ok(())
218    }
219
220    async fn compile_grammar(
221        &self,
222        extension_dir: &Path,
223        grammar_name: &str,
224        grammar_metadata: &GrammarManifestEntry,
225    ) -> Result<()> {
226        let clang_path = self.install_wasi_sdk_if_needed().await?;
227
228        let mut grammar_repo_dir = extension_dir.to_path_buf();
229        grammar_repo_dir.extend(["grammars", grammar_name]);
230
231        let mut grammar_wasm_path = grammar_repo_dir.clone();
232        grammar_wasm_path.set_extension("wasm");
233
234        log::info!("checking out {grammar_name} parser");
235        self.checkout_repo(
236            &grammar_repo_dir,
237            &grammar_metadata.repository,
238            &grammar_metadata.rev,
239        )
240        .await?;
241
242        let base_grammar_path = grammar_metadata
243            .path
244            .as_ref()
245            .map(|path| grammar_repo_dir.join(path))
246            .unwrap_or(grammar_repo_dir);
247
248        let src_path = base_grammar_path.join("src");
249        let parser_path = src_path.join("parser.c");
250        let scanner_path = src_path.join("scanner.c");
251
252        log::info!("compiling {grammar_name} parser");
253        let clang_output = util::command::new_smol_command(&clang_path)
254            .args(["-fPIC", "-shared", "-Os"])
255            .arg(format!("-Wl,--export=tree_sitter_{grammar_name}"))
256            .arg("-o")
257            .arg(&grammar_wasm_path)
258            .arg("-I")
259            .arg(&src_path)
260            .arg(&parser_path)
261            .args(scanner_path.exists().then_some(scanner_path))
262            .output()
263            .await
264            .context("failed to run clang")?;
265
266        if !clang_output.status.success() {
267            bail!(
268                "failed to compile {} parser with clang: {}",
269                grammar_name,
270                String::from_utf8_lossy(&clang_output.stderr),
271            );
272        }
273
274        Ok(())
275    }
276
277    async fn checkout_repo(&self, directory: &Path, url: &str, rev: &str) -> Result<()> {
278        let git_dir = directory.join(".git");
279
280        if directory.exists() {
281            let remotes_output = util::command::new_smol_command("git")
282                .arg("--git-dir")
283                .arg(&git_dir)
284                .args(["remote", "-v"])
285                .output()
286                .await?;
287            let has_remote = remotes_output.status.success()
288                && String::from_utf8_lossy(&remotes_output.stdout)
289                    .lines()
290                    .any(|line| {
291                        let mut parts = line.split(|c: char| c.is_whitespace());
292                        parts.next() == Some("origin") && parts.any(|part| part == url)
293                    });
294            if !has_remote {
295                bail!(
296                    "grammar directory '{}' already exists, but is not a git clone of '{}'",
297                    directory.display(),
298                    url
299                );
300            }
301        } else {
302            fs::create_dir_all(directory).with_context(|| {
303                format!("failed to create grammar directory {}", directory.display(),)
304            })?;
305            let init_output = util::command::new_smol_command("git")
306                .arg("init")
307                .current_dir(directory)
308                .output()
309                .await?;
310            if !init_output.status.success() {
311                bail!(
312                    "failed to run `git init` in directory '{}'",
313                    directory.display()
314                );
315            }
316
317            let remote_add_output = util::command::new_smol_command("git")
318                .arg("--git-dir")
319                .arg(&git_dir)
320                .args(["remote", "add", "origin", url])
321                .output()
322                .await
323                .context("failed to execute `git remote add`")?;
324            if !remote_add_output.status.success() {
325                bail!(
326                    "failed to add remote {url} for git repository {}",
327                    git_dir.display()
328                );
329            }
330        }
331
332        let fetch_output = util::command::new_smol_command("git")
333            .arg("--git-dir")
334            .arg(&git_dir)
335            .args(["fetch", "--depth", "1", "origin", rev])
336            .output()
337            .await
338            .context("failed to execute `git fetch`")?;
339
340        let checkout_output = util::command::new_smol_command("git")
341            .arg("--git-dir")
342            .arg(&git_dir)
343            .args(["checkout", rev])
344            .current_dir(directory)
345            .output()
346            .await
347            .context("failed to execute `git checkout`")?;
348        if !checkout_output.status.success() {
349            if !fetch_output.status.success() {
350                bail!(
351                    "failed to fetch revision {} in directory '{}'",
352                    rev,
353                    directory.display()
354                );
355            }
356            bail!(
357                "failed to checkout revision {} in directory '{}': {}",
358                rev,
359                directory.display(),
360                String::from_utf8_lossy(&checkout_output.stderr)
361            );
362        }
363
364        Ok(())
365    }
366
367    async fn install_rust_wasm_target_if_needed(&self) -> Result<()> {
368        let rustc_output = util::command::new_smol_command("rustc")
369            .arg("--print")
370            .arg("sysroot")
371            .output()
372            .await
373            .context("failed to run rustc")?;
374        if !rustc_output.status.success() {
375            bail!(
376                "failed to retrieve rust sysroot: {}",
377                String::from_utf8_lossy(&rustc_output.stderr)
378            );
379        }
380
381        let sysroot = PathBuf::from(String::from_utf8(rustc_output.stdout)?.trim());
382        if sysroot.join("lib/rustlib").join(RUST_TARGET).exists() {
383            return Ok(());
384        }
385
386        let output = util::command::new_smol_command("rustup")
387            .args(["target", "add", RUST_TARGET])
388            .stderr(Stdio::piped())
389            .stdout(Stdio::inherit())
390            .output()
391            .await
392            .context("failed to run `rustup target add`")?;
393        if !output.status.success() {
394            bail!(
395                "failed to install the `{RUST_TARGET}` target: {}",
396                String::from_utf8_lossy(&rustc_output.stderr)
397            );
398        }
399
400        Ok(())
401    }
402
403    async fn install_wasi_sdk_if_needed(&self) -> Result<PathBuf> {
404        let url = if let Some(asset_name) = WASI_SDK_ASSET_NAME {
405            format!("{WASI_SDK_URL}{asset_name}")
406        } else {
407            bail!("wasi-sdk is not available for platform {}", env::consts::OS);
408        };
409
410        let wasi_sdk_dir = self.cache_dir.join("wasi-sdk");
411        let mut clang_path = wasi_sdk_dir.clone();
412        clang_path.extend(["bin", &format!("clang{}", env::consts::EXE_SUFFIX)]);
413
414        log::info!("downloading wasi-sdk to {}", wasi_sdk_dir.display());
415
416        if fs::metadata(&clang_path).is_ok_and(|metadata| metadata.is_file()) {
417            return Ok(clang_path);
418        }
419
420        let mut tar_out_dir = wasi_sdk_dir.clone();
421        tar_out_dir.set_extension("archive");
422
423        fs::remove_dir_all(&wasi_sdk_dir).ok();
424        fs::remove_dir_all(&tar_out_dir).ok();
425
426        log::info!("downloading wasi-sdk to {}", wasi_sdk_dir.display());
427        let mut response = self.http.get(&url, AsyncBody::default(), true).await?;
428        let body = GzipDecoder::new({
429            // stream the entire request into memory at once as the artifact is quite big (100MB+)
430            let mut b = vec![];
431            response.body_mut().read_to_end(&mut b).await?;
432            Cursor::new(b)
433        });
434        let tar = Archive::new(body);
435
436        log::info!("un-tarring wasi-sdk to {}", wasi_sdk_dir.display());
437        tar.unpack(&tar_out_dir)
438            .await
439            .context("failed to unpack wasi-sdk archive")?;
440        log::info!("finished downloading wasi-sdk");
441
442        let inner_dir = fs::read_dir(&tar_out_dir)?
443            .next()
444            .context("no content")?
445            .context("failed to read contents of extracted wasi archive directory")?
446            .path();
447        fs::rename(&inner_dir, &wasi_sdk_dir).context("failed to move extracted wasi dir")?;
448        fs::remove_dir_all(&tar_out_dir).ok();
449
450        Ok(clang_path)
451    }
452
453    // This was adapted from:
454    // https://github.com/bytecodealliance/wasm-tools/blob/e8809bb17fcf69aa8c85cd5e6db7cff5cf36b1de/src/bin/wasm-tools/strip.rs
455    fn strip_custom_sections(&self, input: &Vec<u8>) -> Result<Vec<u8>> {
456        use wasmparser::Payload::*;
457
458        let strip_custom_section = |name: &str| {
459            // Default strip everything but:
460            // * the `name` section
461            // * any `component-type` sections
462            // * the `dylink.0` section
463            // * our custom version section
464            name != "name"
465                && !name.starts_with("component-type:")
466                && name != "dylink.0"
467                && name != "zed:api-version"
468        };
469
470        let mut output = Vec::new();
471        let mut stack = Vec::new();
472
473        for payload in Parser::new(0).parse_all(input) {
474            let payload = payload?;
475
476            // Track nesting depth, so that we don't mess with inner producer sections:
477            match payload {
478                Version { encoding, .. } => {
479                    output.extend_from_slice(match encoding {
480                        wasmparser::Encoding::Component => &wasm_encoder::Component::HEADER,
481                        wasmparser::Encoding::Module => &wasm_encoder::Module::HEADER,
482                    });
483                }
484                ModuleSection { .. } | ComponentSection { .. } => {
485                    stack.push(mem::take(&mut output));
486                    continue;
487                }
488                End { .. } => {
489                    let mut parent = match stack.pop() {
490                        Some(c) => c,
491                        None => break,
492                    };
493                    if output.starts_with(&wasm_encoder::Component::HEADER) {
494                        parent.push(ComponentSectionId::Component as u8);
495                        output.encode(&mut parent);
496                    } else {
497                        parent.push(ComponentSectionId::CoreModule as u8);
498                        output.encode(&mut parent);
499                    }
500                    output = parent;
501                }
502                _ => {}
503            }
504
505            if let CustomSection(c) = &payload
506                && strip_custom_section(c.name())
507            {
508                continue;
509            }
510            if let Some((id, range)) = payload.as_section() {
511                RawSection {
512                    id,
513                    data: &input[range],
514                }
515                .append_to(&mut output);
516            }
517        }
518
519        Ok(output)
520    }
521}
522
523fn populate_defaults(manifest: &mut ExtensionManifest, extension_path: &Path) -> Result<()> {
524    // For legacy extensions on the v0 schema (aka, using `extension.json`), clear out any existing
525    // contents of the computed fields, since we don't care what the existing values are.
526    if manifest.schema_version.is_v0() {
527        manifest.languages.clear();
528        manifest.grammars.clear();
529        manifest.themes.clear();
530    }
531
532    let cargo_toml_path = extension_path.join("Cargo.toml");
533    if cargo_toml_path.exists() {
534        manifest.lib.kind = Some(ExtensionLibraryKind::Rust);
535    }
536
537    let languages_dir = extension_path.join("languages");
538    if languages_dir.exists() {
539        for entry in fs::read_dir(&languages_dir).context("failed to list languages dir")? {
540            let entry = entry?;
541            let language_dir = entry.path();
542            let config_path = language_dir.join("config.toml");
543            if config_path.exists() {
544                let relative_language_dir =
545                    language_dir.strip_prefix(extension_path)?.to_path_buf();
546                if !manifest.languages.contains(&relative_language_dir) {
547                    manifest.languages.push(relative_language_dir);
548                }
549            }
550        }
551    }
552
553    let themes_dir = extension_path.join("themes");
554    if themes_dir.exists() {
555        for entry in fs::read_dir(&themes_dir).context("failed to list themes dir")? {
556            let entry = entry?;
557            let theme_path = entry.path();
558            if theme_path.extension() == Some("json".as_ref()) {
559                let relative_theme_path = theme_path.strip_prefix(extension_path)?.to_path_buf();
560                if !manifest.themes.contains(&relative_theme_path) {
561                    manifest.themes.push(relative_theme_path);
562                }
563            }
564        }
565    }
566
567    let icon_themes_dir = extension_path.join("icon_themes");
568    if icon_themes_dir.exists() {
569        for entry in fs::read_dir(&icon_themes_dir).context("failed to list icon themes dir")? {
570            let entry = entry?;
571            let icon_theme_path = entry.path();
572            if icon_theme_path.extension() == Some("json".as_ref()) {
573                let relative_icon_theme_path =
574                    icon_theme_path.strip_prefix(extension_path)?.to_path_buf();
575                if !manifest.icon_themes.contains(&relative_icon_theme_path) {
576                    manifest.icon_themes.push(relative_icon_theme_path);
577                }
578            }
579        }
580    }
581
582    let snippets_json_path = extension_path.join("snippets.json");
583    if snippets_json_path.exists() {
584        manifest.snippets = Some(snippets_json_path);
585    }
586
587    // For legacy extensions on the v0 schema (aka, using `extension.json`), we want to populate the grammars in
588    // the manifest using the contents of the `grammars` directory.
589    if manifest.schema_version.is_v0() {
590        let grammars_dir = extension_path.join("grammars");
591        if grammars_dir.exists() {
592            for entry in fs::read_dir(&grammars_dir).context("failed to list grammars dir")? {
593                let entry = entry?;
594                let grammar_path = entry.path();
595                if grammar_path.extension() == Some("toml".as_ref()) {
596                    #[derive(Deserialize)]
597                    struct GrammarConfigToml {
598                        pub repository: String,
599                        pub commit: String,
600                        #[serde(default)]
601                        pub path: Option<String>,
602                    }
603
604                    let grammar_config = fs::read_to_string(&grammar_path)?;
605                    let grammar_config: GrammarConfigToml = toml::from_str(&grammar_config)?;
606
607                    let grammar_name = grammar_path
608                        .file_stem()
609                        .and_then(|stem| stem.to_str())
610                        .context("no grammar name")?;
611                    if !manifest.grammars.contains_key(grammar_name) {
612                        manifest.grammars.insert(
613                            grammar_name.into(),
614                            GrammarManifestEntry {
615                                repository: grammar_config.repository,
616                                rev: grammar_config.commit,
617                                path: grammar_config.path,
618                            },
619                        );
620                    }
621                }
622            }
623        }
624    }
625
626    Ok(())
627}