feat: treeshake imports based on guest content (#117)

Author: karthik2804

README.md

@@ -152,6 +152,9 @@ export function componentize(jsSource: string, opts: {
 
 Converts a JS source into a component binary.
 
+Imports provides the list of used guest imports only, while the StarlingMonkey engine may pull in additional
+imports. Direct component analysis should be used to correctly infer the real imports list.
+
 ## Contributing
 
 ### Pre-requisites

crates/spidermonkey-embedding-splicer/src/bindgen.rs

@@ -1,4 +1,5 @@
 use crate::{uwrite, uwriteln};
+use anyhow::{bail, Result};
 use heck::*;
 use js_component_bindgen::function_bindgen::{
     ErrHandling, FunctionBindgen, ResourceData, ResourceMap, ResourceTable,
@@ -124,7 +125,13 @@ pub struct Componentization {
     pub resource_imports: Vec<(String, String, u32)>,
 }
 
-pub fn componentize_bindgen(resolve: &Resolve, id: WorldId, name: &str) -> Componentization {
+pub fn componentize_bindgen(
+    resolve: &Resolve,
+    id: WorldId,
+    name: &str,
+    guest_imports: &Vec<String>,
+    guest_exports: &Vec<String>,
+) -> Result<Componentization> {
     let mut bindgen = JsBindgen {
         src: Source::default(),
         esm_bindgen: EsmBindgen::default(),
@@ -147,9 +154,9 @@ pub fn componentize_bindgen(resolve: &Resolve, id: WorldId, name: &str) -> Compo
         .local_names
         .exclude_globals(Intrinsic::get_global_names());
 
-    bindgen.imports_bindgen();
+    bindgen.imports_bindgen(&guest_imports);
 
-    bindgen.exports_bindgen();
+    bindgen.exports_bindgen(&guest_exports)?;
     bindgen.esm_bindgen.populate_export_aliases();
 
     // consolidate import specifiers and generate wrappers
@@ -348,13 +355,13 @@ pub fn componentize_bindgen(resolve: &Resolve, id: WorldId, name: &str) -> Compo
     output.push_str(&js_intrinsics);
     output.push_str(&bindgen.src);
 
-    Componentization {
+    Ok(Componentization {
         js_bindings: output.to_string(),
         exports: bindgen.exports,
         imports: bindgen.imports,
         import_wrappers,
         resource_imports,
-    }
+    })
 }
 
 impl JsBindgen<'_> {
@@ -363,9 +370,39 @@ impl JsBindgen<'_> {
         return intrinsic.name().to_string();
     }
 
-    fn exports_bindgen(&mut self) {
+    fn exports_bindgen(&mut self, guest_exports: &Vec<String>) -> Result<()> {
         for (key, export) in &self.resolve.worlds[self.world].exports {
             let name = self.resolve.name_world_key(key);
+
+            // Do not generate exports when the guest export is not implemented.
+            // We check both the full interface name - "ns:pkg@v/my-interface" and the
+            // aliased interface name "myInterface". All other names are always
+            // camel-case in the check.
+            match key {
+                WorldKey::Interface(iface) => {
+                    if !guest_exports.contains(&name) {
+                        let iface = &self.resolve.interfaces[*iface];
+                        if let Some(name) = iface.name.as_ref() {
+                            let camel_case_name = name.to_lower_camel_case();
+                            if !guest_exports.contains(&camel_case_name) {
+                                bail!("Expected a JS export definition for '{}'", camel_case_name);
+                            }
+                            // TODO: move populate_export_aliases to a preprocessing
+                            // step that doesn't require esm_bindgen, so that we can
+                            // do alias deduping here as well.
+                        } else {
+                            continue;
+                        }
+                    }
+                }
+                WorldKey::Name(export_name) => {
+                    let camel_case_name = export_name.to_lower_camel_case();
+                    if !guest_exports.contains(&camel_case_name) {
+                        bail!("Expected a JS export definition for '{}'", camel_case_name);
+                    }
+                }
+            }
+
             match export {
                 WorldItem::Function(func) => {
                     let local_name = self.local_names.create_once(&func.name).to_string();
@@ -449,11 +486,15 @@ impl JsBindgen<'_> {
                 WorldItem::Type(_) => {}
             }
         }
+        Ok(())
     }
 
-    fn imports_bindgen(&mut self) {
+    fn imports_bindgen(&mut self, guest_imports: &Vec<String>) {
         for (key, impt) in &self.resolve.worlds[self.world].imports {
             let import_name = self.resolve.name_world_key(key);
+            if !guest_imports.contains(&import_name) {
+                continue;
+            }
             match &impt {
                 WorldItem::Function(f) => {
                     self.import_bindgen(import_name, f, false, None);

crates/spidermonkey-embedding-splicer/src/lib.rs

@@ -1,6 +1,9 @@
 use anyhow::{bail, Context, Result};
 use bindgen::BindingItem;
-use std::path::{Path, PathBuf};
+use std::{
+    path::{Path, PathBuf},
+    vec,
+};
 
 mod bindgen;
 mod splice;
@@ -110,6 +113,8 @@ impl Guest for SpidermonkeyEmbeddingSplicerComponent {
         wit_source: Option<String>,
         wit_path: Option<String>,
         world_name: Option<String>,
+        mut guest_imports: Vec<String>,
+        guest_exports: Vec<String>,
         debug: bool,
     ) -> Result<SpliceResult, String> {
         let source_name = source_name.unwrap_or("source.js".to_string());
@@ -131,7 +136,6 @@ impl Guest for SpidermonkeyEmbeddingSplicerComponent {
             .map_err(|e| e.to_string())?;
 
         let mut wasm_bytes = wit_component::dummy_module(&resolve, world);
-        let componentized = bindgen::componentize_bindgen(&resolve, world, &source_name);
 
         // merge the engine world with the target world, retaining the engine producers
         let producers = if let Ok((
@@ -144,6 +148,11 @@ impl Guest for SpidermonkeyEmbeddingSplicerComponent {
             },
         )) = decode(&engine)
         {
+            // merge the imports from the engine with the imports from the guest content
+            for (k, _) in &engine_resolve.worlds[engine_world].imports {
+                guest_imports.push(engine_resolve.name_world_key(k));
+            }
+
             // we disable the engine run and incoming handler as we recreate these exports
             // when needed, so remove these from the world before initiating the merge
             let maybe_run = engine_resolve.worlds[engine_world]
@@ -174,19 +183,29 @@ impl Guest for SpidermonkeyEmbeddingSplicerComponent {
             resolve
                 .merge(engine_resolve)
                 .expect("unable to merge with engine world");
-            let (engine_world, _) = resolve
-                .worlds
-                .iter()
-                .find(|(world, _)| resolve.worlds[*world].name == "root")
-                .unwrap();
-            resolve
-                .merge_worlds(engine_world, world)
-                .expect("unable to merge with engine world");
             producers
         } else {
             None
         };
 
+        let componentized = bindgen::componentize_bindgen(
+            &resolve,
+            world,
+            &source_name,
+            &guest_imports,
+            &guest_exports,
+        )
+        .map_err(|err| err.to_string())?;
+
+        let (engine_world, _) = resolve
+            .worlds
+            .iter()
+            .find(|(world, _)| resolve.worlds[*world].name == "root")
+            .unwrap();
+        resolve
+            .merge_worlds(engine_world, world)
+            .expect("unable to merge with engine world");
+
         let encoded = wit_component::metadata::encode(
             &resolve,
             world,
@@ -327,8 +346,8 @@ impl Guest for SpidermonkeyEmbeddingSplicerComponent {
             ));
         }
 
-        // println!("{:?}", &imports);
         // println!("{:?}", &componentized.imports);
+        // println!("{:?}", &componentized.resource_imports);
         // println!("{:?}", &exports);
         let mut wasm =
             splice::splice(engine, imports, exports, debug).map_err(|e| format!("{:?}", e))?;

crates/spidermonkey-embedding-splicer/src/splice.rs

@@ -38,20 +38,30 @@ pub fn splice(
     let mut module = config.parse(&engine)?;
 
     // since StarlingMonkey implements CLI Run and incoming handler,
-    // we override these in ComponentizeJS, removing them from the
-    // core function exports
-    if let Ok(run) = module.exports.get_func("wasi:cli/[email protected]#run") {
-        let expt = module.exports.get_exported_func(run).unwrap();
-        module.exports.delete(expt.id());
-        module.funcs.delete(run);
+    // we override them only if the guest content exports those functions
+    if exports
+        .iter()
+        .any(|(name, _)| name == "wasi:cli/[email protected]#run")
+    {
+        if let Ok(run) = module.exports.get_func("wasi:cli/[email protected]#run") {
+            let expt = module.exports.get_exported_func(run).unwrap();
+            module.exports.delete(expt.id());
+            module.funcs.delete(run);
+        }
     }
-    if let Ok(serve) = module
-        .exports
-        .get_func("wasi:http/[email protected]#handle")
+
+    if exports
+        .iter()
+        .any(|(name, _)| name == "wasi:http/[email protected]#handle")
     {
-        let expt = module.exports.get_exported_func(serve).unwrap();
-        module.exports.delete(expt.id());
-        module.funcs.delete(serve);
+        if let Ok(serve) = module
+            .exports
+            .get_func("wasi:http/[email protected]#handle")
+        {
+            let expt = module.exports.get_exported_func(serve).unwrap();
+            module.exports.delete(expt.id());
+            module.funcs.delete(serve);
+        }
     }
 
     // we reencode the WASI world component data, so strip it out from the

crates/spidermonkey-embedding-splicer/wit/spidermonkey-embedding-splicer.wit

@@ -33,5 +33,5 @@ world spidermonkey-embedding-splicer {
 
   export stub-wasi: func(engine: list<u8>, features: list<features>, wit-world: option<string>, wit-path: option<string>, world-name: option<string>) -> result<list<u8>, string>;
 
-  export splice-bindings: func(source-name: option<string>, spidermonkey-engine: list<u8>, wit-world: option<string>, wit-path: option<string>, world-name: option<string>, debug: bool) -> result<splice-result, string>;
+  export splice-bindings: func(source-name: option<string>, spidermonkey-engine: list<u8>, wit-world: option<string>, wit-path: option<string>, world-name: option<string>, guest-imports: list<string>, guest-exports: list<string>, debug: bool) -> result<splice-result, string>;
 }

src/componentize.js

@@ -47,12 +47,33 @@ export async function componentize(jsSource, witWorld, opts) {
     enableFeatures = [],
   } = opts || {};
 
+  await lexerInit;
+  let jsImports = [];
+  let jsExports = [];
+  try {
+    [jsImports, jsExports] = parse(jsSource);
+  } catch {
+    // ignore parser errors - will show up as engine parse errors shortly
+  }
+
+  let guestImports = [];
+  jsImports.map((k) => {
+    guestImports.push(k.n);
+  });
+
+  let guestExports = [];
+  jsExports.map((k) => {
+    guestExports.push(k.n);
+  });
+
   let { wasm, jsBindings, importWrappers, exports, imports } = spliceBindings(
     sourceName,
     await readFile(engine),
     witWorld,
     maybeWindowsPath(witPath),
     worldName,
+    guestImports,
+    guestExports,
     false
   );
 
@@ -103,13 +124,6 @@ export async function componentize(jsSource, witWorld, opts) {
   await writeFile(input, Buffer.from(wasm));
 
   // rewrite the JS source import specifiers to reference import wrappers
-  await lexerInit;
-  let jsImports = [];
-  try {
-    [jsImports] = parse(jsSource);
-  } catch {
-    // ignore parser errors - will show up as engine parse errors shortly
-  }
   let source = '',
     curIdx = 0;
   for (const jsImpt of jsImports) {
@@ -183,7 +197,9 @@ export async function componentize(jsSource, witWorld, opts) {
       {
         stdio: [null, stdout, stderr],
         env,
-        input: maybeWindowsPath(join(sourceDir, sourceName.slice(0, -3) + '.bindings.js')),
+        input: maybeWindowsPath(
+          join(sourceDir, sourceName.slice(0, -3) + '.bindings.js')
+        ),
         shell: true,
         encoding: 'utf-8',
       }
@@ -299,9 +315,13 @@ export async function componentize(jsSource, witWorld, opts) {
   }
 
   // after wizering, stub out the wasi imports depending on what features are enabled
-  const finalBin = stubWasi(bin, features, witWorld,
+  const finalBin = stubWasi(
+    bin,
+    features,
+    witWorld,
     maybeWindowsPath(witPath),
-    worldName,);
+    worldName
+  );
 
   const component = await metadataAdd(
     await componentNew(

test/cases/missing-export/source.js

@@ -0,0 +1,3 @@
+export function notFunc () {
+
+}

test/cases/missing-export/test.js

@@ -0,0 +1,5 @@
+import { strictEqual } from 'node:assert';
+
+export function err(e) {
+  strictEqual(e.message, 'Expected a JS export definition for \'expected\'');
+}

test/cases/missing-export/world.wit

@@ -0,0 +1,5 @@
+package local:smoke;
+
+world the-world {
+  export expected: func();
+}

test/cases/smoke/source.js

@@ -1,5 +1,3 @@
-import { y } from 'imports';
-
 export const exports = {
   hello () {
     return 'world (' + getNum('world') + ')';

test/cases/smoke/imports.js test/cases/treeshake/imports.js

@@ -1,3 +1,3 @@
 export function y () {
-  // console.log('y');
+  globalThis.y = true;
 }

test/cases/treeshake/source.js

@@ -0,0 +1,5 @@
+import * as imports from 'imports';
+
+export function hello () {
+    imports.y();
+}

test/cases/treeshake/test.js

@@ -0,0 +1,7 @@
+import { strictEqual, deepStrictEqual } from 'node:assert';
+
+export function test(instance, { imports }) {
+  deepStrictEqual(imports, [['imports', 'y']]);
+  instance.hello();
+  strictEqual(globalThis.y, true);
+}

test/cases/treeshake/world.wit

@@ -0,0 +1,9 @@
+package local:smoke;
+
+world the-world {
+  import unused: func();
+  import imports: interface {
+    y: func();
+  }
+  export hello: func();
+}