Simplify string arguments for control statements

Some control statements (e.g. '.incbin', '.asciiz') only require a
double-quoted string as an argument. In fact, for these functions
there's only one argument required, which is this string one. Given this
fact, the parsing on these functions don't have to go through
the (expensive) general argument parsing function, and they can simply
assume that a double-quoted string will be provided.

Signed-off-by: Miquel Sabaté Solà <mikisabate@gmail.com>

2 files changed, 51 insertions, 28 deletions

diff --git a/lib/xixanta/src/opcodes.rs b/lib/xixanta/src/opcodes.rs
index f27129d..0e23594 100644
--- a/lib/xixanta/src/opcodes.rs
+++ b/lib/xixanta/src/opcodes.rs
@@ -67,6 +67,11 @@ pub struct Control {
     /// (e.g. .repeat).
     pub has_identifier: Option<bool>,
 
+    /// Whether the control statement only has one mandatory argument which is a
+    /// double-quoted string. If set to true, then `required_args` is ignored in
+    /// favor of this.
+    pub only_string: bool,
+
     /// Minimum and maximum number of arguments accepted by this control
     /// statement, or None if undefined (e.g. a .macro which has an undefined
     /// number of arguments).
@@ -731,27 +736,27 @@ lazy_static! {
     pub static ref CONTROL_FUNCTIONS: HashMap<String, Control> = {
         let mut functions = HashMap::new();
 
-        functions.insert(String::from(".hibyte"), Control { control_type: ControlType::Hibyte, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
-        functions.insert(String::from(".lobyte"), Control { control_type: ControlType::Lobyte, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
-        functions.insert(String::from(".macro"), Control { control_type: ControlType::StartMacro, has_identifier: Some(false), required_args: None, touches_context: true });
-        functions.insert(String::from(".proc"), Control { control_type: ControlType::StartProc, has_identifier: Some(false), required_args: Some((0, 0)), touches_context: true });
-        functions.insert(String::from(".scope"), Control { control_type: ControlType::StartScope, has_identifier: Some(false), required_args: Some((0, 0)), touches_context: true });
-        functions.insert(String::from(".endscope"), Control { control_type: ControlType::EndScope, has_identifier: None, required_args: Some((0, 0)), touches_context: true });
-        functions.insert(String::from(".endproc"), Control { control_type: ControlType::EndProc, has_identifier: None, required_args: Some((0, 0)), touches_context: true });
-        functions.insert(String::from(".endmacro"), Control { control_type: ControlType::EndMacro, has_identifier: None, required_args: Some((0, 0)), touches_context: true });
-        functions.insert(String::from(".segment"), Control { control_type: ControlType::Segment, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
-        functions.insert(String::from(".byte"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".byt"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".db"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".word"), Control { control_type: ControlType::Word, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".dw"), Control { control_type: ControlType::Word, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".addr"), Control { control_type: ControlType::Addr, has_identifier: None, required_args: None, touches_context: false });
-        functions.insert(String::from(".incbin"), Control { control_type: ControlType::IncBin, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
-        functions.insert(String::from(".repeat"), Control { control_type: ControlType::StartRepeat, has_identifier: Some(true), required_args: Some((1, 2)), touches_context: true });
-        functions.insert(String::from(".endrepeat"), Control { control_type: ControlType::EndRepeat, has_identifier: None, required_args: None, touches_context: true });
-        functions.insert(String::from(".include"), Control { control_type: ControlType::IncludeSource, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
-        functions.insert(String::from(".res"), Control { control_type: ControlType::ReserveMemory, has_identifier: None, required_args: Some((1, 2)), touches_context: false });
-        functions.insert(String::from(".asciiz"), Control { control_type: ControlType::Asciiz, has_identifier: None, required_args: Some((1, 1)), touches_context: false });
+        functions.insert(String::from(".hibyte"), Control { control_type: ControlType::Hibyte, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
+        functions.insert(String::from(".lobyte"), Control { control_type: ControlType::Lobyte, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: false });
+        functions.insert(String::from(".macro"), Control { control_type: ControlType::StartMacro, has_identifier: Some(false), required_args: None, touches_context: true, only_string: false });
+        functions.insert(String::from(".proc"), Control { control_type: ControlType::StartProc, has_identifier: Some(false), required_args: Some((0, 0)), touches_context: true, only_string: false });
+        functions.insert(String::from(".scope"), Control { control_type: ControlType::StartScope, has_identifier: Some(false), required_args: Some((0, 0)), touches_context: true, only_string: false });
+        functions.insert(String::from(".endscope"), Control { control_type: ControlType::EndScope, has_identifier: None, required_args: Some((0, 0)), touches_context: true , only_string: false});
+        functions.insert(String::from(".endproc"), Control { control_type: ControlType::EndProc, has_identifier: None, required_args: Some((0, 0)), touches_context: true , only_string: false});
+        functions.insert(String::from(".endmacro"), Control { control_type: ControlType::EndMacro, has_identifier: None, required_args: Some((0, 0)), touches_context: true, only_string: false });
+        functions.insert(String::from(".segment"), Control { control_type: ControlType::Segment, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
+        functions.insert(String::from(".byte"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".byt"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".db"), Control { control_type: ControlType::Byte, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".word"), Control { control_type: ControlType::Word, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".dw"), Control { control_type: ControlType::Word, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".addr"), Control { control_type: ControlType::Addr, has_identifier: None, required_args: None, touches_context: false, only_string: false });
+        functions.insert(String::from(".incbin"), Control { control_type: ControlType::IncBin, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
+        functions.insert(String::from(".repeat"), Control { control_type: ControlType::StartRepeat, has_identifier: Some(true), required_args: Some((1, 2)), touches_context: true, only_string: false });
+        functions.insert(String::from(".endrepeat"), Control { control_type: ControlType::EndRepeat, has_identifier: None, required_args: None, touches_context: true, only_string: false });
+        functions.insert(String::from(".include"), Control { control_type: ControlType::IncludeSource, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
+        functions.insert(String::from(".res"), Control { control_type: ControlType::ReserveMemory, has_identifier: None, required_args: Some((1, 2)), touches_context: false, only_string: false });
+        functions.insert(String::from(".asciiz"), Control { control_type: ControlType::Asciiz, has_identifier: None, required_args: Some((1, 1)), touches_context: false, only_string: true });
 
         functions
     };
diff --git a/lib/xixanta/src/parser.rs b/lib/xixanta/src/parser.rs
index cad0901..f21445d 100644
--- a/lib/xixanta/src/parser.rs
+++ b/lib/xixanta/src/parser.rs
@@ -969,8 +969,6 @@ impl Parser {
         let first = line.chars().next().unwrap_or_default();
         if first == '(' {
             return self.extract_parenthesized_expression(line);
-        } else if first == '"' {
-            return self.parse_quoted_string(line);
         } else if let Some(node_type) = self.get_unary_from_line(line) {
             // Only treat this as a unary operator if the next character is not
             // another unary operator (e.g. disambiguate between '<<' and '<').
@@ -1200,7 +1198,12 @@ impl Parser {
         // by the control function has already been parsed and set in `left`).
         // Then, just parse the arguments and ensure that it matches the amount
         // required by the function.
-        let args = self.parse_arguments(line)?;
+        let args = if control.only_string {
+            self.skip_whitespace(line);
+            vec![self.parse_quoted_string(line)?]
+        } else {
+            self.parse_arguments(line)?
+        };
         if let Some(args_required) = control.required_args {
             if args.len() < args_required.0 || args.len() > args_required.1 {
                 return Err(self.parser_error(
@@ -1584,17 +1587,32 @@ mod tests {
     #[test]
     fn parse_string() {
         let mut parser = Parser::default();
-        let line = ".asciiz \"a: b\"";
+        let line = ".asciiz \"a: b, c\"";
 
         assert!(parser.parse(line.as_bytes(), SourceInfo::default()).is_ok());
 
         let stmt = parser.nodes.last().unwrap().last().unwrap();
         let inner = stmt.args.as_ref().unwrap().first().unwrap();
-        assert_eq!(inner.node_type, NodeType::Literal);
-        assert_eq!(inner.value.value, "\"a: b\"");
+        assert_eq!(inner.node_type, NodeType::Value);
+        assert_eq!(inner.value.value, "\"a: b, c\"");
         assert_eq!(
             line.get(inner.value.start..inner.value.end).unwrap(),
-            "\"a: b\""
+            "\"a: b, c\""
+        );
+    }
+
+    #[test]
+    fn error_on_non_ascii_string() {
+        let mut parser = Parser::default();
+        let line = ".asciiz \"à\"";
+
+        let err = parser
+            .parse(line.as_bytes(), SourceInfo::default())
+            .unwrap_err();
+
+        assert_eq!(
+            err.first().unwrap().message,
+            "using non-ASCII characters in a string"
         );
     }