Implement jump tables.

- Add a ir::jumptable module with a JumpTableData struct representing the vector
  of destinations.
- Add an entity map of jump tables to the Function.
- Parse and write jump tables in the function preamble.
- Rewrite EBB references in jumptables after parsing.

src/libcretonne/ir/jumptable.rs

@@ -0,0 +1,156 @@
+//! Jump table representation.
+//!
+//! Jump tables are declared in the preamble and assigned an `ir::entities::JumpTable` reference.
+//! The actual table of destinations is stored in a `JumpTableData` struct defined in this module.
+
+use ir::entities::{Ebb, NO_EBB};
+use std::iter;
+use std::slice;
+use std::fmt::{self, Display, Formatter};
+
+/// Contents of a jump table.
+///
+/// All jump tables use 0-based indexing and are expected to be densely populated. They don't need
+/// to be completely populated, though. Individual entries can be missing.
+pub struct JumpTableData {
+    // Table entries, using NO_EBB as a placeholder for missing entries.
+    table: Vec<Ebb>,
+
+    // How many `NO_EBB` holes in table?
+    holes: usize,
+}
+
+impl JumpTableData {
+    /// Create a new empty jump table.
+    pub fn new() -> JumpTableData {
+        JumpTableData {
+            table: Vec::new(),
+            holes: 0,
+        }
+    }
+
+    /// Set a table entry.
+    ///
+    /// The table will grow as needed to fit 'idx'.
+    pub fn set_entry(&mut self, idx: usize, dest: Ebb) {
+        assert!(dest != NO_EBB);
+        // Resize table to fit `idx`.
+        if idx >= self.table.len() {
+            self.holes += idx - self.table.len();
+            self.table.resize(idx + 1, NO_EBB);
+        } else if self.table[idx] == NO_EBB {
+            // We're filling in an existing hole.
+            self.holes -= 1;
+        }
+        self.table[idx] = dest;
+    }
+
+    /// Clear a table entry.
+    ///
+    /// The `br_table` instruction will fall through if given an index corresponding to a cleared
+    /// table entry.
+    pub fn clear_entry(&mut self, idx: usize) {
+        if idx < self.table.len() && self.table[idx] != NO_EBB {
+            self.holes += 1;
+            self.table[idx] = NO_EBB;
+        }
+    }
+
+    /// Get the entry for `idx`, or `None`.
+    pub fn get_entry(&self, idx: usize) -> Option<Ebb> {
+        if idx < self.table.len() && self.table[idx] != NO_EBB {
+            Some(self.table[idx])
+        } else {
+            None
+        }
+    }
+
+    /// Enumerate over all `(idx, dest)` pairs in the table in order.
+    ///
+    /// This returns an iterator that skips any empty slots in the table.
+    pub fn entries<'a>(&'a self) -> Entries {
+        Entries(self.table.iter().cloned().enumerate())
+    }
+
+    /// Access the whole table as a mutable slice.
+    pub fn as_mut_slice(&mut self) -> &mut [Ebb] {
+        self.table.as_mut_slice()
+    }
+}
+
+/// Enumerate `(idx, dest)` pairs in order.
+pub struct Entries<'a>(iter::Enumerate<iter::Cloned<slice::Iter<'a, Ebb>>>);
+
+impl<'a> Iterator for Entries<'a> {
+    type Item = (usize, Ebb);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some((idx, dest)) = self.0.next() {
+                if dest != NO_EBB {
+                    return Some((idx, dest));
+                }
+            } else {
+                return None;
+            }
+        }
+    }
+}
+
+impl Display for JumpTableData {
+    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
+        let first = self.table.first().cloned().unwrap_or_default();
+        if first == NO_EBB {
+            try!(write!(fmt, "jump_table 0"));
+        } else {
+            try!(write!(fmt, "jump_table {}", first));
+        }
+
+        for dest in self.table.iter().cloned().skip(1) {
+            if dest == NO_EBB {
+                try!(write!(fmt, ", 0"));
+            } else {
+                try!(write!(fmt, ", {}", dest));
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::JumpTableData;
+    use ir::entities::Ebb;
+    use entity_map::EntityRef;
+
+    #[test]
+    fn empty() {
+        let jt = JumpTableData::new();
+
+        assert_eq!(jt.get_entry(0), None);
+        assert_eq!(jt.get_entry(10), None);
+
+        assert_eq!(jt.to_string(), "jump_table 0");
+
+        let v: Vec<(usize, Ebb)> = jt.entries().collect();
+        assert_eq!(v, []);
+    }
+
+    #[test]
+    fn insert() {
+        let e1 = Ebb::new(1);
+        let e2 = Ebb::new(2);
+
+        let mut jt = JumpTableData::new();
+
+        jt.set_entry(0, e1);
+        jt.set_entry(0, e2);
+        jt.set_entry(10, e1);
+
+        assert_eq!(jt.to_string(),
+                   "jump_table ebb2, 0, 0, 0, 0, 0, 0, 0, 0, 0, ebb1");
+
+        let v: Vec<(usize, Ebb)> = jt.entries().collect();
+        assert_eq!(v, [(0, e2), (10, e1)]);
+    }
+}

src/libcretonne/ir/mod.rs

@@ -5,12 +5,14 @@ pub mod types;
 pub mod condcodes;
 pub mod immediates;
 pub mod instructions;
+pub mod jumptable;
 pub mod dfg;
 pub mod layout;
 
 use ir::types::{FunctionName, Signature};
-use entity_map::EntityRef;
+use entity_map::{EntityRef, EntityMap};
-use ir::entities::StackSlot;
+use ir::entities::{StackSlot, JumpTable};
+use ir::jumptable::JumpTableData;
 use ir::dfg::DataFlowGraph;
 use ir::layout::Layout;
 use std::fmt::{self, Debug, Display, Formatter};
@@ -27,6 +29,9 @@ pub struct Function {
     /// Stack slots allocated in this function.
     stack_slots: Vec<StackSlotData>,
 
+    /// Jump tables used in this function.
+    pub jump_tables: EntityMap<JumpTable, JumpTableData>,
+
     /// Data flow graph containing the primary definition of all instructions, EBBs and values.
     pub dfg: DataFlowGraph,
 
@@ -41,6 +46,7 @@ impl Function {
             name: name,
             signature: sig,
             stack_slots: Vec::new(),
+            jump_tables: EntityMap::new(),
             dfg: DataFlowGraph::new(),
             layout: Layout::new(),
         }

src/libcretonne/write.rs

@@ -75,6 +75,11 @@ fn write_preamble(w: &mut Write, func: &Function) -> io::Result<bool> {
         try!(writeln!(w, "    {} = {}", ss, func[ss]));
     }
 
+    for jt in func.jump_tables.keys() {
+        any = true;
+        try!(writeln!(w, "    {} = {}", jt, func.jump_tables[jt]));
+    }
+
     Ok(any)
 }
 

src/libreader/lexer.rs

@@ -38,6 +38,7 @@ pub enum Token<'a> {
     Value(Value), // v12, vx7
     Ebb(Ebb), // ebb3
     StackSlot(u32), // ss3
+    JumpTable(u32), // jt2
     Identifier(&'a str), // Unrecognized identifier (opcode, enumerator, ...)
 }
 
@@ -291,6 +292,7 @@ impl<'a> Lexer<'a> {
             "vx" => Value::table_with_number(value).map(|v| Token::Value(v)),
             "ebb" => Ebb::with_number(value).map(|ebb| Token::Ebb(ebb)),
             "ss" => Some(Token::StackSlot(value)),
+            "jt" => Some(Token::JumpTable(value)),
             _ => None,
         }
     }

src/libreader/parser.rs

@@ -17,6 +17,7 @@ use cretonne::ir::entities::*;
 use cretonne::ir::instructions::{Opcode, InstructionFormat, InstructionData, VariableArgs,
                                  JumpData, BranchData, ReturnData};
 use cretonne::ir::{Function, StackSlotData};
+use cretonne::ir::jumptable::JumpTableData;
 
 pub use lexer::Location;
 
@@ -71,6 +72,7 @@ pub struct Parser<'a> {
 struct Context {
     function: Function,
     stack_slots: HashMap<u32, StackSlot>, // ssNN
+    jump_tables: HashMap<u32, JumpTable>, // jtNN
     ebbs: HashMap<Ebb, Ebb>, // ebbNN
     values: HashMap<Value, Value>, // vNN, vxNN
 
@@ -83,6 +85,7 @@ impl Context {
         Context {
             function: f,
             stack_slots: HashMap::new(),
+            jump_tables: HashMap::new(),
             ebbs: HashMap::new(),
             values: HashMap::new(),
             inst_locs: Vec::new(),
@@ -98,6 +101,15 @@ impl Context {
         }
     }
 
+    // Allocate a new jump table and add a mapping number -> JumpTable.
+    fn add_jt(&mut self, number: u32, data: JumpTableData, loc: &Location) -> Result<()> {
+        if self.jump_tables.insert(number, self.function.jump_tables.push(data)).is_some() {
+            err!(loc, "duplicate jump table: jt{}", number)
+        } else {
+            Ok(())
+        }
+    }
+
     // Allocate a new EBB and add a mapping src_ebb -> Ebb.
     fn add_ebb(&mut self, src_ebb: Ebb, loc: &Location) -> Result<Ebb> {
         let ebb = self.function.dfg.make_ebb();
@@ -211,8 +223,15 @@ impl Context {
             }
         }
 
-        // TODO: Rewrite EBB references in jump tables. (Once jump table data structures are
+        // Rewrite EBB references in jump tables.
-        // defined).
+        let loc = Location { line_number: 0 };
+        for jt in self.function.jump_tables.keys() {
+            for ebb in self.function.jump_tables[jt].as_mut_slice() {
+                if *ebb != NO_EBB {
+                    try!(Self::rewrite_ebb(&self.ebbs, ebb, &loc));
+                }
+            }
+        }
 
         Ok(())
     }
@@ -312,6 +331,16 @@ impl<'a> Parser<'a> {
         }
     }
 
+    // Match and consume a jump table reference.
+    fn match_jt(&mut self, err_msg: &str) -> Result<u32> {
+        if let Some(Token::JumpTable(jt)) = self.token() {
+            self.consume();
+            Ok(jt)
+        } else {
+            err!(self.loc, err_msg)
+        }
+    }
+
     // Match and consume an ebb reference.
     fn match_ebb(&mut self, err_msg: &str) -> Result<Ebb> {
         if let Some(Token::Ebb(ebb)) = self.token() {
@@ -530,6 +559,7 @@ impl<'a> Parser<'a> {
     // preamble-decl ::= * stack-slot-decl
     //                   * function-decl
     //                   * signature-decl
+    //                   * jump-table-decl
     //
     // The parsed decls are added to `ctx` rather than returned.
     fn parse_preamble(&mut self, ctx: &mut Context) -> Result<()> {
@@ -539,13 +569,17 @@ impl<'a> Parser<'a> {
                     self.parse_stack_slot_decl()
                         .and_then(|(num, dat)| ctx.add_ss(num, dat, &self.loc))
                 }
+                Some(Token::JumpTable(..)) => {
+                    self.parse_jump_table_decl()
+                        .and_then(|(num, dat)| ctx.add_jt(num, dat, &self.loc))
+                }
                 // More to come..
                 _ => return Ok(()),
             });
         }
     }
 
-    // Parse a stack slot decl, add to `func`.
+    // Parse a stack slot decl.
     //
     // stack-slot-decl ::= * StackSlot(ss) "=" "stack_slot" Bytes {"," stack-slot-flag}
     fn parse_stack_slot_decl(&mut self) -> Result<(u32, StackSlotData)> {
@@ -564,6 +598,48 @@ impl<'a> Parser<'a> {
         Ok((number, data))
     }
 
+    // Parse a jump table decl.
+    //
+    // jump-table-decl ::= * JumpTable(jt) "=" "jump_table" jt-entry {"," jt-entry}
+    fn parse_jump_table_decl(&mut self) -> Result<(u32, JumpTableData)> {
+        let number = try!(self.match_jt("expected jump table number: jt«n»"));
+        try!(self.match_token(Token::Equal, "expected '=' in jump_table decl"));
+        try!(self.match_identifier("jump_table", "expected 'jump_table'"));
+
+        let mut data = JumpTableData::new();
+
+        // jump-table-decl ::= JumpTable(jt) "=" "jump_table" * jt-entry {"," jt-entry}
+        for idx in 0usize.. {
+            if let Some(dest) = try!(self.parse_jump_table_entry()) {
+                data.set_entry(idx, dest);
+            }
+            if !self.optional(Token::Comma) {
+                return Ok((number, data));
+            }
+        }
+
+        err!(self.loc, "jump_table too long")
+    }
+
+    // jt-entry ::= * Ebb(dest) | "0"
+    fn parse_jump_table_entry(&mut self) -> Result<Option<Ebb>> {
+        match self.token() {
+            Some(Token::Integer(s)) => {
+                if s == "0" {
+                    self.consume();
+                    Ok(None)
+                } else {
+                    err!(self.loc, "invalid jump_table entry '{}'", s)
+                }
+            }
+            Some(Token::Ebb(dest)) => {
+                self.consume();
+                Ok(Some(dest))
+            }
+            _ => err!(self.loc, "expected jump_table entry"),
+        }
+    }
+
     // Parse a function body, add contents to `ctx`.
     //
     // function-body ::= * { extended-basic-block }

tests/parser/branch.cton

@@ -43,3 +43,17 @@ ebb0(vx0: i32, vx1: f32):
 ebb1(vx2: i32, vx3: f32):
     brnz vx0, ebb0(vx2, vx3)
 }
+
+function jumptable() {
+    jt200 = jump_table 0, 0
+    jt2 = jump_table 0, 0, ebb10, ebb40, ebb20, ebb30
+
+ebb10:
+    trap
+ebb20:
+    trap
+ebb30:
+    trap
+ebb40:
+    trap
+}

tests/parser/branch.cton.ref

@@ -37,3 +37,20 @@ ebb0(vx0: i32, vx1: f32):
 ebb1(vx2: i32, vx3: f32):
     brnz vx0, ebb0(vx2, vx3)
 }
+
+function jumptable() {
+    jt0 = jump_table 0
+    jt1 = jump_table 0, 0, ebb0, ebb3, ebb1, ebb2
+
+ebb0:
+    trap
+
+ebb1:
+    trap
+
+ebb2:
+    trap
+
+ebb3:
+    trap
+}