resolve #11. Switch statement implemented!

plugins/b-mode.el

@@ -2,7 +2,7 @@
 
 (define-generic-mode 'b-mode
   '(("/*" . "*/"))
-  '("return" "auto" "extrn" "eof" "while" "if" "else" "goto" "switch" "break")
+  '("return" "auto" "extrn" "eof" "while" "if" "else" "goto" "switch" "break" "case" "default")
   '(("\\b[0-9]+\\b" . font-lock-constant-face))
   '("\\.b$")
   '() ;; TODO: Highlight string & char literals here

src/ast.rs

@@ -92,6 +92,16 @@ pub enum Statement {
     Block(Vec<Statement>),
     If(Expr, Box<Statement>, Option<Box<Statement>>),
     While(Expr, Box<Statement>),
+    Switch(Expr, Vec<SwInner>),
+}
+
+/// Statements allowed inside switch statements
+/// Switch needs special rules for "case" and "default"
+#[derive(Debug)]
+pub enum SwInner {
+    Default(Pos),
+    Case(Pos, i64),
+    Statement(Statement),
 }
 
 #[derive(Debug)]

src/codegen.rs

@@ -1,5 +1,6 @@
 use std::sync::Condvar;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::io::BufWriter;
 use std::io::Write;
 use std::sync::{Arc, Mutex};
@@ -763,15 +764,15 @@ fn gen_expr_deref_ass(
     Ok((Loc::Register(rhs_reg), used_registers))
 }
 
+/// # Arguments
+/// * `dest_reg` - For values > 2^32, this is the reg that will be used.
 fn gen_int(
-    instructions: &mut Vec<String>, signed: i64
+    instructions: &mut Vec<String>, signed: i64, dest_reg: Reg
 ) -> (Loc, RegSet) {
     if signed < i32::MAX as i64 && signed >= i32::MIN as i64 {
         (Loc::Immediate(signed), RegSet::empty())
     } else {
         let unsigned = signed as u64;
-        let dest_reg = Reg::Rax;
-
         let upper = (unsigned & !((1 << 32) - 1)) >> 32;
         let lower = unsigned & ((1 << 32) - 1);
 
@@ -825,7 +826,7 @@ fn gen_expr(
     expr: &Expr
 ) -> Result<(Loc, RegSet), CompErr> {
     match expr {
-        Expr::Int(_, value) => Ok(gen_int(instructions, *value)),
+        Expr::Int(_, value) => Ok(gen_int(instructions, *value, Reg::Rax)),
         Expr::Id(pos, name) => {
             match c.find_in_scope(name) {
                 Some(ScopeEntry::Var(loc)) =>
@@ -977,6 +978,73 @@ fn gen_if_else(
     Ok(())
 }
 
+fn gen_switch(
+    c: &mut FunContext, instructions: &mut Vec<String>,
+    cond: &Expr, body: &Vec<SwInner>
+) -> Result<(), CompErr> {
+    let (expr_loc, _) = gen_expr(c, instructions, cond)?;
+    // cmp requires the dest to be in a register
+    let cond_loc = match expr_loc {
+        loc @ Loc::Register(_) => loc,
+        other => {
+            instructions.push(format!("movq {},%rax", other));
+            Loc::Register(Reg::Rax)
+        },
+    };
+
+    // The register to store case values
+    let case_reg = match cond_loc {
+        Loc::Register(Reg::Rax) => Reg::Rcx,
+        _                       => Reg::Rax,
+    };
+
+    let mut used_case_values = HashSet::new();
+    let mut default_label: Option<String> = None;
+    let mut body_inst = vec!();
+
+    let switch_end_label = c.new_label("SW_END");
+
+    c.break_dest_stack.push(switch_end_label.clone());
+    for inner in body {
+        match inner {
+            SwInner::Default(pos) => {
+                if !default_label.is_none() {
+                    return CompErr::err(
+                        pos,
+                        format!("`default` label is already defined in switch"));
+                }
+                let label_name = c.new_label("SW_DEFAULT");
+                body_inst.push(format!("{}:", label_name));
+                default_label = Some(label_name);
+            },
+            SwInner::Case(pos, value) => {
+                if used_case_values.contains(value) {
+                    return CompErr::err(
+                        pos,
+                        format!("case {} is already defined in switch", value));
+                }
+                used_case_values.insert(value);
+
+                let (case_loc, _) = gen_int(&mut body_inst, *value, case_reg);
+                let label_name = c.new_label("SW_CASE");
+                body_inst.push(format!("{}:", label_name));
+                instructions.push(format!("cmpq {},{}", case_loc, cond_loc));
+                instructions.push(format!("je {}", label_name));
+            },
+            SwInner::Statement(body) => gen_statement(c, &mut body_inst, body)?,
+        }
+    }
+    c.break_dest_stack.pop();
+    // Default jump point
+    instructions.push(format!("jmp {}", match default_label {
+        Some(label_name) => label_name,
+        None             => switch_end_label.clone(),
+    }));
+    instructions.append(&mut body_inst);
+    instructions.push(format!("{}:", switch_end_label));
+    Ok(())
+}
+
 fn gen_auto(
     c: &mut FunContext, instructions: &mut Vec<String>,
     pos: &Pos, vars: &Vec<Var>
@@ -997,13 +1065,13 @@ fn gen_auto(
                     let value = initial[i];
                     let val_dest_loc = Loc::Stack(-size - 1 + i as i64);
 
-                    let (val_loc, _) = gen_int(instructions, value);
+                    let (val_loc, _) = gen_int(instructions, value, Reg::Rax);
                     instructions.push(format!(
                         "movq {},{}", val_loc, val_dest_loc));
                 }
             },
             Var::Single(_, Some(value)) => {
-                let (val_loc, _) = gen_int(instructions, *value);
+                let (val_loc, _) = gen_int(instructions, *value, Reg::Rax);
                 instructions.push(format!("movq {},{}", val_loc, dest_loc));
             },
             Var::Single(_, None) => {},
@@ -1091,6 +1159,7 @@ fn gen_statement(
         Statement::While(cond, body)       => gen_while(c, instructions, cond, body),
         Statement::If(cond, if_body, Some(else_body)) =>
             gen_if_else(c, instructions, cond, if_body, else_body),
+        Statement::Switch(cond, body) => gen_switch(c, instructions, cond, body),
     }
 }
 

src/parser.rs

@@ -275,6 +275,7 @@ fn parse_statement(c: &mut ParseContext) -> Result<Statement, CompErr> {
         Token::Extern      => parse_statement_extern(c, pos),
         Token::If          => parse_statement_if(c),
         Token::While       => parse_statement_while(c),
+        Token::Switch      => parse_statement_switch(c),
         Token::Semicolon   => Ok(Statement::Null),
         Token::Label(name) => Ok(Statement::Label(pos, name)),
         Token::Goto        => parse_statement_goto(c),
@@ -389,6 +390,42 @@ fn parse_statement_goto(c: &mut ParseContext) -> Result<Statement, CompErr> {
     }
 }
 
+// Expect "switch" to have been parsed already
+fn parse_statement_switch(c: &mut ParseContext) -> Result<Statement, CompErr> {
+    parse_tok(c, Token::LParen)?;
+    let cond_expr = parse_expr(c)?;
+    parse_tok(c, Token::RParen)?;
+    // As per the B spec, switch bodies must always be compound statements
+    parse_tok(c, Token::LBrace)?;
+
+    let mut inner_statements = vec!();
+    loop {
+        let inner = match pop_tok(c)? {
+            (pos, Token::Default) => {
+                parse_tok(c, Token::Colon)?;
+                SwInner::Default(pos)
+            },
+            (pos, Token::Case) => {
+                let value = match pop_tok(c)? {
+                    (_, Token::Int(value))  => value,
+                    (_, Token::Char(chars)) => pack_chars(&chars)[0],
+                    (pos, tok) => return CompErr::err(
+                        &pos, format!("Int or char expected, {:?} given.", tok)),
+                };
+                parse_tok(c, Token::Colon)?;
+                SwInner::Case(pos, value)
+            },
+            (_, Token::RBrace) => break,
+            other => {
+                push_tok(c, other);
+                SwInner::Statement(parse_statement(c)?)
+            },
+        };
+        inner_statements.push(inner);
+    }
+    Ok(Statement::Switch(cond_expr, inner_statements))
+}
+
 // Expect "while" to have been parsed already
 fn parse_statement_while(c: &mut ParseContext) -> Result<Statement, CompErr> {
     parse_tok(c, Token::LParen)?;

src/tokenizer.rs

@@ -15,6 +15,8 @@ pub enum Token {
     Import,
     Return,
     Auto,
+    Case,
+    Default,
     Extern,
     Eof,
     While,
@@ -370,15 +372,17 @@ fn get_tok_word(c: &mut ParseContext) -> Result<(Pos, Token), CompErr> {
 
     // Safe to assume it's valid utf8 since we enforce ASCII
     let tok = match slice {
-        "auto"   => Token::Auto,
-        "break"  => Token::Break,
-        "else"   => Token::Else,
-        "extrn"  => Token::Extern,
-        "goto"   => Token::Goto,
-        "if"     => Token::If,
-        "return" => Token::Return,
-        "switch" => Token::Switch,
-        "while"  => Token::While,
+        "auto"    => Token::Auto,
+        "break"   => Token::Break,
+        "case"    => Token::Case,
+        "default" => Token::Default,
+        "else"    => Token::Else,
+        "extrn"   => Token::Extern,
+        "goto"    => Token::Goto,
+        "if"      => Token::If,
+        "return"  => Token::Return,
+        "switch"  => Token::Switch,
+        "while"   => Token::While,
         word => {
             let name: String = word.to_string();
 

test/hello.b

@@ -1,17 +1,5 @@
 #import "../assets/stdlib.b";
-#import "../assets/memory.b";
 
 main() {
-    auto c 10, i 3;
-    putnum(i << 3);
-    putstr("*n");
-    putnum(i * 8);
-    putstr("*n");
     putstr("Hello world!*n");
-
-    putstr(num2str(-321));
-    putstr("*n");
-
-    putnum(0 ? c : 0 ? 42 : 1337);
-    putstr("*n");
 }

test/switch.b

@@ -0,0 +1,14 @@
+#import "../assets/stdlib.b";
+
+main() {
+    switch ('b' - 1) {
+    case 1:
+        putstr("One*n");
+        break;
+    case 'a':
+        putstr("a is for apple*n");
+        break;
+    default:
+        putstr("Something else*n");
+    }
+}