From dba4a2e2ecfa11a002c710c84f0cf73291bd0e6d Mon Sep 17 00:00:00 2001
From: Nilstrieb <nilstrieb@gmail.com>
Date: Sat, 18 Sep 2021 21:02:06 +0200
Subject: [PATCH] added label support

---
 src/stmt.rs | 95 +++++++++++++++++++++++++++++++++++++++--------------
 test.m8     | 19 +++++------
 2 files changed, 78 insertions(+), 36 deletions(-)

diff --git a/src/stmt.rs b/src/stmt.rs
index 4c02a4c..32e2d3a 100644
--- a/src/stmt.rs
+++ b/src/stmt.rs
@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::fmt::Formatter;
 use std::num::ParseIntError;
 
@@ -10,21 +11,67 @@ pub enum Stmt {
     Stop,
 }
 
+enum IrStmt<'a> {
+    Inc(usize),
+    Dec(usize),
+    IsZero(usize, &'a str),
+    Jump(&'a str),
+    Label(&'a str),
+    Stop,
+}
+
 pub fn parse(text: &str) -> Result<Vec<Stmt>, String> {
-    text.lines()
-        .filter(|line| line.split_whitespace().next().is_some())
-        .map(parse_line)
+    let mut labels = HashMap::new();
+
+    let mut statements = Vec::new();
+    let mut statement_number = 0;
+
+    for (line_number, line) in text.lines().enumerate() {
+        if line.split_whitespace().next().is_none() {
+            continue;
+        }
+        let result = parse_line(line);
+        match result {
+            Ok(IrStmt::Label(name)) => {
+                labels.insert(name, statement_number);
+            }
+            Ok(stmt) => {
+                statement_number += 1;
+                statements.push((stmt, line_number));
+            }
+            Err(msg) => return Err(format!("error on line {}: {}", line_number, msg)),
+        }
+    }
+
+    statements
+        .iter()
+        .map(|stmt| match stmt.0 {
+            IrStmt::Inc(r) => Ok(Stmt::Inc(r)),
+            IrStmt::Dec(r) => Ok(Stmt::Dec(r)),
+            IrStmt::IsZero(r, label) => Ok(Stmt::IsZero(
+                r,
+                match labels.get(label) {
+                    Some(line) => *line,
+                    None => return Err(format!("Label '{}' not found on line {}", label, stmt.1)),
+                },
+            )),
+            IrStmt::Jump(label) => Ok(Stmt::Jump(match labels.get(label) {
+                Some(line) => *line,
+                None => return Err(format!("Label '{}' not found on line {}", label, stmt.1)),
+            })),
+            IrStmt::Stop => Ok(Stmt::Stop),
+            IrStmt::Label(_) => unreachable!(),
+        })
         .collect()
 }
 
-fn parse_line(line: &str) -> Result<Stmt, String> {
-    let no_register = || "No register".to_string();
-    let no_line_number = || "No line number".to_string();
-    let empty_line = || "Empty line not allowed".to_string();
+fn parse_line(line: &str) -> Result<IrStmt, String> {
+    let no_register = || "No register provided".to_string();
+    let no_label = || "No label provided".to_string();
     let display_err = |parse_err: ParseIntError| parse_err.to_string();
 
-    let mut iter = line.split_ascii_whitespace();
-    let first = iter.next().ok_or_else(empty_line)?;
+    let mut iter = line.split_whitespace();
+    let first = iter.next().expect("Empty lines filtered out");
 
     Ok(match first {
         "INC" => {
@@ -33,7 +80,7 @@ fn parse_line(line: &str) -> Result<Stmt, String> {
                 .ok_or_else(no_register)?
                 .parse()
                 .map_err(display_err)?;
-            Stmt::Inc(register)
+            IrStmt::Inc(register)
         }
         "DEC" => {
             let register = iter
@@ -41,7 +88,7 @@ fn parse_line(line: &str) -> Result<Stmt, String> {
                 .ok_or_else(no_register)?
                 .parse()
                 .map_err(display_err)?;
-            Stmt::Dec(register)
+            IrStmt::Dec(register)
         }
         "IS_ZERO" => {
             let register = iter
@@ -49,23 +96,21 @@ fn parse_line(line: &str) -> Result<Stmt, String> {
                 .ok_or_else(no_register)?
                 .parse()
                 .map_err(display_err)?;
-            let line_number = iter
-                .next()
-                .ok_or_else(no_line_number)?
-                .parse()
-                .map_err(display_err)?;
-            Stmt::IsZero(register, line_number)
+            let label = iter.next().ok_or_else(no_label)?;
+            IrStmt::IsZero(register, label)
         }
         "JUMP" => {
-            let line_number = iter
-                .next()
-                .ok_or_else(no_line_number)?
-                .parse()
-                .map_err(display_err)?;
-            Stmt::Jump(line_number)
+            let label = iter.next().ok_or_else(no_label)?;
+            IrStmt::Jump(label)
+        }
+        "STOP" => IrStmt::Stop,
+        stmt => {
+            if stmt.starts_with('.') {
+                IrStmt::Label(&stmt[1..])
+            } else {
+                return Err(format!("Illegal instruction: '{}'", stmt));
+            }
         }
-        "STOP" => Stmt::Stop,
-        stmt => return Err(format!("Illegal instruction: '{}'", stmt)),
     })
 }
 
diff --git a/test.m8 b/test.m8
index 64cc87e..d640005 100644
--- a/test.m8
+++ b/test.m8
@@ -1,11 +1,8 @@
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
-INC 1
+INC 0
+INC 0
+JUMP end
+DEC 0
+DEC 0
+DEC 0
+.end
+STOP
\ No newline at end of file