Handle Pseudo-Instructions, ensure labels are properly setup before assembling

src/env.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use crate::{
     err::RuntimeErr,
-    instructions::{instruction, with, Arg},
+    instructions::{handle_pseudo, instruction, kind::Kind, with, Arg},
     parser::{Loc, Token},
 };
 
@@ -18,8 +18,8 @@ pub enum SymbolValue {
 pub struct Env {
     pub register_alias: HashMap<String, u32>,
     labels: HashMap<String, u32>,
-    registers: [u64; 32],
+    registers: [u32; 32],
-    pub stack: Vec<u64>, // TODO: Find the size of the stack
+    pub stack: Vec<u32>, // TODO: Find the size of the stack
     pub instructions: Vec<u32>,
 }
 
@@ -74,11 +74,11 @@ impl Env {
         }
     }
 
-    pub fn set_register(&mut self, reg: u32, value: u64) {
+    pub fn set_register(&mut self, reg: u32, value: u32) {
         self.registers[reg as usize] = value;
     }
 
-    pub fn get_register(&self, reg: u32) -> u64 {
+    pub fn get_register(&self, reg: u32) -> u32 {
         self.registers[reg as usize]
     }
 
@@ -121,9 +121,17 @@ impl Env {
     pub fn assemble_op(
         &self,
         op: (Token, Loc),
-    ) -> Result<u32, (RuntimeErr, Loc, Option<String>)> {
+    ) -> Result<Vec<u32>, (RuntimeErr, Loc, Option<String>)> {
         if let (Token::Op(name, args), loc) = op {
-            let i = instruction(&name);
+            let i = if let Some(i) = instruction(&name) {
+                i
+            } else {
+                return Err((
+                    RuntimeErr::InvalidOp,
+                    loc,
+                    Some("no implementation exists".to_string()),
+                ));
+            };
             let mut imm = 0u32;
             let mut regs = vec![0; 4];
             if args.len() != i.1.len() {
@@ -145,7 +153,10 @@ impl Env {
                                 Ok(())
                             } else {
                                 Err((
-                                    RuntimeErr::InvalidType("immediate".to_string(), v.kind()),
+                                    RuntimeErr::InvalidType(
+                                        Arg::from(args[k].0.clone()).kind(),
+                                        v.kind(),
+                                    ),
                                     args[k].1,
                                     None,
                                 ))
@@ -157,7 +168,10 @@ impl Env {
                                 Ok(())
                             } else {
                                 Err((
-                                    RuntimeErr::InvalidType("register".to_string(), v.kind()),
+                                    RuntimeErr::InvalidType(
+                                        Arg::from(args[k].0.clone()).kind(),
+                                        v.kind(),
+                                    ),
                                     args[k].1,
                                     None,
                                 ))
@@ -184,17 +198,80 @@ impl Env {
                                 Ok(())
                             } else {
                                 Err((
-                                    RuntimeErr::InvalidType("memory".to_string(), v.kind()),
+                                    RuntimeErr::InvalidType(
+                                        Arg::from(args[k].0.clone()).kind(),
+                                        v.kind(),
+                                    ),
+                                    args[k].1,
+                                    None,
+                                ))
+                            }
+                        }
+                        Arg::Symbol => {
+                            if let Token::Symbol(s) = &args[k].0 {
+                                if let Some(v) = self.get_label(&s) {
+                                    imm = v;
+                                    Ok(())
+                                } else {
+                                    Err((
+                                        RuntimeErr::InvalidType(
+                                            Arg::from(args[k].0.clone()).kind(),
+                                            v.kind(),
+                                        ),
+                                        args[k].1,
+                                        None,
+                                    ))
+                                }
+                            } else {
+                                Err((
+                                    RuntimeErr::InvalidType(
+                                        Arg::from(args[k].0.clone()).kind(),
+                                        v.kind(),
+                                    ),
                                     args[k].1,
                                     None,
                                 ))
                             }
                         }
-                        _ => unimplemented!(),
                     })?;
-            Ok(u32::from_str_radix(dbg!(&with(i, imm, regs).0.to_string()), 2).unwrap())
+            Ok({
+                if let Kind::Pseudo(_) = i.0 {
+                    handle_pseudo(i, imm, regs)
+                        .into_iter()
+                        .map(|x| u32::from_str_radix(&x.0.to_string(), 2).unwrap())
+                        .collect()
+                } else {
+                    vec![u32::from_str_radix(&with(i, imm, regs).0.to_string(), 2).unwrap()]
+                }
+            })
         } else {
             unreachable!()
         }
     }
+
+    pub fn handle_labels(&mut self, tokens: Vec<(Token, Loc)>) {
+        let mut i = 0;
+        // Calculate the instruction position for all opcodes to
+        // allow for labels to be used before they are defined
+        tokens.into_iter().for_each(|(token, _)| match token {
+            Token::Op(name, _) => {
+                if let Some((kind, args)) = instruction(&name) {
+                    if let Kind::Pseudo(_) = kind {
+                        handle_pseudo((kind, args), 0, vec![0; 4])
+                            .into_iter()
+                            .for_each(|_| i += 1);
+                    } else {
+                        i += 1;
+                    }
+                }
+            }
+            Token::Label(name) => {
+                self.add_label(&name, i * 4);
+            }
+            other => {
+                dbg!(other);
+                unreachable!()
+            }
+        });
+    }
 }

src/err.rs

@@ -54,7 +54,7 @@ pub enum RuntimeErr {
 impl Display for RuntimeErr {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         match self {
-            RuntimeErr::InvalidOp => write!(f, "invalid operation"),
+            RuntimeErr::InvalidOp => write!(f, "invalid opcode"),
             RuntimeErr::InvalidOpArity(op, actual, expected) => {
                 write!(f, "`{}` expected {} args, got {}", op, expected, actual)
             }
@@ -70,7 +70,7 @@ impl RuntimeErr {
         match self {
             RuntimeErr::InvalidOp => "check the ref sheet for the avaliable opcodes".to_string(),
             RuntimeErr::InvalidOpArity(op, actual, expected) => {
-                let args = instruction(op).1;
+                let args = instruction(op).unwrap().1;
                 match actual.cmp(expected) {
                     Ordering::Equal => unreachable!(),
                     Ordering::Greater if actual - expected == 1 => {

src/instructions.rs

@@ -4,7 +4,7 @@ pub mod kind {
     use crate::instructions::to_u32;
 
     /// will be converted by the engine to a real instruction
-    pub struct Pseudo {}
+    pub struct Pseudo(pub &'static str);
 
     pub struct R {
         pub funct7: [bool; 7],
@@ -110,7 +110,6 @@ pub mod kind {
 
     impl Display for Pseudo {
         fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-            // (pseudo) padded on either side with - to make it 32 characters
             write!(f, "{:0^32}", 0)
         }
     }
@@ -229,6 +228,7 @@ pub enum Arg {
     Immediate,
     /// always ra
     Memory,
+    // Apperently a symbol is a label but in respect to the current pc
     Symbol,
 }
 
@@ -244,17 +244,31 @@ impl Arg {
     }
 }
 
+impl From<Token> for Arg {
+    fn from(token: Token) -> Self {
+        match token {
+            Token::Immediate(_) => Arg::Immediate,
+            Token::Register(_) => Arg::Register(0),
+            Token::Memory(_, _) => Arg::Memory,
+            Token::Symbol(_) => Arg::Symbol,
+            _ => unreachable!(),
+        }
+    }
+}
+
 use kind::*;
 
+use crate::parser::Token;
+
 /// (kind, (arity, Vec<token kind>))
-pub fn instruction(op: &str) -> (Kind, Vec<Arg>) {
+pub fn instruction(op: &str) -> Option<(Kind, Vec<Arg>)> {
-    match op {
+    Some(match op {
         // -
-        "nop" => (Kind::Pseudo(Pseudo {}), vec![]),
+        "nop" => (Kind::Pseudo(Pseudo("nop")), vec![]),
 
         // Move
         "li" => (
-            Kind::Pseudo(Pseudo {}),
+            Kind::Pseudo(Pseudo("li")),
             vec![Arg::Register(0), Arg::Immediate],
         ),
         "lui" => (
@@ -301,7 +315,30 @@ pub fn instruction(op: &str) -> (Kind, Vec<Arg>) {
             }),
             vec![Arg::Register(0), Arg::Register(1), Arg::Immediate],
         ),
+
         // Multiply, Divide
+        "mul" => (
+            Kind::R(R {
+                funct7: to_bits(0b0000001),
+                rb: to_bits(0),
+                ra: to_bits(0),
+                funct3: to_bits(0b000),
+                rd: to_bits(0),
+                opcode: to_bits(0b0110011),
+            }),
+            vec![Arg::Register(0), Arg::Register(1), Arg::Register(2)],
+        ),
+        "div" => (
+            Kind::R(R {
+                funct7: to_bits(0b0000001),
+                rb: to_bits(0),
+                ra: to_bits(0),
+                funct3: to_bits(0b100),
+                rd: to_bits(0),
+                opcode: to_bits(0b0110011),
+            }),
+            vec![Arg::Register(0), Arg::Register(1), Arg::Register(2)],
+        ),
 
         // Compare
 
@@ -317,8 +354,31 @@ pub fn instruction(op: &str) -> (Kind, Vec<Arg>) {
             }),
             vec![Arg::Register(1), Arg::Register(2), Arg::Immediate],
         ),
+        "bne" => (
+            Kind::B(B {
+                imm: to_bits(0),
+                rb: to_bits(0),
+                ra: to_bits(0),
+                funct3: to_bits(0b001),
+                imm2: to_bits(0),
+                opcode: to_bits(0b1100011),
+            }),
+            vec![Arg::Register(1), Arg::Register(2), Arg::Immediate],
+        ),
+        "beqz" => (
+            Kind::Pseudo(Pseudo("beqz")),
+            vec![Arg::Register(1), Arg::Symbol],
+        ),
+        "bnez" => (
+            Kind::Pseudo(Pseudo("bnez")),
+            vec![Arg::Register(1), Arg::Symbol],
+        ),
         op => unimplemented!("{}", op),
+    })
 }
+
+pub fn get_instruction(op: &str) -> (Kind, Vec<Arg>) {
+    unsafe { instruction(op).unwrap_unchecked() }
 }
 
 /// regs order: rd, ra, rb, rc
@@ -447,6 +507,55 @@ pub fn with((kind, args): (Kind, Vec<Arg>), imm: u32, regs: Vec<u32>) -> (Kind,
     }
 }
 
+/// regs order: rd, ra, rb, rc
+pub fn handle_pseudo(
+    (kind, args): (Kind, Vec<Arg>),
+    imm: u32,
+    regs: Vec<u32>,
+) -> Vec<(Kind, Vec<Arg>)> {
+    let op = if let Kind::Pseudo(Pseudo(op)) = kind {
+        op
+    } else {
+        return vec![(kind, args)];
+    };
+
+    match op {
+        "nop" => vec![
+            // addi x0, x0, 0
+            with(get_instruction("addi"), 0, vec![0, 0]),
+        ],
+        "li" => {
+            match imm {
+                // if the immediate is small enough (12 bits), use addi
+                _ if imm >> 12 == 0 => {
+                    vec![with(get_instruction("addi"), imm, regs)]
+                }
+                // if the immediate is a multiple of 0x1000, use lui
+                _ if imm & 0xfff == 0 => {
+                    vec![with(get_instruction("lui"), imm, regs)]
+                }
+                // otherwise, use lui and addi
+                _ => vec![
+                    with(get_instruction("lui"), imm >> 12, regs.clone()),
+                    with(get_instruction("addi"), imm & 0xfff, regs),
+                ],
+            }
+        }
+        "beqz" => vec![
+            // beq ra, x0, imm
+            with(get_instruction("beq"), imm, vec![0, regs[0], 0]),
+        ],
+        "bnez" => vec![
+            // bne ra, x0, imm
+            with(get_instruction("bne"), imm, vec![0, regs[0], 0]),
+        ],
+        other => {
+            dbg!(other);
+            unimplemented!()
+        }
+    }
+}
+
 fn to_bits<const N: usize>(val: u32) -> [bool; N] {
     let mut bits = [false; N];
     for i in 0..N {

src/main.rs

@@ -19,7 +19,7 @@ fn main() -> anyhow::Result<()> {
 
     let file = SimpleFile::new("test.s", input.clone());
 
-    let env = Env::new();
+    let mut env = Env::new();
 
     match parse(&env, &input) {
         Ok(tokens) => {
@@ -28,18 +28,29 @@ fn main() -> anyhow::Result<()> {
 
             tokens.iter().enumerate().for_each(|(line, token)| {
                 let token = token.clone();
-                let mut formatted = format!("{:<1$}", lines[line].to_string() + ":", size + 3);
+                env.handle_labels(tokens.clone());
 
                 match token.0 {
                     Token::Op(..) => match env.assemble_op(token) {
                         Ok(op) => {
-                            formatted += &format!("{:032b}", op);
+                            let mut formatted = format!(
+                                "{:<1$} {2:032b}",
+                                lines[line].to_string() + ":",
+                                size + 3,
+                                op[0]
+                            );
+
+                            if op.len() > 1 {
+                                for op in op[1..].iter() {
+                                    formatted += &format!("{:<1$} {2:032b}", "", size + 3, op);
+                                }
+                            }
                             println!("{}", formatted);
                         }
                         Err(err) => {
                             let diagnostic = Diagnostic::error()
                                 .with_message("Runtime Error")
-                                .with_labels(vec![Label::primary((), err.1.start..err.1.end)
+                                .with_labels(vec![Label::primary((), err.1.start..(err.1.end + 1))
                                     .with_message(err.0.to_string())])
                                 .with_notes({
                                     let mut notes = Vec::new();
@@ -53,7 +64,7 @@ fn main() -> anyhow::Result<()> {
                             term::emit(&mut writer.lock(), &config, &file, &diagnostic).unwrap();
                         }
                     },
-                    Token::LabelDef(name) => {
+                    Token::Label(name) => {
                         println!("{name}:");
                     }
                     _ => unreachable!(),

src/parser.rs

@@ -4,10 +4,8 @@ use itertools::Itertools;
 
 #[derive(Debug, Clone)]
 pub enum Token {
-    /// ' ', '\t', '\r'
+    /// ' ', '\t', '\r', \# blablabla
     Spacing,
-    /// \# blablabla,
-    Comment,
     /// 1, 2, -1
     Immediate(u32),
     /// zero, r1, pc
@@ -16,8 +14,8 @@ pub enum Token {
     Register(String),
     /// add, xor, j
     Op(String, Vec<(Token, Loc)>),
-    /// <label>:
+    /// \<label>:
-    LabelDef(String),
+    Label(String),
     /// 0(a0)
     Memory(Box<Token>, Option<Box<Token>>),
     /// symbol
@@ -34,11 +32,10 @@ impl Token {
         use Token::*;
         match self {
             Spacing => "spacing",
-            Comment => "comment",
             Immediate(_) => "immediate",
             Register(_) => "register",
             Op(_, _) => "op",
-            LabelDef(_) => "label def",
+            Label(_) => "label",
             Memory(_, _) => "memory",
             Symbol(_) => "symbol",
             String(_) => "string",
@@ -71,7 +68,7 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                     chars.next();
                     loc.end += 1;
                 }
-                Comment
+                Spacing
             }
 
             '0'..='9' => {
@@ -83,7 +80,7 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                 if let Some('(') | Some(' ') | None = chars.peek() {
                     Immediate(num.parse().unwrap())
                 } else {
-                    return Err((
+                    let err = Err((
                         SyntaxErr::UnexpectedChar,
                         Loc {
                             line: loc.line,
@@ -93,6 +90,8 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                         tokens.clone(),
                         None,
                     ));
+                    advance_to_next_line(&mut chars, loc);
+                    return err;
                 }
             }
             '-' => {
@@ -110,12 +109,14 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                     imm = Box::new(tokens.pop().unwrap());
                     loc.start = imm.1.start;
                 } else {
-                    return Err((
+                    let err = Err((
                         SyntaxErr::UnexpectedChar,
                         loc.clone(),
                         tokens.clone(),
                         Some("a memory index must be of the form imm(reg) or imm".to_string()),
                     ));
+                    advance_to_next_line(&mut chars, loc);
+                    return err;
                 }
 
                 let mut reg = std::string::String::new();
@@ -129,7 +130,7 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
 
                 let reg = reg.trim();
                 if !env.is_valid_register(reg) {
-                    return Err((
+                    let err = Err((
                         SyntaxErr::MemoryInvalidRegister,
                         Loc {
                             line: loc.line,
@@ -139,14 +140,18 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                         tokens.clone(),
                         None,
                     ));
+                    advance_to_next_line(&mut chars, loc);
+                    return err;
                 }
                 if chars.next() != Some(')') {
-                    return Err((
+                    let err = Err((
                         SyntaxErr::UnmatchedParen(false),
                         loc.clone(),
                         tokens.clone(),
                         None,
                     ));
+                    advance_to_next_line(&mut chars, loc);
+                    return err;
                 }
                 loc.end += 2;
 
@@ -156,12 +161,14 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                 )
             }
             ')' => {
-                return Err((
+                let err = Err((
                     SyntaxErr::UnmatchedParen(true),
                     loc.clone(),
                     tokens.clone(),
                     None,
-                ))
+                ));
+                advance_to_next_line(&mut chars, loc);
+                return err;
             }
 
             // Opcode or Label definition
@@ -176,14 +183,18 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                 if let Some(':') = chars.peek() {
                     chars.next();
                     loc.end += 1;
-                    LabelDef(str[..str.len()].to_string())
+                    Label(str[..str.len()].to_string())
                 } else {
                     // These Registers may actually be ops, label references or symbols, but there's ambiguity
                     // between them and registers, so we'll just assume they're registers for now
                     Register(str.trim().to_owned())
                 }
             }
-            _ => return Err((SyntaxErr::UnexpectedChar, loc.clone(), tokens.clone(), None)),
+            _ => {
+                let err = Err((SyntaxErr::UnexpectedChar, loc.clone(), tokens.clone(), None));
+                advance_to_next_line(&mut chars, loc);
+                return err;
+            }
         };
         tokens.push((token, loc.clone()));
         loc.end += 1;
@@ -232,8 +243,18 @@ fn parse_line(env: &Env, input: &str, loc: &mut Loc) -> Result<Vec<(Token, Loc)>
                         loc.clone(),
                     )];
                 }
-
+                for (token, loc) in group[1..].iter() {
-                args.extend_from_slice(&group[1..]);
+                    match token.clone() {
+                        Token::Register(name) => {
+                            if env.is_valid_register(&name) {
+                                args.push((token.clone(), loc.clone()));
+                            } else {
+                                args.push((Token::Symbol(name.to_owned()), *loc))
+                            }
+                        }
+                        others => args.push((others, *loc)),
+                    }
+                }
 
                 vec![(Op(name, args), loc)]
             } else {
@@ -281,3 +302,12 @@ pub fn parse(env: &Env, input: &str) -> Result<Vec<(Token, Loc)>, Vec<ParseErr>>
         Err(err.into_iter().map(|line| line.unwrap_err()).collect())
     }
 }
+
+fn advance_to_next_line(chars: &mut std::iter::Peekable<std::str::Chars>, loc: &mut Loc) {
+    while let Some(_) = chars.peek() {
+        chars.next();
+        loc.end += 1;
+    }
+    loc.end += 1; // Newline
+    loc.start = loc.end;
+}

src/tests.rs

@@ -2,7 +2,7 @@
 /// Test values come from Ripes
 use crate::{
     env::Env,
-    instructions::{instruction, with},
+    instructions::{get_instruction, with},
 };
 
 #[test]
@@ -19,7 +19,7 @@ fn nop() {
     assert_eq!(
         u32::from_str_radix(
             &with(
-                instruction("nop"),
+                get_instruction("nop"),
                 0, // imm
                 vec![]
             )
@@ -46,7 +46,7 @@ fn sb() {
     assert_eq!(
         u32::from_str_radix(
             &with(
-                instruction("sb"),
+                get_instruction("sb"),
                 -4i32 as u32, // imm
                 vec![
                     0,                                    // rd
@@ -77,7 +77,7 @@ fn add() {
     assert_eq!(
         u32::from_str_radix(
             &with(
-                instruction("add"),
+                get_instruction("add"),
                 0, // imm
                 vec![
                     env.alias_to_register("a0").unwrap(), // rd
@@ -109,7 +109,7 @@ fn addi() {
     assert_eq!(
         u32::from_str_radix(
             with(
-                instruction("addi"),
+                get_instruction("addi"),
                 1,
                 vec![
                     env.alias_to_register("a0").unwrap(),
@@ -141,7 +141,7 @@ fn beq() {
     assert_eq!(
         u32::from_str_radix(
             &with(
-                instruction("beq"),
+                get_instruction("beq"),
                 4,
                 vec![
                     0, // no rd

test.s

@@ -2,8 +2,9 @@
     li a1 1
 
 factorial:
-    beqz a0 end
+    # beqz a0 end
-    mul a1 a1 a0
+    mul a1 a1 a0 -1
-    addi a0 a0 -1
+    a0 addi a0 a0 -1
-    bneqz factorial
+
+    bnez factorial
 end: