refactor: cleanup code & include logging

src/cli/build.rs

@@ -1,3 +1,4 @@
+use console::Style;
 use std::{
     env,
     path::{Path, PathBuf},
@@ -11,10 +12,26 @@ use tokio::{
 use anyhow::Result;
 use mlua::Compiler as LuaCompiler;
 
+/**
+    Compiles and embeds the bytecode of a requested lua file to form a standalone binary,
+    then writes it to an output file, with the required permissions.
+*/
+#[allow(clippy::similar_names)]
 pub async fn build_standalone<T: AsRef<Path> + Into<PathBuf>>(
+    script_path: String,
     output_path: T,
     code: impl AsRef<[u8]>,
 ) -> Result<ExitCode> {
+    let log_output_path = output_path.as_ref().display();
+
+    let prefix_style = Style::new().green().bold();
+    let compile_prefix = prefix_style.apply_to("Compile");
+    let bytecode_prefix = prefix_style.apply_to("Bytecode");
+    let write_prefix = prefix_style.apply_to("Write");
+    let compiled_prefix = prefix_style.apply_to("Compiled");
+
+    println!("{compile_prefix} {script_path}");
+
     // First, we read the contents of the lune interpreter as our starting point
     let mut patched_bin = fs::read(env::current_exe()?).await?;
     let base_bin_offset = u64::try_from(patched_bin.len())?;
@@ -30,13 +47,15 @@ pub async fn build_standalone<T: AsRef<Path> + Into<PathBuf>>(
         .set_debug_level(0)
         .compile(code);
 
+    println!("  {bytecode_prefix} {script_path}");
+
     patched_bin.append(&mut bytecode.clone());
 
     let mut meta = base_bin_offset.to_ne_bytes().to_vec();
 
     // Include metadata in the META chunk, each field is 8 bytes
     meta.append(&mut (bytecode.len() as u64).to_ne_bytes().to_vec()); // Size of bytecode, used to calculate end offset at runtime
-    meta.append(&mut 1_u64.to_ne_bytes().to_vec()); // Number of files, padded with null bytes
+    meta.append(&mut 1_u64.to_ne_bytes().to_vec()); // Number of files, padded with null bytes - for future use
 
     patched_bin.append(&mut meta);
 
@@ -57,5 +76,9 @@ pub async fn build_standalone<T: AsRef<Path> + Into<PathBuf>>(
     #[cfg(target_family = "windows")]
     fs::write(&output_path, &patched_bin).await?;
 
+    println!("  {write_prefix} {log_output_path}");
+
+    println!("{compiled_prefix} {log_output_path}");
+
     Ok(ExitCode::SUCCESS)
 }

src/cli/mod.rs

@@ -1,13 +1,9 @@
-use std::{env, fmt::Write as _, ops::ControlFlow, path::PathBuf, process::ExitCode, sync::Mutex};
+use std::{env, fmt::Write as _, path::PathBuf, process::ExitCode};
 
 use anyhow::{Context, Result};
 use clap::Parser;
 
 use lune::Lune;
-use rayon::{
-    iter::{IndexedParallelIterator, ParallelIterator},
-    slice::ParallelSlice,
-};
 use tokio::{
     fs::read as read_to_vec,
     io::{stdin, AsyncReadExt},
@@ -89,121 +85,9 @@ impl Cli {
 
     #[allow(clippy::too_many_lines)]
     pub async fn run(self) -> Result<ExitCode> {
-        // Signature which is only present in standalone lune binaries
-        let signature: Vec<u8> = vec![0x4f, 0x3e, 0xf8, 0x41, 0xc3, 0x3a, 0x52, 0x16];
-        // Read the current lune binary to memory
-        let bin = read_to_vec(env::current_exe()?).await?;
-
-        let is_standalone = bin[bin.len() - signature.len()..bin.len()] == signature;
-
-        if is_standalone {
-            let bytecode_offset = Mutex::new(0);
-            let bytecode_size = Mutex::new(0);
-
-            // standalone binary structure (reversed, 8 bytes per field)
-            // [0] => signature
-            // ----------------
-            // -- META Chunk --
-            // [1] => file count
-            // [2] => bytecode size
-            // [3] => bytecode offset
-            // ----------------
-            // -- MISC Chunk --
-            // [4..n] => bytecode (variable size)
-            // ----------------
-            // NOTE: All integers are 8 byte unsigned 64 bit (u64's).
-
-            // The rchunks will have unequally sized sections in the beginning
-            // but that doesn't matter to us because we don't need anything past the
-            // middle chunks where the bytecode is stored
-            bin.par_rchunks(signature.len())
-                .enumerate()
-                .try_for_each(|(idx, chunk)| {
-                    let mut bytecode_offset = bytecode_offset.lock().unwrap();
-                    let mut bytecode_size = bytecode_size.lock().unwrap();
-
-                    if *bytecode_offset != 0 && *bytecode_size != 0 {
-                        return ControlFlow::Break(());
-                    }
-
-                    if idx == 0 && chunk != signature {
-                        // Binary is guaranteed to be standalone, we've confirmed this before
-                        unreachable!("expected proper signature for standalone binary")
-                    }
-
-                    if idx == 3 {
-                        *bytecode_offset = u64::from_ne_bytes(chunk.try_into().unwrap());
-                    }
-
-                    if idx == 2 {
-                        *bytecode_size = u64::from_ne_bytes(chunk.try_into().unwrap());
-                    }
-
-                    ControlFlow::Continue(())
-                });
-
-            let bytecode_offset_inner = bytecode_offset.into_inner().unwrap();
-            let bytecode_size_inner = bytecode_size.into_inner().unwrap();
-
-            // If we were able to retrieve the required metadata, we load
-            // and execute the bytecode
-            // if bytecode_offset_inner != 0 && bytecode_size_inner != 0 {
-            // FIXME: Passing arguments does not work like it should, because the first
-            // argument provided is treated as the script path. We should probably also not
-            // allow any runner functionality within standalone binaries
-
-            let mut reserved_args = Vec::new();
-
-            macro_rules! include_reserved_args {
-                    ($($arg_bool:expr=> $mapping:literal),*) => {
-                        $(
-                            if $arg_bool {
-                                reserved_args.push($mapping.to_string())
-                            }
-                        )*
-                    };
-                }
-
-            let mut real_args = Vec::new();
-
-            if let Some(first_arg) = self.script_path {
-                real_args.push(first_arg);
-            }
-
-            include_reserved_args! {
-                self.setup => "--setup",
-                self.generate_docs_file => "--generate-docs-file",
-                self.generate_selene_types => "--generate-selene-types",
-                self.generate_luau_types => "--generate-luau-types",
-                self.list => "--list",
-                self.build => "--build"
-            }
-
-            real_args.append(&mut reserved_args);
-            real_args.append(&mut self.script_args.clone());
-
-            let result = Lune::new()
-                .with_args(real_args) // TODO: args should also include lune reserved ones
-                .run(
-                    "STANDALONE",
-                    &bin[usize::try_from(bytecode_offset_inner)?
-                        ..usize::try_from(bytecode_offset_inner + bytecode_size_inner)?],
-                )
-                .await;
-
-            return Ok(match result {
-                Err(err) => {
-                    eprintln!("{err}");
-                    ExitCode::FAILURE
-                }
-                Ok(code) => code,
-            });
-            // }
-        }
-
         // List files in `lune` and `.lune` directories, if wanted
         // This will also exit early and not run anything else
-        if self.list && !is_standalone {
+        if self.list {
             let sorted_relative = find_lune_scripts(false).await.map(sort_lune_scripts);
 
             let sorted_home_dir = find_lune_scripts(true).await.map(sort_lune_scripts);
@@ -271,7 +155,6 @@ impl Cli {
             return repl::show_interface().await;
         }
 
-        if !is_standalone {
         // Figure out if we should read from stdin or from a file,
         // reading from stdin is marked by passing a single "-"
         // (dash) as the script name to run to the cli
@@ -295,18 +178,18 @@ impl Cli {
         if self.build {
             let output_path =
                 PathBuf::from(script_path.clone()).with_extension(env::consts::EXE_EXTENSION);
-                println!(
-                    "Building {script_path} to {}",
-                    output_path.to_string_lossy()
-                );
 
-                return Ok(match build_standalone(output_path, script_contents).await {
+            println!("Building {script_path} to {}...\n", output_path.display());
+
+            return Ok(
+                match build_standalone(script_path, output_path, script_contents).await {
                     Ok(exitcode) => exitcode,
                     Err(err) => {
                         eprintln!("{err}");
                         ExitCode::FAILURE
                     }
-                });
+                },
+            );
         }
 
         // Create a new lune object with all globals & run the script
@@ -314,15 +197,12 @@ impl Cli {
             .with_args(self.script_args)
             .run(&script_display_name, strip_shebang(script_contents))
             .await;
-            return Ok(match result {
+        Ok(match result {
             Err(err) => {
                 eprintln!("{err}");
                 ExitCode::FAILURE
             }
             Ok(code) => code,
-            });
+        })
-        }
-
-        Ok(ExitCode::SUCCESS)
     }
 }

src/executor.rs

@@ -0,0 +1,123 @@
+use std::{
+    collections::VecDeque,
+    env,
+    ops::ControlFlow,
+    process::{self, ExitCode},
+    sync::Mutex,
+};
+
+use lune::Lune;
+
+use anyhow::Result;
+use rayon::{
+    iter::{IndexedParallelIterator, ParallelIterator},
+    slice::ParallelSlice,
+};
+use tokio::fs::read as read_to_vec;
+
+/**
+    Returns information about whether the execution environment is standalone
+    or not, the standalone binary signature, and the contents of the binary.
+*/
+pub async fn check_env() -> (bool, Vec<u8>, Vec<u8>) {
+    // Signature which is only present in standalone lune binaries
+    let signature: Vec<u8> = vec![0x4f, 0x3e, 0xf8, 0x41, 0xc3, 0x3a, 0x52, 0x16];
+
+    // Read the current lune binary to memory
+    let bin = if let Ok(contents) = read_to_vec(match env::current_exe() {
+        Ok(path) => path,
+        Err(err) => {
+            eprintln!(
+                "WARN: Couldn't get path to currently running lune executable; err: {}",
+                err.kind()
+            );
+            process::exit(1);
+        }
+    })
+    .await
+    {
+        contents
+    } else {
+        Vec::new()
+    };
+
+    let is_standalone = !bin.is_empty() && bin[bin.len() - signature.len()..bin.len()] == signature;
+
+    (is_standalone, signature, bin)
+}
+
+/**
+    Discovers, loads and executes the bytecode contained in a standalone binary.
+*/
+pub async fn run_standalone(signature: Vec<u8>, bin: Vec<u8>) -> Result<ExitCode> {
+    let bytecode_offset = Mutex::new(0);
+    let bytecode_size = Mutex::new(0);
+
+    // standalone binary structure (reversed, 8 bytes per field)
+    // [0] => signature
+    // ----------------
+    // -- META Chunk --
+    // [1] => file count
+    // [2] => bytecode size
+    // [3] => bytecode offset
+    // ----------------
+    // -- MISC Chunk --
+    // [4..n] => bytecode (variable size)
+    // ----------------
+    // NOTE: All integers are 8 byte unsigned 64 bit (u64's).
+
+    // The rchunks will have unequally sized sections in the beginning
+    // but that doesn't matter to us because we don't need anything past the
+    // middle chunks where the bytecode is stored
+    bin.par_rchunks(signature.len())
+        .enumerate()
+        .try_for_each(|(idx, chunk)| {
+            let mut bytecode_offset = bytecode_offset.lock().unwrap();
+            let mut bytecode_size = bytecode_size.lock().unwrap();
+
+            if *bytecode_offset != 0 && *bytecode_size != 0 {
+                return ControlFlow::Break(());
+            }
+
+            if idx == 0 && chunk != signature {
+                // Binary is guaranteed to be standalone, we've confirmed this before
+                unreachable!("expected proper signature for standalone binary")
+            }
+
+            if idx == 3 {
+                *bytecode_offset = u64::from_ne_bytes(chunk.try_into().unwrap());
+            }
+
+            if idx == 2 {
+                *bytecode_size = u64::from_ne_bytes(chunk.try_into().unwrap());
+            }
+
+            ControlFlow::Continue(())
+        });
+
+    let bytecode_offset_inner = bytecode_offset.into_inner().unwrap();
+    let bytecode_size_inner = bytecode_size.into_inner().unwrap();
+
+    // If we were able to retrieve the required metadata, we load
+    // and execute the bytecode
+
+    let mut args = env::args().collect::<VecDeque<String>>();
+    args.pop_front();
+
+    let result = Lune::new()
+        .with_args(args)
+        .run(
+            "STANDALONE",
+            &bin[usize::try_from(bytecode_offset_inner)?
+                ..usize::try_from(bytecode_offset_inner + bytecode_size_inner)?],
+        )
+        .await;
+
+    Ok(match result {
+        Err(err) => {
+            eprintln!("{err}");
+            ExitCode::FAILURE
+        }
+        Ok(code) => code,
+    })
+}

src/main.rs

@@ -13,9 +13,11 @@ use std::process::ExitCode;
 use clap::Parser;
 
 pub(crate) mod cli;
+pub(crate) mod executor;
 
 use cli::Cli;
 use console::style;
+
 #[cfg(not(target_env = "msvc"))]
 use tikv_jemallocator::Jemalloc;
 
@@ -32,6 +34,15 @@ async fn main() -> ExitCode {
         .with_timer(tracing_subscriber::fmt::time::uptime())
         .with_level(true)
         .init();
+
+    let (is_standalone, signature, bin) = executor::check_env().await;
+
+    if is_standalone {
+        // It's fine to unwrap here since we don't want to continue
+        // if something fails
+        return executor::run_standalone(signature, bin).await.unwrap();
+    }
+
     match Cli::parse().run().await {
         Ok(code) => code,
         Err(err) => {