dioxus/packages/cli/src/assets.rs

use crate::builder::{
    BuildMessage, BuildRequest, MessageSource, MessageType, Stage, UpdateBuildProgress, UpdateStage,
};
use crate::Result;
use anyhow::Context;
use brotli::enc::BrotliEncoderParams;
use futures_channel::mpsc::UnboundedSender;
use manganis_cli_support::{process_file, AssetManifest, AssetManifestExt, AssetType};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::fs;
use std::path::Path;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use std::{ffi::OsString, path::PathBuf};
use std::{fs::File, io::Write};
use tracing::Level;
use walkdir::WalkDir;

/// The temp file name for passing manganis json from linker to current exec.
pub const MG_JSON_OUT: &str = "mg-out";

pub fn asset_manifest(build: &BuildRequest) -> AssetManifest {
    let file_path = build.target_out_dir().join(MG_JSON_OUT);
    let read = fs::read_to_string(&file_path).unwrap();
    _ = fs::remove_file(file_path);
    let json: Vec<String> = serde_json::from_str(&read).unwrap();

    AssetManifest::load(json)
}

/// Create a head file that contains all of the imports for assets that the user project uses
pub fn create_assets_head(build: &BuildRequest, manifest: &AssetManifest) -> Result<()> {
    let out_dir = build.target_out_dir();
    std::fs::create_dir_all(&out_dir)?;
    let mut file = File::create(out_dir.join("__assets_head.html"))?;
    file.write_all(manifest.head().as_bytes())?;
    Ok(())
}

/// Process any assets collected from the binary
pub(crate) fn process_assets(
    build: &BuildRequest,
    manifest: &AssetManifest,
    progress: &mut UnboundedSender<UpdateBuildProgress>,
) -> anyhow::Result<()> {
    let static_asset_output_dir = build.target_out_dir();

    std::fs::create_dir_all(&static_asset_output_dir)
        .context("Failed to create static asset output directory")?;

    let assets_finished = Arc::new(AtomicUsize::new(0));
    let assets = manifest.assets();
    let asset_count = assets.len();
    assets.par_iter().try_for_each_init(
        || progress.clone(),
        move |progress, asset| {
            if let AssetType::File(file_asset) = asset {
                match process_file(file_asset, &static_asset_output_dir) {
                    Ok(_) => {
                        // Update the progress
                        _ = progress.start_send(UpdateBuildProgress {
                            stage: Stage::OptimizingAssets,
                            update: UpdateStage::AddMessage(BuildMessage {
                                level: Level::INFO,
                                message: MessageType::Text(format!(
                                    "Optimized static asset {}",
                                    file_asset
                                )),
                                source: MessageSource::Build,
                            }),
                        });
                        let assets_finished =
                            assets_finished.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                        _ = progress.start_send(UpdateBuildProgress {
                            stage: Stage::OptimizingAssets,
                            update: UpdateStage::SetProgress(
                                assets_finished as f64 / asset_count as f64,
                            ),
                        });
                    }
                    Err(err) => {
                        tracing::error!("Failed to copy static asset: {}", err);
                        return Err(err);
                    }
                }
            }
            Ok::<(), anyhow::Error>(())
        },
    )?;

    Ok(())
}

/// A guard that sets up the environment for the web renderer to compile in. This guard sets the location that assets will be served from
pub(crate) struct AssetConfigDropGuard;

impl AssetConfigDropGuard {
    pub fn new(base_path: Option<&str>) -> Self {
        // Set up the collect asset config
        let base = match base_path {
            Some(base) => format!("/{}/", base.trim_matches('/')),
            None => "/".to_string(),
        };
        manganis_cli_support::Config::default()
            .with_assets_serve_location(base)
            .save();
        Self {}
    }
}

impl Drop for AssetConfigDropGuard {
    fn drop(&mut self) {
        // Reset the config
        manganis_cli_support::Config::default().save();
    }
}

pub(crate) fn copy_dir_to(
    src_dir: PathBuf,
    dest_dir: PathBuf,
    pre_compress: bool,
) -> std::io::Result<()> {
    let entries = std::fs::read_dir(&src_dir)?;
    let mut children: Vec<std::thread::JoinHandle<std::io::Result<()>>> = Vec::new();

    for entry in entries.flatten() {
        let entry_path = entry.path();
        let path_relative_to_src = entry_path.strip_prefix(&src_dir).unwrap();
        let output_file_location = dest_dir.join(path_relative_to_src);
        children.push(std::thread::spawn(move || {
            if entry.file_type()?.is_dir() {
                // If the file is a directory, recursively copy it into the output directory
                if let Err(err) =
                    copy_dir_to(entry_path.clone(), output_file_location, pre_compress)
                {
                    tracing::error!(
                        "Failed to pre-compress directory {}: {}",
                        entry_path.display(),
                        err
                    );
                }
            } else {
                // Make sure the directory exists
                std::fs::create_dir_all(output_file_location.parent().unwrap())?;
                // Copy the file to the output directory
                std::fs::copy(&entry_path, &output_file_location)?;

                // Then pre-compress the file if needed
                if pre_compress {
                    if let Err(err) = pre_compress_file(&output_file_location) {
                        tracing::error!(
                            "Failed to pre-compress static assets {}: {}",
                            output_file_location.display(),
                            err
                        );
                    }
                    // If pre-compression isn't enabled, we should remove the old compressed file if it exists
                } else if let Some(compressed_path) = compressed_path(&output_file_location) {
                    _ = std::fs::remove_file(compressed_path);
                }
            }

            Ok(())
        }));
    }
    for child in children {
        child.join().unwrap()?;
    }
    Ok(())
}

/// Get the path to the compressed version of a file
fn compressed_path(path: &Path) -> Option<PathBuf> {
    let new_extension = match path.extension() {
        Some(ext) => {
            if ext.to_string_lossy().to_lowercase().ends_with("br") {
                return None;
            }
            let mut ext = ext.to_os_string();
            ext.push(".br");
            ext
        }
        None => OsString::from("br"),
    };
    Some(path.with_extension(new_extension))
}

/// pre-compress a file with brotli
pub(crate) fn pre_compress_file(path: &Path) -> std::io::Result<()> {
    let Some(compressed_path) = compressed_path(path) else {
        return Ok(());
    };
    let file = std::fs::File::open(path)?;
    let mut stream = std::io::BufReader::new(file);
    let mut buffer = std::fs::File::create(compressed_path)?;
    let params = BrotliEncoderParams::default();
    brotli::BrotliCompress(&mut stream, &mut buffer, &params)?;
    Ok(())
}

/// pre-compress all files in a folder
pub(crate) fn pre_compress_folder(path: &Path, pre_compress: bool) -> std::io::Result<()> {
    let walk_dir = WalkDir::new(path);
    for entry in walk_dir.into_iter().filter_map(|e| e.ok()) {
        let entry_path = entry.path();
        if entry_path.is_file() {
            if pre_compress {
                if let Err(err) = pre_compress_file(entry_path) {
                    tracing::error!("Failed to pre-compress file {entry_path:?}: {err}");
                }
            }
            // If pre-compression isn't enabled, we should remove the old compressed file if it exists
            else if let Some(compressed_path) = compressed_path(entry_path) {
                _ = std::fs::remove_file(compressed_path);
            }
        }
    }
    Ok(())
}