build.rs

use std::{
    collections::HashSet,
    env,
    error::Error,
    fs::{self, File},
    io::{self, BufWriter},
    path::{Path, PathBuf, MAIN_SEPARATOR_STR},
    process::Command,
    result::Result as StdResult,
};

use std::io::Write;

use jwalk::WalkDir;
use rquickjs::{Context, Module, Runtime};

const BUNDLE_DIR: &str = "bundle";

macro_rules! info {
    ($($tokens: tt)*) => {
        println!("cargo:info={}", format!($($tokens)*))
    }
}

macro_rules! rerun_if_changed {
    ($file: expr) => {
        println!("cargo:rerun-if-changed={}", $file)
    };
}

include!("src/compiler_common.rs");

#[tokio::main]
async fn main() -> StdResult<(), Box<dyn Error>> {
    rerun_if_changed!(BUNDLE_DIR);

    let resolver = (DummyResolver,);
    let loader = (DummyLoader,);

    let rt = Runtime::new().unwrap();
    rt.set_loader(resolver, loader);
    rt.set_max_stack_size(512 * 1024);
    let ctx = Context::full(&rt).unwrap();

    let sdk_bytecode_path = Path::new("src").join("bytecode_cache.rs");
    let mut sdk_bytecode_file = BufWriter::new(File::create(sdk_bytecode_path).unwrap());

    let mut ph_map = phf_codegen::Map::<String>::new();
    let mut filenames = vec![];
    let mut total_bytes: usize = 0;

    fs::write("VERSION", env!("CARGO_PKG_VERSION")).expect("Unable to write VERSION file");

    ctx.with(|ctx| {
        for dir_ent in WalkDir::new(BUNDLE_DIR).into_iter().flatten() {
            let path = dir_ent.path();

            let path = path.strip_prefix(BUNDLE_DIR).unwrap().to_owned();
            let path_str = path.to_string_lossy().to_string();

            if path_str.starts_with("__tests__") || path.extension().unwrap_or_default() != "js" {
                continue;
            }

            #[cfg(feature = "lambda")]
            {
                if path == PathBuf::new().join("@llrt").join("test.js") {
                    continue;
                }
            }

            #[cfg(feature = "no-sdk")]
            {
                if path_str.starts_with("@aws-sdk")
                    || path_str.starts_with("@smithy")
                    || path_str.starts_with("llrt-chunk-sdk")
                {
                    continue;
                }
            }

            let source = fs::read_to_string(dir_ent.path())
                .unwrap_or_else(|_| panic!("Unable to load: {}", dir_ent.path().to_string_lossy()));

            let module_name = if !path_str.starts_with("llrt-chunk-") {
                path.with_extension("").to_string_lossy().to_string()
            } else {
                path.to_string_lossy().to_string()
            };

            info!("Compiling module: {}", module_name);

            let filename = dir_ent
                .path()
                .with_extension("lrt")
                .to_string_lossy()
                .to_string();
            filenames.push(filename.clone());

            let module = unsafe {
                Module::unsafe_declare(ctx.clone(), module_name.clone(), source).unwrap()
            };
            let bytes = module.write_object(false).unwrap();
            total_bytes += bytes.len();
            fs::write(&filename, bytes).unwrap();
            info!("Done!");

            ph_map.entry(
                module_name,
                &format!("include_bytes!(\"..{}{}\")", MAIN_SEPARATOR_STR, &filename),
            );
        }
    });

    write!(
        &mut sdk_bytecode_file,
        "// @generated by build.rs\n\npub static BYTECODE_CACHE: phf::Map<&'static str, &[u8]> = {}",
        ph_map.build()
    )
    .unwrap();
    writeln!(&mut sdk_bytecode_file, ";").unwrap();

    info!(
        "\n===============================\nUncompressed bytecode size: {}\n===============================",
        human_file_size(total_bytes)
    );

    total_bytes = compress_bytecode(BUNDLE_DIR, filenames).unwrap();

    info!(
        "\n===============================\nCompressed bytecode size: {}\n===============================",
        human_file_size(total_bytes)
    );

    Ok(())
}

fn compress_bytecode(bundles_dir: &str, source_files: Vec<String>) -> io::Result<usize> {
    let bundle_path = Path::new(bundles_dir);
    info!("Generating compression dictionary...");
    let dictionary_path = bundle_path
        .join("compression.dict")
        .to_string_lossy()
        .to_string();

    let file_count = source_files.len();
    let mut dictionary_filenames = source_files.clone();
    let mut dictionary_file_set: HashSet<String> = HashSet::from_iter(dictionary_filenames.clone());

    let mut cmd = Command::new("zstd");
    cmd.args([
        "--train",
        "--train-fastcover=steps=40",
        "--maxdict=20K",
        "-o",
        &dictionary_path,
    ]);
    if file_count < 5 {
        dictionary_file_set.retain(|file_path| {
            let metadata = fs::metadata(file_path).unwrap();
            let file_size = metadata.len();
            file_size >= 1024 // 1 kilobyte = 1024 bytes
        });
        cmd.arg("-B1K");
        dictionary_filenames = dictionary_file_set.into_iter().collect();
    }
    cmd.args(&dictionary_filenames);

    let mut cmd = cmd.args(&source_files).spawn()?;

    let exit_status = cmd.wait()?;

    if !exit_status.success() {
        return Err(io::Error::new(
            io::ErrorKind::Other,
            "Failed to generate compression dictionary",
        ));
    }

    let mut total_size = 0;
    let tmp_dir = env::temp_dir();

    for filename in source_files {
        info!("Compressing {}...", filename);

        let tmp_filename = tmp_dir
            .join(nanoid::nanoid!())
            .to_string_lossy()
            .to_string();

        fs::copy(&filename, &tmp_filename)?;

        let uncompressed_file_size = PathBuf::from(&filename).metadata().unwrap().len() as u32;

        let output = Command::new("zstd")
            .args([
                "--ultra",
                "-22",
                "-f",
                "-D",
                &dictionary_path,
                &tmp_filename,
                "-o",
                &filename,
            ])
            .output()?;

        if !output.status.success() {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "Failed to compress file",
            ));
        }

        let bytes = fs::read(&filename)?;
        let mut compressed = Vec::with_capacity(4);
        compressed.extend_from_slice(&uncompressed_file_size.to_le_bytes());
        compressed.extend_from_slice(&bytes);
        fs::write(&filename, compressed)?;

        let compressed_file_size = PathBuf::from(&filename).metadata().unwrap().len() as usize;

        total_size += compressed_file_size;
    }

    Ok(total_size)
}