checksmix 0.4.3

use checksmix::{
    MMix, MMixAssembler, MmoDecoder, Stop, ValueFormat, entry_point, start_program, write_image,
};
use clap::{Parser, Subcommand};
use std::fs;
use std::io::{self, Write};
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::process;
use tracing_subscriber::{EnvFilter, fmt};

/// The status of a process whose stdout reader has gone: 128 plus SIGPIPE,
/// what a shell reports for a C tool that SIGPIPE killed.
const BROKEN_PIPE_STATUS: i32 = 141;

/// Ends the process after a failed write to stdout: quietly with
/// `BROKEN_PIPE_STATUS` once the reader has gone, else with the error on
/// stderr and status 1.
fn stdout_failed(err: &io::Error) -> ! {
    if err.kind() == io::ErrorKind::BrokenPipe {
        process::exit(BROKEN_PIPE_STATUS);
    }
    eprintln!("error writing to stdout: {err}");
    process::exit(1);
}

macro_rules! outln {
    ($($arg:tt)*) => {
        if let Err(err) = writeln!(io::stdout(), $($arg)*) {
            stdout_failed(&err);
        }
    };
}

#[derive(Parser, Debug)]
#[command(
    name = "checksmix",
    about = "Run MMIX programs and assemblers",
    version,
    author,
    args_conflicts_with_subcommands = true,
    subcommand_negates_reqs = true
)]
struct Cli {
    #[command(subcommand)]
    command: Option<Command>,

    /// Display register values as unsigned decimals (run mode)
    #[arg(long)]
    unsigned: bool,

    /// Stop after N instructions if the program has not halted, and exit 124 (run mode)
    #[arg(long, value_name = "N")]
    max_steps: Option<NonZeroUsize>,

    /// Program file(s) to execute. A single file dispatches by extension
    /// (.mms/.mmo); multiple files must all be .mms and are assembled into
    /// one shared symbol space before execution.
    #[arg(required = true, num_args = 1..)]
    program_files: Vec<String>,
}

#[derive(Subcommand, Debug)]
enum Command {
    /// Run a program (equivalent to omitting the subcommand)
    Run {
        /// Display register values as unsigned decimals
        #[arg(long)]
        unsigned: bool,
        /// Stop after N instructions if the program has not halted, and exit 124
        #[arg(long, value_name = "N")]
        max_steps: Option<NonZeroUsize>,
        /// Program file(s) to execute
        #[arg(required = true, num_args = 1..)]
        program_files: Vec<String>,
    },
    /// Parse and encode .mms source(s); prints any assembler warning to
    /// stderr, errors on failure, and exits 0 on success regardless
    Check {
        /// MMIX assembly source file(s)
        #[arg(required = true, num_args = 1.., value_name = "FILE.mms")]
        files: Vec<PathBuf>,
    },
    /// Assemble one or more .mms sources into a .mmo object file
    Build {
        /// Output .mmo file (default: first input's path with .mmo
        /// extension, beside the source)
        #[arg(short = 'o', long, value_name = "OUT.mmo")]
        output: Option<PathBuf>,
        /// MMIX assembly source file(s)
        #[arg(required = true, num_args = 1.., value_name = "FILE.mms")]
        files: Vec<PathBuf>,
    },
}

fn main() {
    fmt().with_env_filter(EnvFilter::from_default_env()).init();

    let cli = Cli::parse();
    match cli.command {
        Some(Command::Run {
            unsigned,
            max_steps,
            program_files,
        }) => {
            let vfmt = if unsigned {
                ValueFormat::Unsigned
            } else {
                ValueFormat::Signed
            };
            dispatch_run(&program_files, vfmt, max_steps);
        }
        Some(Command::Check { files }) => cmd_check(&files),
        Some(Command::Build { output, files }) => cmd_build(&files, output.as_deref()),
        None => {
            let vfmt = if cli.unsigned {
                ValueFormat::Unsigned
            } else {
                ValueFormat::Signed
            };
            dispatch_run(&cli.program_files, vfmt, cli.max_steps);
        }
    }
}

fn dispatch_run(
    program_files: &[String],
    value_format: ValueFormat,
    max_steps: Option<NonZeroUsize>,
) {
    if program_files.len() == 1 {
        let program_file = &program_files[0];
        let path = Path::new(program_file);
        let extension = path.extension().and_then(|s| s.to_str()).unwrap_or("");

        match extension {
            "mms" => run_mms(program_files, value_format, max_steps),
            "mmo" => run_mmo(program_file, value_format, max_steps),
            _ => {
                eprintln!(
                    "Unknown file extension: .{}",
                    if extension.is_empty() {
                        "(none)"
                    } else {
                        extension
                    }
                );
                eprintln!("Supported extensions: .mms, .mmo");
                process::exit(1);
            }
        }
    } else {
        for f in program_files {
            let ext = Path::new(f)
                .extension()
                .and_then(|s| s.to_str())
                .unwrap_or("");
            if ext != "mms" {
                eprintln!(
                    "When passing multiple inputs, all files must be .mms (got '{}')",
                    f
                );
                process::exit(1);
            }
        }
        run_mms(program_files, value_format, max_steps);
    }
}

fn assemble_sources(paths: &[PathBuf]) -> Result<MMixAssembler, String> {
    let reader = |p: &Path| std::fs::read_to_string(p);
    let mut sources: Vec<(String, String)> = Vec::new();
    for path in paths {
        let src = fs::read_to_string(path)
            .map_err(|err| format!("error reading '{}': {}", path.display(), err))?;
        let base = path.parent().unwrap_or_else(|| Path::new("."));
        let units = MMixAssembler::resolve_includes(&src, &path.to_string_lossy(), base, &reader)?;
        // A trimmed-away segment (blank space, or an INCLUDE resolving to
        // nothing) leaves this input with no translation unit at all.
        if units.is_empty() {
            return Err(format!("'{}' contributed no source", path.display()));
        }
        sources.extend(units);
    }
    let (first_name, first_src) = &sources[0];
    let mut asm = MMixAssembler::new(first_src, first_name);
    for (name, src) in sources.iter().skip(1) {
        asm.add_source(src, name);
    }
    asm.parse()?;
    for warning in asm.warnings() {
        eprintln!("{}", warning);
    }
    Ok(asm)
}

fn cmd_check(files: &[PathBuf]) {
    let asm = assemble_sources(files).unwrap_or_else(|e| {
        eprintln!("{}", e);
        process::exit(1);
    });
    let _ = asm.generate_object_code();
}

fn cmd_build(files: &[PathBuf], output: Option<&Path>) {
    let asm = assemble_sources(files).unwrap_or_else(|e| {
        eprintln!("{}", e);
        process::exit(1);
    });
    if asm.instructions.is_empty() {
        eprintln!("error: no instructions to assemble");
        process::exit(1);
    }
    let object_code = asm.generate_object_code();
    let out_path = output
        .map(|p| p.to_path_buf())
        .unwrap_or_else(|| files[0].with_extension("mmo"));
    fs::write(&out_path, &object_code).unwrap_or_else(|err| {
        eprintln!("error writing '{}': {}", out_path.display(), err);
        process::exit(1);
    });
    outln!("{}", out_path.display());
}

/// Run to completion, or to an instruction budget, and exit. A halted
/// program exits with its own status; a budget that runs out first prints
/// the final state without "Execution completed." and exits 124.
fn finish_run(mmix: &mut MMix, value_format: ValueFormat, max_steps: Option<NonZeroUsize>) -> ! {
    let (count, stop) = match max_steps {
        Some(budget) => mmix.run_bounded(budget.get()),
        // Calling `run` here, rather than `run_bounded(usize::MAX)`, keeps
        // its own tracing span distinct from a budgeted run's.
        None => (mmix.run(), Stop::Halted),
    };
    outln!();
    outln!("Executed {} instructions", count);
    outln!();

    outln!("=== Final Machine State ===");
    outln!("{}", mmix.display_with(value_format));
    outln!();

    match stop {
        Stop::Halted => {
            outln!("Execution completed.");
            process::exit(mmix.get_exit_code() as i32);
        }
        Stop::BudgetExhausted => {
            let noun = if count == 1 {
                "instruction"
            } else {
                "instructions"
            };
            eprintln!(
                "program did not halt within {} {}; @ = #{:x}",
                count,
                noun,
                mmix.get_pc()
            );
            process::exit(124);
        }
        _ => unreachable!("run_bounded stops only by halting or exhausting its budget"),
    }
}

fn run_mms(filenames: &[String], value_format: ValueFormat, max_steps: Option<NonZeroUsize>) {
    outln!("=== MMIX Assembler ===");
    if filenames.len() == 1 {
        outln!("=== Parsing assembly from: {} ===", filenames[0]);
    } else {
        outln!("=== Parsing {} assembly inputs ===", filenames.len());
        for f in filenames {
            outln!("  {}", f);
        }
    }
    outln!();

    let paths: Vec<PathBuf> = filenames.iter().map(PathBuf::from).collect();
    let assembler = assemble_sources(&paths).unwrap_or_else(|e| {
        eprintln!("Error: {}", e);
        process::exit(1);
    });

    outln!("Assembly parsed successfully");
    outln!();

    let mut mmix = MMix::new();

    write_image(&mut mmix, &assembler);
    start_program(&mut mmix, entry_point(&assembler));

    outln!("=== Initial Machine State ===");
    outln!("{}", mmix.display_with(value_format));
    outln!();

    outln!("=== Executing Program ===");
    finish_run(&mut mmix, value_format, max_steps);
}

fn run_mmo(filename: &str, value_format: ValueFormat, max_steps: Option<NonZeroUsize>) {
    let data = fs::read(filename).unwrap_or_else(|err| {
        eprintln!("Error reading file '{}': {}", filename, err);
        process::exit(1);
    });

    outln!("=== MMIX Computer ===");
    outln!("=== Loading program from: {} ===", filename);
    outln!();

    let mut mmix = MMix::new();

    let decoder = MmoDecoder::new(data);
    let entry = decoder.load(&mut mmix).unwrap_or_else(|err| {
        eprintln!("{}", err);
        process::exit(1);
    });

    start_program(&mut mmix, entry);

    outln!("Loaded object file (entry point: 0x{:X})", entry);
    outln!();

    outln!("=== Initial Machine State ===");
    outln!("{}", mmix.display_with(value_format));
    outln!();

    outln!("=== Executing Program ===");
    finish_run(&mut mmix, value_format, max_steps);
}

#[cfg(test)]
mod tests {
    use super::*;
    use clap::Parser;

    #[test]
    fn no_subcommand_routes_to_run() {
        let cli = Cli::try_parse_from(["checksmix", "file.mms"]).unwrap();
        assert!(cli.command.is_none());
        assert_eq!(cli.program_files, vec!["file.mms"]);
        assert!(!cli.unsigned);
    }

    #[test]
    fn no_subcommand_unsigned_flag() {
        let cli = Cli::try_parse_from(["checksmix", "--unsigned", "file.mms"]).unwrap();
        assert!(cli.command.is_none());
        assert!(cli.unsigned);
        assert_eq!(cli.program_files, vec!["file.mms"]);
    }

    #[test]
    fn run_subcommand_parses() {
        let cli = Cli::try_parse_from(["checksmix", "run", "file.mms"]).unwrap();
        match cli.command {
            Some(Command::Run {
                unsigned,
                program_files,
                ..
            }) => {
                assert!(!unsigned);
                assert_eq!(program_files, vec!["file.mms"]);
            }
            _ => panic!("expected Run"),
        }
    }

    #[test]
    fn run_subcommand_unsigned() {
        let cli = Cli::try_parse_from(["checksmix", "run", "--unsigned", "file.mms"]).unwrap();
        match cli.command {
            Some(Command::Run { unsigned, .. }) => assert!(unsigned),
            _ => panic!("expected Run"),
        }
    }

    #[test]
    fn run_subcommand_multiple_files() {
        let cli = Cli::try_parse_from(["checksmix", "run", "a.mms", "b.mms"]).unwrap();
        match cli.command {
            Some(Command::Run { program_files, .. }) => {
                assert_eq!(program_files, vec!["a.mms", "b.mms"]);
            }
            _ => panic!("expected Run"),
        }
    }

    #[test]
    fn check_subcommand_parses() {
        let cli = Cli::try_parse_from(["checksmix", "check", "a.mms", "b.mms"]).unwrap();
        match cli.command {
            Some(Command::Check { files }) => {
                assert_eq!(files, vec![PathBuf::from("a.mms"), PathBuf::from("b.mms")]);
            }
            _ => panic!("expected Check"),
        }
    }

    #[test]
    fn build_subcommand_short_output() {
        let cli = Cli::try_parse_from(["checksmix", "build", "-o", "out.mmo", "a.mms"]).unwrap();
        match cli.command {
            Some(Command::Build { output, files }) => {
                assert_eq!(output, Some(PathBuf::from("out.mmo")));
                assert_eq!(files, vec![PathBuf::from("a.mms")]);
            }
            _ => panic!("expected Build"),
        }
    }

    #[test]
    fn build_subcommand_long_output() {
        let cli =
            Cli::try_parse_from(["checksmix", "build", "--output", "out.mmo", "a.mms"]).unwrap();
        match cli.command {
            Some(Command::Build { output, .. }) => {
                assert_eq!(output, Some(PathBuf::from("out.mmo")));
            }
            _ => panic!("expected Build"),
        }
    }

    #[test]
    fn build_subcommand_default_output() {
        let cli = Cli::try_parse_from(["checksmix", "build", "a.mms"]).unwrap();
        match cli.command {
            Some(Command::Build { output, .. }) => {
                assert_eq!(output, None);
            }
            _ => panic!("expected Build"),
        }
    }

    #[test]
    fn unsigned_rejected_on_check() {
        // --unsigned is not defined on the check subcommand
        assert!(Cli::try_parse_from(["checksmix", "check", "--unsigned", "a.mms"]).is_err());
    }

    #[test]
    fn output_flag_rejected_on_run() {
        // -o is only defined on build
        assert!(Cli::try_parse_from(["checksmix", "run", "-o", "out.mmo", "a.mms"]).is_err());
    }
}