Compile IfFalse, IfTrue, and Jump instructions (Shopify/zjit#72)

* Compile IfFalse instruction

* Add a TODO comment

* Rename *s_len to num_*s

* Run only gen_param() against block.params

* Add a few more tests

* Wrap label indexes with Label

* Compile blocks in reverse post-order

* Simplify a nested test

* s/get_block/block/

* Return a number instead of an iterator

* Clarify the allocator uses disjoint sets of registers

* Use Display for Block and Insn

* Compile IfTrue and Jump

* Avoid resolving Param instructions

* Always compile Insn::Param as basic block arguments

* Remove an obsoleted variable

* Change it back to use find

* Use find for params too

* Use Display more

* Add more tests

* nested if
* if after if
* if elsif else
* loop after loop
* nested loops
* if in loop
* loop in if

Author: k0kubun

test/ruby/test_zjit.rb

@@ -161,39 +161,191 @@ def test(a, b) = a >= b
     }, call_threshold: 2
   end
 
+  def test_if
+    assert_compiles '[0, nil]', %q{
+      def test(n)
+        if n < 5
+          0
+        end
+      end
+      [test(3), test(7)]
+    }
+  end
 
+  def test_if_else
+    assert_compiles '[0, 1]', %q{
+      def test(n)
+        if n < 5
+          0
+        else
+          1
+        end
+      end
+      [test(3), test(7)]
+    }
+  end
 
-  # FIXME: missing IfFalse insn
-  #def test_if_else
-  #  assert_compiles '[0, 1]', %q{
-  #    def test(n)
-  #      if n < 5
-  #        0
-  #      else
-  #        1
-  #      end
-  #    end
-  #    [test(3), test(7)]
-  #  }, call_threshold: 2
-  #end
+  def test_if_else_params
+    assert_compiles '[1, 20]', %q{
+      def test(n, a, b)
+        if n < 5
+          a
+        else
+          b
+        end
+      end
+      [test(3, 1, 2), test(7, 10, 20)]
+    }
+  end
 
+  def test_if_else_nested
+    assert_compiles '[3, 8, 9, 14]', %q{
+      def test(a, b, c, d, e)
+        if 2 < a
+          if a < 4
+            b
+          else
+            c
+          end
+        else
+          if a < 0
+            d
+          else
+            e
+          end
+        end
+      end
+      [
+        test(-1,  1,  2,  3,  4),
+        test( 0,  5,  6,  7,  8),
+        test( 3,  9, 10, 11, 12),
+        test( 5, 13, 14, 15, 16),
+      ]
+    }
+  end
 
+  def test_if_else_chained
+    assert_compiles '[12, 11, 21]', %q{
+      def test(a)
+        (if 2 < a then 1 else 2 end) + (if a < 4 then 10 else 20 end)
+      end
+      [test(0), test(3), test(5)]
+    }
+  end
+
+  def test_if_elsif_else
+    assert_compiles '[0, 2, 1]', %q{
+      def test(n)
+        if n < 5
+          0
+        elsif 8 < n
+          1
+        else
+          2
+        end
+      end
+      [test(3), test(7), test(9)]
+    }
+  end
+
+  def test_ternary_operator
+    assert_compiles '[1, 20]', %q{
+      def test(n, a, b)
+        n < 5 ? a : b
+      end
+      [test(3, 1, 2), test(7, 10, 20)]
+    }
+  end
 
+  def test_ternary_operator_nested
+    assert_compiles '[2, 21]', %q{
+      def test(n, a, b)
+        (n < 5 ? a : b) + 1
+      end
+      [test(3, 1, 2), test(7, 10, 20)]
+    }
+  end
 
-  # FIXME: need to call twice because of call threshold 2, but
-  # then this fails because of missing FixnumLt
   def test_while_loop
     assert_compiles '10', %q{
-      def loop_fun(n)
+      def test(n)
         i = 0
         while i < n
           i = i + 1
         end
         i
       end
-      loop_fun(10)
-      #loop_fun(10)
-    }, call_threshold: 2
+      test(10)
+    }
+  end
+
+  def test_while_loop_chain
+    assert_compiles '[135, 270]', %q{
+      def test(n)
+        i = 0
+        while i < n
+          i = i + 1
+        end
+        while i < n * 10
+          i = i * 3
+        end
+        i
+      end
+      [test(5), test(10)]
+    }
+  end
+
+  def test_while_loop_nested
+    assert_compiles '[0, 4, 12]', %q{
+      def test(n, m)
+        i = 0
+        while i < n
+          j = 0
+          while j < m
+            j += 2
+          end
+          i += j
+        end
+        i
+      end
+      [test(0, 0), test(1, 3), test(10, 5)]
+    }
+  end
+
+  def test_while_loop_if_else
+    assert_compiles '[9, -1]', %q{
+      def test(n)
+        i = 0
+        while i < n
+          if n >= 10
+            return -1
+          else
+            i = i + 1
+          end
+        end
+        i
+      end
+      [test(9), test(10)]
+    }
+  end
+
+  def test_if_while_loop
+    assert_compiles '[9, 12]', %q{
+      def test(n)
+        i = 0
+        if n < 10
+          while i < n
+            i += 1
+          end
+        else
+          while i < n
+            i += 3
+          end
+        end
+        i
+      end
+      [test(9), test(10)]
+    }
   end
 
   private

zjit/src/asm/arm64/opnd.rs

@@ -1,7 +1,7 @@
 
 
 /// This operand represents a register.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 pub struct A64Reg
 {
     // Size in bits

zjit/src/asm/mod.rs

@@ -2,21 +2,26 @@ use std::collections::BTreeMap;
 //use std::fmt;
 use std::rc::Rc;
 use std::cell::RefCell;
+use std::mem;
 use crate::virtualmem::*;
 
 // Lots of manual vertical alignment in there that rustfmt doesn't handle well.
 #[rustfmt::skip]
 pub mod x86_64;
 pub mod arm64;
 
+/// Index to a label created by cb.new_label()
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct Label(pub usize);
+
 /// Reference to an ASM label
 #[derive(Clone)]
 pub struct LabelRef {
     // Position in the code block where the label reference exists
     pos: usize,
 
     // Label which this refers to
-    label_idx: usize,
+    label: Label,
 
     /// The number of bytes that this label reference takes up in the memory.
     /// It's necessary to know this ahead of time so that when we come back to
@@ -32,9 +37,21 @@ pub struct CodeBlock {
     // Memory for storing the encoded instructions
     mem_block: Rc<RefCell<VirtualMem>>,
 
+    // Memory block size
+    mem_size: usize,
+
     // Current writing position
     write_pos: usize,
 
+    // Table of registered label addresses
+    label_addrs: Vec<usize>,
+
+    // Table of registered label names
+    label_names: Vec<String>,
+
+    // References to labels
+    label_refs: Vec<LabelRef>,
+
     // A switch for keeping comments. They take up memory.
     keep_comments: bool,
 
@@ -50,9 +67,14 @@ pub struct CodeBlock {
 impl CodeBlock {
     /// Make a new CodeBlock
     pub fn new(mem_block: Rc<RefCell<VirtualMem>>, keep_comments: bool) -> Self {
+        let mem_size = mem_block.borrow().virtual_region_size();
         Self {
             mem_block,
+            mem_size,
             write_pos: 0,
+            label_addrs: Vec::new(),
+            label_names: Vec::new(),
+            label_refs: Vec::new(),
             keep_comments,
             asm_comments: BTreeMap::new(),
             dropped_bytes: false,
@@ -144,21 +166,70 @@ impl CodeBlock {
         self.dropped_bytes
     }
 
+    /// Allocate a new label with a given name
+    pub fn new_label(&mut self, name: String) -> Label {
+        assert!(!name.contains(' '), "use underscores in label names, not spaces");
+
+        // This label doesn't have an address yet
+        self.label_addrs.push(0);
+        self.label_names.push(name);
+
+        Label(self.label_addrs.len() - 1)
+    }
+
+    /// Write a label at the current address
+    pub fn write_label(&mut self, label: Label) {
+        self.label_addrs[label.0] = self.write_pos;
+    }
+
     // Add a label reference at the current write position
-    pub fn label_ref(&mut self, _label_idx: usize, _num_bytes: usize, _encode: fn(&mut CodeBlock, i64, i64)) {
-        // TODO: copy labels
+    pub fn label_ref(&mut self, label: Label, num_bytes: usize, encode: fn(&mut CodeBlock, i64, i64)) {
+        assert!(label.0 < self.label_addrs.len());
 
-        //assert!(label_idx < self.label_addrs.len());
+        // Keep track of the reference
+        self.label_refs.push(LabelRef { pos: self.write_pos, label, num_bytes, encode });
 
-        //// Keep track of the reference
-        //self.label_refs.push(LabelRef { pos: self.write_pos, label_idx, num_bytes, encode });
+        // Move past however many bytes the instruction takes up
+        if self.write_pos + num_bytes < self.mem_size {
+            self.write_pos += num_bytes;
+        } else {
+            self.dropped_bytes = true; // retry emitting the Insn after next_page
+        }
+    }
+
+    // Link internal label references
+    pub fn link_labels(&mut self) {
+        let orig_pos = self.write_pos;
+
+        // For each label reference
+        for label_ref in mem::take(&mut self.label_refs) {
+            let ref_pos = label_ref.pos;
+            let label_idx = label_ref.label.0;
+            assert!(ref_pos < self.mem_size);
+
+            let label_addr = self.label_addrs[label_idx];
+            assert!(label_addr < self.mem_size);
+
+            self.write_pos = ref_pos;
+            (label_ref.encode)(self, (ref_pos + label_ref.num_bytes) as i64, label_addr as i64);
+
+            // Assert that we've written the same number of bytes that we
+            // expected to have written.
+            assert!(self.write_pos == ref_pos + label_ref.num_bytes);
+        }
+
+        self.write_pos = orig_pos;
+
+        // Clear the label positions and references
+        self.label_addrs.clear();
+        self.label_names.clear();
+        assert!(self.label_refs.is_empty());
+    }
 
-        //// Move past however many bytes the instruction takes up
-        //if self.has_capacity(num_bytes) {
-        //    self.write_pos += num_bytes;
-        //} else {
-        //    self.dropped_bytes = true; // retry emitting the Insn after next_page
-        //}
+    pub fn clear_labels(&mut self) {
+        self.label_addrs.clear();
+        self.label_names.clear();
+        self.label_refs.clear();
     }
 
     /// Make all the code in the region executable. Call this at the end of a write session.