rasp/src/main.rs

#![allow(dead_code)]

mod common;
mod vm;

use structopt::StructOpt;
use snafu::Snafu;
use std::{fs, io::{stdin, stdout, Write, Read}, path::{Path, PathBuf}, process};

const DEFAULT_MAX_MEM: usize = 64 * 1024 * 1024;

#[derive(Snafu, Debug, Clone, PartialEq)]
enum ArgParseError<'a> {
    #[snafu(display("invalid specified size: {} (sizes may end in G, M, or K)", text))]
    InvalidSize { text: &'a str }
}

fn from_bytes_size<'a>(mut text: &'a str) -> std::result::Result<usize, ArgParseError<'a>> {
    if text.is_empty() {
        return Err(ArgParseError::InvalidSize { text })
    }
    let multiplier = match text.chars().last().unwrap() {
        'k' | 'K' => 1024,
        'm' | 'M' => 1024 * 1024,
        'g' | 'G' => 1024 * 1024 * 1024,
        _ => 1,
    };
    // trim the last character
    if multiplier != 1 {
        text = &text[0 .. text.len() - 1];
    }
    let value = text.parse::<usize>().map_err(|_| ArgParseError::InvalidSize { text })?;
    Ok(value * multiplier)
}

#[test]
fn test_bytes_size() {
    assert_eq!(from_bytes_size("64G"), Ok(64 * 1024 * 1024 * 1024));
    assert_eq!(from_bytes_size("64M"), Ok(64 * 1024 * 1024));
    assert_eq!(from_bytes_size("64K"), Ok(64 * 1024));
    assert_eq!(from_bytes_size("64g"), Ok(64 * 1024 * 1024 * 1024));
    assert_eq!(from_bytes_size("64m"), Ok(64 * 1024 * 1024));
    assert_eq!(from_bytes_size("64k"), Ok(64 * 1024));
    assert_eq!(from_bytes_size("64"), Ok(64));

    assert!(from_bytes_size("64B").is_err());
    assert!(from_bytes_size("64MB").is_err());
    assert!(from_bytes_size("64GB").is_err());
    assert!(from_bytes_size("64KB").is_err());
    assert!(from_bytes_size("64b").is_err());
    assert!(from_bytes_size("64mb").is_err());
    assert!(from_bytes_size("64gb").is_err());
    assert!(from_bytes_size("64kb").is_err());
}

#[derive(StructOpt, Debug)]
struct Options {
    // maybe some other options:
    // * debug

    /// The input file to work with.
    ///
    /// By default, the file will be executed. If -c is passed, it will not run and only compile.
    /// Supplying - for the path will read from STDIN.
    #[structopt(name = "FILE", parse(from_os_str))]
    input: PathBuf,

    /// The maximum amount of virtual memory allowed.
    #[structopt(long, parse(try_from_str = from_bytes_size))]
    max_mem: Option<usize>,

    /// The output file for generated files (preprocessed, compiled, etc).
    ///
    /// If not supplied, a name will be inferred from the input file. Supplying - for the path will
    /// output to STDOUT.
    #[structopt(short = "o", long)]
    out: Option<PathBuf>,

    /// Only run the preprocessor.
    #[structopt(short = "E", long)]
    preprocess_only: bool,

    /// Only compile the input file to an object.
    #[structopt(short = "c", long)]
    compile_only: bool,

}

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

fn get_reader(path: impl AsRef<Path>) -> Result<Box<dyn Read>> {
    if let Some("-") = path.as_ref().to_str() {
        Ok(Box::new(stdin()))
    } else {
        Ok(Box::new(fs::File::open(path.as_ref())?))
    }
}

fn get_writer(path: impl AsRef<Path>) -> Result<Box<dyn Write>> {
    if let Some("-") = path.as_ref().to_str() {
        Ok(Box::new(stdout()))
    } else {
        Ok(Box::new(fs::File::create(path.as_ref())?))
    }
}

fn main() -> Result<()> {
    use vm::{
        state::State,
        obj::assemble::AsmSession,
    };

    let opt = Options::from_args();
    let mut asm_session = AsmSession::default();
    asm_session.include(&opt.input)?;
    let object = asm_session.assemble()?;

    if opt.compile_only {
        let outfile = opt.out.clone().unwrap_or_else(|| {
            let mut outfile = opt.input.clone();
            assert!(outfile.set_extension("obj"));
            outfile
        });
        let bytes = object.to_bytes();
        let mut writer = get_writer(&outfile)?;
        writer.write(&bytes)?;
        Ok(())
    } else {
        let mut state = State::new();
        state.load_object(object, opt.max_mem.unwrap_or(DEFAULT_MAX_MEM))?;
        let status = state.exec()?;
        process::exit((status & 0xffff_ffff) as i32);
    }
}