From ef83cf7ef3221e6d7f94d1103a213de9a87f795d Mon Sep 17 00:00:00 2001
From: Alek Ratzloff <alekratz@gmail.com>
Date: Tue, 25 Feb 2020 14:55:14 -0500
Subject: [PATCH] Add max_mem and size parsing for arg parsing

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>
---
 src/main.rs | 60 ++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 53 insertions(+), 7 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 155b225..5509556 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,7 +4,54 @@ mod common;
 mod vm;
 
 use structopt::StructOpt;
-use std::{fs, io, process, path::{Path, PathBuf}};
+use snafu::Snafu;
+use std::{fs, process, path::PathBuf};
+
+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 {
@@ -12,14 +59,13 @@ struct Options {
     // * debug
     #[structopt(name = "FILE", parse(from_os_str))]
     input: PathBuf,
+
+    #[structopt(short, long, parse(try_from_str = from_bytes_size))]
+    max_mem: Option<usize>,
 }
 
 type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
-fn get_input_string(path: impl AsRef<Path>) -> io::Result<String> {
-    fs::read_to_string(path.as_ref())
-}
-
 fn main() -> Result<()> {
     use vm::{
         state::State,
@@ -29,7 +75,7 @@ fn main() -> Result<()> {
         }
     };
     let opt = Options::from_args();
-    let text = get_input_string(&opt.input)?;
+    let text = fs::read_to_string(&opt.input)?;
     let lexerdef = lexer::lexerdef();
     let lexer = lexerdef.lexer(&text);
     let (res, errors) = parser::parse(&lexer);
@@ -44,7 +90,7 @@ fn main() -> Result<()> {
     let mut asm = Asm::default();
     let object = res.assemble(&mut asm)?;
     let mut state = State::new();
-    state.load_object(object, 64 * 1024 * 1024)?;
+    state.load_object(object, opt.max_mem.unwrap_or(DEFAULT_MAX_MEM))?;
     let status = state.exec()?;
     process::exit((status & 0xffff_ffff) as i32);
 }