feat: Add compression support for wasm binary

src/pipelines/rust/compress.rs

@@ -0,0 +1,99 @@
+use std::{
+    io::{BufRead, Read},
+    ops::Deref,
+    str::FromStr,
+};
+
+use anyhow::bail;
+use flate2::{
+    bufread::{DeflateEncoder, GzEncoder, ZlibEncoder},
+    Compression,
+};
+
+#[derive(PartialEq, Eq, Debug, Default)]
+pub enum CompressionAlgorithm {
+    #[default]
+    Gzip,
+    Zlib,
+    Deflate,
+}
+
+impl CompressionAlgorithm {
+    pub fn encoder<'a, R: BufRead + Send + 'a>(
+        &self,
+        reader: R,
+        level: Compression,
+    ) -> Box<dyn Read + Send + 'a> {
+        match self {
+            Self::Gzip => Box::new(GzEncoder::new(reader, level)),
+            Self::Zlib => Box::new(ZlibEncoder::new(reader, level)),
+            Self::Deflate => Box::new(DeflateEncoder::new(reader, level)),
+        }
+    }
+
+    pub fn as_string(&self) -> String {
+        match self {
+            Self::Gzip => "gzip".to_string(),
+            Self::Zlib => "deflate".to_string(),
+            Self::Deflate => "deflate-raw".to_string(),
+        }
+    }
+}
+
+impl FromStr for CompressionAlgorithm {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> anyhow::Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "gzip" => Ok(Self::Gzip),
+            "deflate" => Ok(Self::Zlib),
+            "deflate-raw" => Ok(Self::Deflate),
+            _ => bail!("unknown compression algorithm `{}`", s),
+        }
+    }
+}
+
+#[derive(Default, PartialEq, Eq)]
+pub struct CompressionLevel(Compression);
+
+impl CompressionLevel {
+    pub const OFF: Self = Self(Compression::new(0));
+}
+
+impl FromStr for CompressionLevel {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> anyhow::Result<Self, Self::Err> {
+        let level = match s {
+            "0" => Compression::new(0),
+            "1" => Compression::new(1),
+            "2" => Compression::new(2),
+            "3" => Compression::new(3),
+            "4" => Compression::new(4),
+            "5" => Compression::new(5),
+            "6" => Compression::new(6),
+            "7" => Compression::new(7),
+            "8" => Compression::new(8),
+            "9" => Compression::new(9),
+            "default" => Compression::default(),
+            "fast" => Compression::fast(),
+            "best" => Compression::best(),
+            _ => bail!("unknown gzip level `{}`", s),
+        };
+        Ok(Self(level))
+    }
+}
+
+impl AsRef<Compression> for CompressionLevel {
+    fn as_ref(&self) -> &Compression {
+        &self.0
+    }
+}
+
+impl Deref for CompressionLevel {
+    type Target = Compression;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}

src/pipelines/rust/initializer.js

@@ -1,4 +1,4 @@
-async function __trunkInitializer(init, source, sourceSize, initializer, initWithObject) {
+async function __trunkInitializer(init, source, sourceSize, initializer, initWithObject, compressAlgorithm) {
   if (initializer === undefined) {
     return await init(initWithObject ? { module_or_path: source } : source);
   }
@@ -11,14 +11,20 @@ async function __trunkInitializer(init, source, sourceSize, initializer, initWit
 
   const response = fetch(source)
       .then((response) => {
-        const reader = response.body.getReader();
         const headers = response.headers;
         const status = response.status;
         const statusText = response.statusText;
 
         const total = sourceSize;
         let current = 0;
 
+        let reader = undefined;
+        if (compressAlgorithm) {
+          reader = response.body.pipeThrough(new DecompressionStream(compressAlgorithm)).getReader();
+        } else {
+          reader = response.body.getReader();
+        }
+
         const stream = new ReadableStream({
           start(controller) {
             function push() {

src/pipelines/rust/mod.rs

@@ -1,5 +1,6 @@
 //! Rust application pipeline.
 
+mod compress;
 mod output;
 mod sri;
 mod wasm_bindgen;
@@ -18,7 +19,10 @@ use crate::{
         types::CrossOrigin,
         CargoMetadata,
     },
-    pipelines::rust::sri::{SriBuilder, SriOptions, SriType},
+    pipelines::rust::{
+        compress::{CompressionAlgorithm, CompressionLevel},
+        sri::{SriBuilder, SriOptions, SriType},
+    },
     processing::{integrity::IntegrityType, minify::minify_js},
     tools::{self, Application, ToolInformation},
 };
@@ -97,6 +101,10 @@ pub struct RustApp {
     import_bindings_name: Option<String>,
     /// The name of the initializer module
     initializer: Option<PathBuf>,
+    /// Compression Algorithm for the WASM file.
+    compression_algorithm: CompressionAlgorithm,
+    /// Compression level for the WASM file. Defaults to `CompressionLevel::Default` on release build.
+    compression_level: CompressionLevel,
 }
 
 /// Describes how the rust application is used.
@@ -185,6 +193,22 @@ impl RustApp {
                 RustAppType::Main => WasmBindgenTarget::Web,
                 RustAppType::Worker => WasmBindgenTarget::NoModules,
             });
+        let compression_algorithm = attrs
+            .get("data-compression-algorithm")
+            .map(|attr| attr.parse())
+            .transpose()?
+            .unwrap_or(CompressionAlgorithm::Gzip);
+        let compression_level = attrs
+            .get("data-compression-level")
+            .map(|attr| attr.parse())
+            .transpose()?
+            .unwrap_or_else(|| {
+                if cfg.release {
+                    Default::default()
+                } else {
+                    CompressionLevel::OFF
+                }
+            });
         let cross_origin = attrs
             .get("data-cross-origin")
             .map(|attr| CrossOrigin::from_str(attr))
@@ -307,6 +331,8 @@ impl RustApp {
             import_bindings_name,
             initializer,
             target_path,
+            compression_algorithm,
+            compression_level,
         })
     }
 
@@ -357,6 +383,8 @@ impl RustApp {
             import_bindings_name: None,
             initializer: None,
             target_path: None,
+            compression_algorithm: CompressionAlgorithm::Gzip,
+            compression_level: CompressionLevel::OFF,
         }))
     }
 
@@ -386,6 +414,11 @@ impl RustApp {
             .await
             .context("running wasm-opt")?;
 
+        // (optionally) gzip the wasm file
+        self.wasm_compression(&output.wasm_output)
+            .await
+            .context("gzip compression")?;
+
         // evaluate wasm integrity after all processing
         self.final_digest(&mut output)
             .await
@@ -748,7 +781,12 @@ impl RustApp {
         };
 
         // return output
-
+        let compression_algorithm =
+            if self.cfg.release && self.compression_level != CompressionLevel::OFF {
+                Some(self.compression_algorithm.as_string())
+            } else {
+                None
+            };
         Ok(RustAppOutput {
             id: self.id,
             cfg: self.cfg.clone(),
@@ -761,6 +799,7 @@ impl RustApp {
             import_bindings: self.import_bindings,
             import_bindings_name: self.import_bindings_name.clone(),
             initializer,
+            compression_algorithm,
             wasm_bindgen_features,
         })
     }
@@ -954,6 +993,64 @@ impl RustApp {
         Ok(())
     }
 
+    #[tracing::instrument(level = "trace", skip(self))]
+    async fn wasm_compression(&self, wasm_name: &str) -> Result<()> {
+        if !self.cfg.release {
+            return Ok(());
+        }
+
+        if self.compression_level == CompressionLevel::OFF {
+            log::debug!("compression is turned off");
+            return Ok(());
+        }
+
+        let compression_name = "wasm-compression";
+        let mode_segment = if self.cfg.release { "release" } else { "debug" };
+        let output = self
+            .manifest
+            .metadata
+            .target_directory
+            .join(compression_name)
+            .join(mode_segment);
+        fs::create_dir_all(&output)
+            .await
+            .context("error creating wasm gzip compression output dir")?;
+
+        tracing::debug!(
+            "compressing with gzip level {}",
+            self.compression_level.level()
+        );
+        let output = output.join(format!("{}_bg.wasm", self.name));
+        let target_wasm = self
+            .cfg
+            .staging_dist
+            .join(wasm_name)
+            .to_string_lossy()
+            .to_string();
+
+        let target_wasm_file = std::fs::File::open(&target_wasm)
+            .context("error opening wasm file for gzip compression")?;
+        let target_wasm_reader = std::io::BufReader::new(target_wasm_file);
+        let mut encoder = self
+            .compression_algorithm
+            .encoder(target_wasm_reader, *self.compression_level);
+
+        let output_file =
+            std::fs::File::create(&output).context("error creating compression output file")?;
+        let mut output_writer = std::io::BufWriter::new(output_file);
+        std::io::copy(&mut encoder, &mut output_writer)
+            .context("error writing compressed wasm file")?;
+
+        // Copy the generated WASM file to the dist dir.
+        tracing::debug!("copying generated wasm-opt artifact from '{output}' to '{target_wasm}'");
+        fs::copy(output, &target_wasm).await.context(format!(
+            "error copying ({} compressed) wasm file to dist dir",
+            self.compression_algorithm.as_string()
+        ))?;
+
+        Ok(())
+    }
+
     /// Build the final WASM digest
     #[tracing::instrument(level = "trace", skip(self, output))]
     async fn final_digest(&self, output: &mut RustAppOutput) -> Result<()> {

src/pipelines/rust/output.rs

@@ -29,6 +29,8 @@ pub struct RustAppOutput {
     pub import_bindings: bool,
     /// The name of the WASM bindings import
     pub import_bindings_name: Option<String>,
+    /// Compression algorithm used for the WASM file
+    pub compression_algorithm: Option<String>,
     /// The target of the initializer module
     pub initializer: Option<String>,
     /// The features supported by the version of wasm-bindgen used
@@ -138,24 +140,51 @@ window.{bindings} = bindings;
 dispatchEvent(new CustomEvent("TrunkApplicationStarted", {detail: {wasm}}));
 "#;
 
+        let compression = if let Some(algorithm) = &self.compression_algorithm {
+            format!(
+                r#"
+const resp = await fetch('{base}{wasm}');
+if (!resp.ok) {{
+    throw new Error('Failed to fetch WASM module: ' + resp.statusText);
+}}
+
+const decompressStream = resp.body.pipeThrough(new DecompressionStream('{algorithm}'));
+const wasmBytes = await new Response(decompressStream).arrayBuffer();
+                "#
+            )
+        } else {
+            String::new()
+        };
+
         let init_with_object = self.wasm_bindgen_features.init_with_object;
 
         match &self.initializer {
-            None => format!(
-                r#"
+            None => {
+                format!(
+                    r#"
 <script type="module"{nonce}>
 import init{import} from '{base}{js}';
+{compression}
 const wasm = await init({init_arg});
 
 {bind}
 {fire}
 </script>"#,
-                init_arg = if init_with_object {
-                    format!("{{ module_or_path: '{base}{wasm}' }}")
-                } else {
-                    format!("'{base}{wasm}'")
-                }
-            ),
+                    init_arg = {
+                        let param = if self.compression_algorithm.is_some() {
+                            "wasmBytes".to_string()
+                        } else {
+                            format!("'{base}{wasm}'")
+                        };
+
+                        if init_with_object {
+                            format!("{{ module_or_path: {param} }}")
+                        } else {
+                            param
+                        }
+                    }
+                )
+            }
             Some(initializer) => format!(
                 r#"
 <script type="module"{nonce}>
@@ -164,13 +193,18 @@ const wasm = await init({init_arg});
 import init{import} from '{base}{js}';
 import initializer from '{base}{initializer}';
 
-const wasm = await __trunkInitializer(init, '{base}{wasm}', {size}, initializer(), {init_with_object});
+const wasm = await __trunkInitializer(init, '{base}{wasm}', {size}, initializer(), {init_with_object}{algorithm});
 
 {bind}
 {fire}
 </script>"#,
                 init = include_str!("initializer.js"),
                 size = self.wasm_size,
+                algorithm = if let Some(algorithm) = &self.compression_algorithm {
+                    format!(", '{algorithm}'")
+                } else {
+                    String::new()
+                },
             ),
         }
     }