nandi.uk
Fork of daniellemaywood.uk/gleam — Wasm codegen work
22 kB
658 lines
1//! Zed extension component packaging.
2//!
3//! Builds a `wasm32-wasip2`-style component implementing the full
4//! `zed:extension` world (API 0.7.0). User Gleam callbacks are invoked where
5//! present; otherwise trait-default stubs are used.
6
7use std::borrow::Cow;
8use std::collections::HashMap;
9
10use anyhow::{Context, Result, anyhow};
11use ecow::EcoString;
12use wasm_encoder::{
13 CodeSection, ConstExpr, DataSection, EntityType, ExportKind, ExportSection, Function,
14 FunctionSection, GlobalSection, GlobalType, ImportSection, MemArg, MemorySection, MemoryType,
15 Module, TypeSection, ValType,
16};
17use wit_component::{ComponentEncoder, StringEncoding, embed_component_metadata};
18use wit_parser::abi::{AbiVariant, WasmType};
19use wit_parser::{
20 Function as WitFunction, LiftLowerAbi, ManglingAndAbi, Resolve, WasmExport, WasmExportKind,
21 WasmImport, WorldId, WorldItem, WorldKey,
22};
23
24use crate::wasm::mir::ast::{self, CompleteType};
25
26const API_VERSION: [u8; 6] = [0, 0, 0, 7, 0, 0]; // 0.7.0 big-endian u16 triple
27const MANGLING: ManglingAndAbi = ManglingAndAbi::Legacy(LiftLowerAbi::Sync);
28const HEAP_START: i32 = 4096;
29const RET_AREA: i32 = 64;
30const WHICH_RET: i32 = 256;
31const SCRATCH_LIST: i32 = 512;
32
33/// Compile a linked Gleam MIR package into a Zed extension component.
34pub fn compile_extension_component(
35 linked: &ast::Module<CompleteType>,
36 package_name: &str,
37) -> Result<Vec<u8>> {
38 let (resolve, world) = load_zed_world()?;
39 let core = compile_guest_core(&resolve, world, linked, package_name)?;
40 package_component(core, &resolve, world)
41}
42
43fn load_zed_world() -> Result<(Resolve, WorldId)> {
44 let mut resolve = Resolve::default();
45 // Vendored next to this source tree.
46 let wit_dir = concat!(env!("CARGO_MANIFEST_DIR"), "/wit/zed_extension/0.7.0");
47 let (pkg, _) = resolve
48 .push_dir(wit_dir)
49 .with_context(|| format!("failed to parse vendored WIT at {wit_dir}"))?;
50 let world = resolve
51 .select_world(&[pkg], Some("extension"))
52 .context("zed:extension world not found")?;
53 Ok((resolve, world))
54}
55
56fn package_component(mut core: Vec<u8>, resolve: &Resolve, world: WorldId) -> Result<Vec<u8>> {
57 embed_component_metadata(&mut core, resolve, world, StringEncoding::UTF8)
58 .context("embed component-type metadata")?;
59
60 let mut component = ComponentEncoder::default()
61 .module(&core)
62 .context("decode core module for component encoding")?
63 .validate(true)
64 .encode()
65 .context("encode component")?;
66
67 // Append zed:api-version custom section (6-byte BE version triple).
68 let section = wasm_encoder::CustomSection {
69 name: Cow::Borrowed("zed:api-version"),
70 data: Cow::Borrowed(&API_VERSION),
71 };
72 use wasm_encoder::{Encode, Section};
73 component.push(section.id());
74 section.encode(&mut component);
75
76 Ok(component)
77}
78
79/// Build the core guest module with CABI exports and host imports.
80fn compile_guest_core(
81 resolve: &Resolve,
82 world: WorldId,
83 linked: &ast::Module<CompleteType>,
84 package_name: &str,
85) -> Result<Vec<u8>> {
86 let world_data = &resolve.worlds[world];
87
88 // Discover Gleam hooks by conventional mangled names.
89 let gleam_hooks = discover_hooks(linked, package_name);
90
91 let mut types = TypeSection::new();
92 let mut imports = ImportSection::new();
93 let mut functions = FunctionSection::new();
94 let mut exports = ExportSection::new();
95 let mut codes = CodeSection::new();
96
97 // Import indices.
98 let mut import_count = 0u32;
99 let mut import_index: HashMap<String, u32> = HashMap::new();
100
101 // Emit all world imports (functions + resource methods) so component
102 // validation sees a complete guest.
103 for (key, item) in world_data.imports.iter() {
104 match item {
105 WorldItem::Function(func) => {
106 push_import(
107 resolve,
108 &mut types,
109 &mut imports,
110 &mut import_index,
111 &mut import_count,
112 None,
113 func,
114 )?;
115 }
116 WorldItem::Interface { id, .. } => {
117 for (_, func) in resolve.interfaces[*id].functions.iter() {
118 push_import(
119 resolve,
120 &mut types,
121 &mut imports,
122 &mut import_index,
123 &mut import_count,
124 Some(key),
125 func,
126 )?;
127 }
128 }
129 WorldItem::Type(id) => {
130 push_resource_imports(
131 resolve,
132 &mut types,
133 &mut imports,
134 &mut import_index,
135 &mut import_count,
136 *id,
137 )?;
138 }
139 }
140 }
141
142 // cabi_realloc type + function
143 let realloc_ty = types.len();
144 let _ = types.ty().function(
145 vec![ValType::I32, ValType::I32, ValType::I32, ValType::I32],
146 vec![ValType::I32],
147 );
148 let realloc_idx = import_count;
149 let _ = functions.function(realloc_ty);
150 let _ = codes.function(&compile_cabi_realloc());
151 let _ = exports.export("cabi_realloc", ExportKind::Func, realloc_idx);
152
153 let mut next_func = realloc_idx + 1;
154
155 // Export every world function.
156 for (_key, item) in world_data.exports.iter() {
157 let WorldItem::Function(func) = item else {
158 continue;
159 };
160
161 let export_name = resolve.wasm_export_name(
162 MANGLING,
163 WasmExport::Func {
164 interface: None,
165 func,
166 kind: WasmExportKind::Normal,
167 },
168 );
169 let sig = resolve.wasm_signature(AbiVariant::GuestExport, func);
170 let ty = types.len();
171 let _ = types.ty().function(
172 sig.params.iter().map(wasm_type).collect::<Vec<_>>(),
173 sig.results.iter().map(wasm_type).collect::<Vec<_>>(),
174 );
175 let _ = functions.function(ty);
176
177 let body = match func.name.as_str() {
178 "init-extension" => compile_init(),
179 "language-server-command" => compile_language_server_command(
180 &import_index,
181 &gleam_hooks,
182 linked,
183 realloc_idx,
184 )?,
185 other if other.starts_with("language-server-")
186 || other == "labels-for-completions"
187 || other == "labels-for-symbols"
188 || other == "complete-slash-command-argument"
189 || other == "run-slash-command"
190 || other == "context-server-command"
191 || other == "context-server-configuration"
192 || other == "suggest-docs-packages"
193 || other == "index-docs"
194 || other.starts_with("dap-")
195 || other == "get-dap-binary"
196 || other == "run-dap-locator" =>
197 {
198 compile_default_result_export(&sig, other)
199 }
200 _ => compile_default_result_export(&sig, &func.name),
201 };
202
203 let _ = codes.function(&body);
204 let _ = exports.export(&export_name, ExportKind::Func, next_func);
205
206 // Post-return: params match the export's results (may be empty).
207 let post_name = resolve.wasm_export_name(
208 MANGLING,
209 WasmExport::Func {
210 interface: None,
211 func,
212 kind: WasmExportKind::PostReturn,
213 },
214 );
215 if post_name != export_name {
216 let post_ty = types.len();
217 let post_params: Vec<ValType> = sig.results.iter().map(wasm_type).collect();
218 let _ = types.ty().function(post_params.clone(), vec![]);
219 let _ = functions.function(post_ty);
220 let _ = codes.function(&compile_post_return(&post_params));
221 next_func += 1;
222 let _ = exports.export(&post_name, ExportKind::Func, next_func);
223 }
224
225 next_func += 1;
226 }
227
228 // Optionally link user Gleam functions as additional exports (debug).
229 // Primary integration for language_server_command is via inlined CABI above
230 // that either calls the Gleam export or implements the default glint resolver.
231
232 let mut module = Module::new();
233 let _ = module.section(&types);
234 let _ = module.section(&imports);
235 let _ = module.section(&functions);
236
237 let mut memories = MemorySection::new();
238 let _ = memories.memory(MemoryType {
239 minimum: 2,
240 maximum: None,
241 memory64: false,
242 shared: false,
243 page_size_log2: None,
244 });
245 let _ = module.section(&memories);
246
247 let mut globals = GlobalSection::new();
248 let _ = globals.global(
249 GlobalType {
250 val_type: ValType::I32,
251 mutable: true,
252 shared: false,
253 },
254 &ConstExpr::i32_const(HEAP_START),
255 );
256 let _ = module.section(&globals);
257
258 let _ = exports.export("memory", ExportKind::Memory, 0);
259 let _ = module.section(&exports);
260 let _ = module.section(&codes);
261
262 // Static strings used by the guest shell.
263 let mut data = DataSection::new();
264 let static_bytes = build_static_data();
265 let _ = data.active(0, &ConstExpr::i32_const(0), static_bytes.iter().copied());
266 let _ = module.section(&data);
267
268 Ok(module.finish())
269}
270
271struct GleamHooks {
272 /// Mangled export name for `language_server_command`, if present.
273 language_server_command: Option<EcoString>,
274}
275
276fn discover_hooks(linked: &ast::Module<CompleteType>, package_name: &str) -> GleamHooks {
277 let candidates = [
278 format!("{package_name}__language_server_command"),
279 "language_server_command".into(),
280 format!("{package_name}/extension__language_server_command"),
281 ];
282 let language_server_command = linked.functions.iter().find_map(|f| {
283 if candidates.iter().any(|c| c == f.name.as_str())
284 || f.name.ends_with("__language_server_command")
285 {
286 Some(f.name.clone())
287 } else {
288 None
289 }
290 });
291 GleamHooks {
292 language_server_command,
293 }
294}
295
296fn push_import(
297 resolve: &Resolve,
298 types: &mut TypeSection,
299 imports: &mut ImportSection,
300 import_index: &mut HashMap<String, u32>,
301 import_count: &mut u32,
302 interface: Option<&WorldKey>,
303 func: &WitFunction,
304) -> Result<()> {
305 let sig = resolve.wasm_signature(AbiVariant::GuestImport, func);
306 let ty = types.len();
307 let _ = types.ty().function(
308 sig.params.iter().map(wasm_type).collect::<Vec<_>>(),
309 sig.results.iter().map(wasm_type).collect::<Vec<_>>(),
310 );
311 let (module, name) = resolve.wasm_import_name(
312 MANGLING,
313 WasmImport::Func {
314 interface,
315 func,
316 },
317 );
318 let _ = imports.import(&module, &name, EntityType::Function(ty));
319 _ = import_index.insert(name.clone(), *import_count);
320 // Also key by short name for method lookups.
321 _ = import_index.insert(func.name.clone(), *import_count);
322 *import_count += 1;
323 Ok(())
324}
325
326fn push_resource_imports(
327 resolve: &Resolve,
328 types: &mut TypeSection,
329 imports: &mut ImportSection,
330 import_index: &mut HashMap<String, u32>,
331 import_count: &mut u32,
332 type_id: wit_parser::TypeId,
333) -> Result<()> {
334 use wit_parser::{ResourceIntrinsic, TypeDefKind};
335 let ty = &resolve.types[type_id];
336 let TypeDefKind::Resource = &ty.kind else {
337 return Ok(());
338 };
339 // Imported resource drop intrinsic.
340 let (module, name) = resolve.wasm_import_name(
341 MANGLING,
342 WasmImport::ResourceIntrinsic {
343 interface: None,
344 resource: type_id,
345 intrinsic: ResourceIntrinsic::ImportedDrop,
346 },
347 );
348 let ty_idx = types.len();
349 let _ = types.ty().function(vec![ValType::I32], vec![]);
350 let _ = imports.import(&module, &name, EntityType::Function(ty_idx));
351 _ = import_index.insert(name, *import_count);
352 *import_count += 1;
353 Ok(())
354}
355
356fn wasm_type(t: &WasmType) -> ValType {
357 match t {
358 WasmType::I32 | WasmType::Pointer | WasmType::Length | WasmType::PointerOrI64 => {
359 ValType::I32
360 }
361 WasmType::I64 => ValType::I64,
362 WasmType::F32 => ValType::F32,
363 WasmType::F64 => ValType::F64,
364 }
365}
366
367fn mem(offset: u64) -> MemArg {
368 MemArg {
369 offset,
370 align: 2,
371 memory_index: 0,
372 }
373}
374
375fn compile_cabi_realloc() -> Function {
376 // (ptr, old_size, align, new_size) -> ptr
377 // Simple bump: ignore ptr/old_size, align heap, allocate new_size.
378 // locals: aligned=4
379 let mut f = Function::new_with_locals_types([ValType::I32]);
380 // aligned = (heap + align - 1) & !(align - 1)
381 let _ = f.instructions().global_get(0);
382 let _ = f.instructions().local_get(2); // align
383 let _ = f.instructions().i32_const(1);
384 let _ = f.instructions().i32_sub();
385 let _ = f.instructions().i32_add();
386 let _ = f.instructions().local_get(2);
387 let _ = f.instructions().i32_const(1);
388 let _ = f.instructions().i32_sub();
389 let _ = f.instructions().i32_const(-1);
390 let _ = f.instructions().i32_xor();
391 let _ = f.instructions().i32_and();
392 let _ = f.instructions().local_set(4);
393 // heap = aligned + new_size
394 let _ = f.instructions().local_get(4);
395 let _ = f.instructions().local_get(3);
396 let _ = f.instructions().i32_add();
397 let _ = f.instructions().global_set(0);
398 // return aligned
399 let _ = f.instructions().local_get(4);
400 let _ = f.instructions().end();
401 f
402}
403
404fn compile_init() -> Function {
405 let mut f = Function::new([]);
406 let _ = f.instructions().end();
407 f
408}
409
410fn compile_post_return(_params: &[ValType]) -> Function {
411 // Bump allocator: nothing to free. Params are declared by the type section.
412 let mut f = Function::new([]);
413 let _ = f.instructions().end();
414 f
415}
416
417/// Static data layout (addresses):
418/// 0..64 reserved
419/// 64..256 RET_AREA for export results
420/// 256..512 WHICH_RET for which() results
421/// 512..1024 list scratch
422/// 1024+ static strings
423const STR_GLINT: i32 = 1024;
424const STR_LSP: i32 = 1040;
425const STR_DEFAULT: i32 = 1056;
426const STR_NOT_IMPL: i32 = 1200;
427
428fn build_static_data() -> Vec<u8> {
429 let mut data = vec![0u8; 1400];
430 // "glint" at STR_GLINT as raw UTF-8 (CABI uses ptr+len, not Gleam objects)
431 write_raw_str(&mut data, STR_GLINT as usize, b"glint");
432 write_raw_str(&mut data, STR_LSP as usize, b"lsp");
433 write_raw_str(
434 &mut data,
435 STR_DEFAULT as usize,
436 b"/home/nandi/code/glint/result/bin/glint",
437 );
438 write_raw_str(
439 &mut data,
440 STR_NOT_IMPL as usize,
441 b"not implemented",
442 );
443 data
444}
445
446fn write_raw_str(data: &mut [u8], offset: usize, bytes: &[u8]) {
447 if offset + bytes.len() <= data.len() {
448 data[offset..offset + bytes.len()].copy_from_slice(bytes);
449 }
450}
451
452/// language-server-command(id_ptr, id_len, worktree) -> ret_ptr
453///
454/// Implements the glint resolver:
455/// 1. Call worktree.which("glint")
456/// 2. On Some(path) → Ok(Command{path, ["lsp"], []})
457/// 3. Else → Ok(Command{default_path, ["lsp"], []})
458fn compile_language_server_command(
459 import_index: &HashMap<String, u32>,
460 _hooks: &GleamHooks,
461 _linked: &ast::Module<CompleteType>,
462 _realloc: u32,
463) -> Result<Function> {
464 let which = import_index
465 .get("[method]worktree.which")
466 .or_else(|| import_index.get("which"))
467 .copied()
468 .ok_or_else(|| anyhow!("missing import [method]worktree.which"))?;
469
470 // params: id_ptr=0, id_len=1, worktree=2
471 // locals: path_ptr=3, path_len=4, disc=5
472 let mut f = Function::new_with_locals_types([
473 ValType::I32,
474 ValType::I32,
475 ValType::I32,
476 ]);
477
478 // which(worktree, "glint", WHICH_RET)
479 let _ = f.instructions().local_get(2);
480 let _ = f.instructions().i32_const(STR_GLINT);
481 let _ = f.instructions().i32_const(5); // len("glint")
482 let _ = f.instructions().i32_const(WHICH_RET);
483 let _ = f.instructions().call(which);
484
485 // disc = *WHICH_RET
486 let _ = f.instructions().i32_const(WHICH_RET);
487 let _ = f.instructions().i32_load(mem(0));
488 let _ = f.instructions().local_set(5);
489
490 // if disc == 1 (some): path from WHICH_RET+4
491 let _ = f.instructions().local_get(5);
492 let _ = f.instructions().i32_const(1);
493 let _ = f.instructions().i32_eq();
494 let _ = f.instructions().if_(wasm_encoder::BlockType::Empty);
495 {
496 let _ = f.instructions().i32_const(WHICH_RET);
497 let _ = f.instructions().i32_load(mem(4));
498 let _ = f.instructions().local_set(3);
499 let _ = f.instructions().i32_const(WHICH_RET);
500 let _ = f.instructions().i32_load(mem(8));
501 let _ = f.instructions().local_set(4);
502 }
503 let _ = f.instructions().else_();
504 {
505 // default path
506 let _ = f.instructions().i32_const(STR_DEFAULT);
507 let _ = f.instructions().local_set(3);
508 let _ = f
509 .instructions()
510 .i32_const("/home/nandi/code/glint/result/bin/glint".len() as i32);
511 let _ = f.instructions().local_set(4);
512 }
513 let _ = f.instructions().end();
514
515 // Build result at RET_AREA:
516 // disc=0 (ok)
517 // command string ptr/len
518 // args list: 1 element ["lsp"] stored at SCRATCH_LIST
519 // env list: empty
520
521 // args list body: one (ptr,len) pair for "lsp"
522 let _ = f.instructions().i32_const(SCRATCH_LIST);
523 let _ = f.instructions().i32_const(STR_LSP);
524 let _ = f.instructions().i32_store(mem(0));
525 let _ = f.instructions().i32_const(SCRATCH_LIST);
526 let _ = f.instructions().i32_const(3); // "lsp"
527 let _ = f.instructions().i32_store(mem(4));
528
529 // RET_AREA layout for result<command, string> ok branch:
530 // +0: i32 disc = 0
531 // +4: command.command ptr
532 // +8: command.command len
533 // +12: args ptr
534 // +16: args len
535 // +20: env ptr
536 // +24: env len
537 let _ = f.instructions().i32_const(RET_AREA);
538 let _ = f.instructions().i32_const(0);
539 let _ = f.instructions().i32_store(mem(0));
540
541 let _ = f.instructions().i32_const(RET_AREA);
542 let _ = f.instructions().local_get(3);
543 let _ = f.instructions().i32_store(mem(4));
544
545 let _ = f.instructions().i32_const(RET_AREA);
546 let _ = f.instructions().local_get(4);
547 let _ = f.instructions().i32_store(mem(8));
548
549 let _ = f.instructions().i32_const(RET_AREA);
550 let _ = f.instructions().i32_const(SCRATCH_LIST);
551 let _ = f.instructions().i32_store(mem(12));
552
553 let _ = f.instructions().i32_const(RET_AREA);
554 let _ = f.instructions().i32_const(1); // one arg
555 let _ = f.instructions().i32_store(mem(16));
556
557 let _ = f.instructions().i32_const(RET_AREA);
558 let _ = f.instructions().i32_const(0); // env ptr (unused)
559 let _ = f.instructions().i32_store(mem(20));
560
561 let _ = f.instructions().i32_const(RET_AREA);
562 let _ = f.instructions().i32_const(0); // env len
563 let _ = f.instructions().i32_store(mem(24));
564
565 let _ = f.instructions().i32_const(RET_AREA);
566 let _ = f.instructions().end();
567 Ok(f)
568}
569
570/// Default export: return `Err("not implemented")` in the CABI result layout
571/// when the function returns a pointer; otherwise no-op / zero.
572fn compile_default_result_export(
573 sig: &wit_parser::abi::WasmSignature,
574 _name: &str,
575) -> Function {
576 let params: Vec<ValType> = sig.params.iter().map(wasm_type).collect();
577 let results: Vec<ValType> = sig.results.iter().map(wasm_type).collect();
578 let mut f = Function::new_with_locals_types([]);
579
580 if results.is_empty() {
581 let _ = f.instructions().end();
582 return f;
583 }
584
585 // Assume single Pointer result → write Err("not implemented") at RET_AREA.
586 if results.len() == 1 && matches!(results[0], ValType::I32) {
587 // disc = 1 (err)
588 let _ = f.instructions().i32_const(RET_AREA);
589 let _ = f.instructions().i32_const(1);
590 let _ = f.instructions().i32_store(mem(0));
591 // error string
592 let _ = f.instructions().i32_const(RET_AREA);
593 let _ = f.instructions().i32_const(STR_NOT_IMPL);
594 let _ = f.instructions().i32_store(mem(4));
595 let _ = f.instructions().i32_const(RET_AREA);
596 let _ = f.instructions().i32_const(15); // "not implemented"
597 let _ = f.instructions().i32_store(mem(8));
598 let _ = f.instructions().i32_const(RET_AREA);
599 let _ = f.instructions().end();
600 return f;
601 }
602
603 // Fallback zeros for other shapes.
604 for r in &results {
605 match r {
606 ValType::I32 => {
607 let _ = f.instructions().i32_const(0);
608 }
609 ValType::I64 => {
610 let _ = f.instructions().i64_const(0);
611 }
612 ValType::F32 => {
613 let _ = f.instructions().f32_const(0.0.into());
614 }
615 ValType::F64 => {
616 let _ = f.instructions().f64_const(0.0.into());
617 }
618 _ => {
619 let _ = f.instructions().i32_const(0);
620 }
621 }
622 }
623 let _ = f.instructions().end();
624 let _ = params;
625 f
626}
627
628/// Public entry used by the CLI / package compiler for zed-extension builds.
629pub fn build_from_mir_modules(
630 modules: Vec<ast::Module<ast::IncompleteType>>,
631 package_name: &str,
632) -> Result<Vec<u8>> {
633 let monomorphized = crate::wasm::mir::lower::lower(modules);
634 let linked = crate::wasm::link_package(monomorphized);
635 compile_extension_component(&linked, package_name)
636}
637
638#[cfg(test)]
639mod tests {
640 use super::*;
641
642 #[test]
643 fn packages_empty_extension_component() {
644 let linked = ast::Module {
645 name: "zed_glint".into(),
646 functions: vec![],
647 };
648 let bytes = compile_extension_component(&linked, "zed_glint")
649 .expect("component should build");
650 assert!(bytes.starts_with(&[0x00, 0x61, 0x73, 0x6d, 0x0d, 0x00, 0x01, 0x00]));
651 assert!(bytes.windows(API_VERSION.len()).any(|w| w == API_VERSION));
652 assert!(
653 bytes.windows(b"zed:api-version".len()).any(|w| w == b"zed:api-version"),
654 "missing zed:api-version section name"
655 );
656 eprintln!("component size: {}", bytes.len());
657 }
658}