Add codegen for NewArray instruction (Shopify/zjit#110)

* Show failing test

* Add second test case

* Add empty NewArray setup

* Update opt_tests and fix NewArray instantiation

* Add code generation for NewArray

* Add NewArray ordering test

Author: aidenfoxivey

test/ruby/test_zjit.rb

@@ -187,6 +187,31 @@ def test(a, b) = a >= b
     }, call_threshold: 2
   end
 
+  def test_new_array_empty
+    assert_compiles '[]', %q{
+      def test = []
+      test
+    }
+  end
+
+  def test_new_array_nonempty
+    assert_compiles '[5]', %q{
+      def a = 5
+      def test = [a]
+      test
+    }
+  end
+
+  def test_new_array_order
+    assert_compiles '[3, 2, 1]', %q{
+      def a = 3
+      def b = 2
+      def c = 1
+      def test = [a, b, c]
+      test
+    }
+  end
+
   def test_array_dup
     assert_compiles '[1, 2, 3]', %q{
       def test = [1,2,3]

zjit/src/codegen.rs

@@ -243,6 +243,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     let out_opnd = match insn {
         Insn::PutSelf => gen_putself(),
         Insn::Const { val: Const::Value(val) } => gen_const(*val),
+        Insn::NewArray { elements, state } => gen_new_array(jit, asm, elements, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
         Insn::Param { idx } => unreachable!("block.insns should not have Insn::Param({idx})"),
         Insn::Snapshot { .. } => return Some(()), // we don't need to do anything for this instruction at the moment
@@ -509,6 +510,37 @@ fn gen_array_dup(
     )
 }
 
+/// Compile a new array instruction
+fn gen_new_array(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    elements: &Vec<InsnId>,
+    state: &FrameState,
+) -> lir::Opnd {
+    asm_comment!(asm, "call rb_ary_new");
+
+    // Save PC
+    gen_save_pc(asm, state);
+
+    let length: ::std::os::raw::c_long = elements.len().try_into().expect("Unable to fit length of elements into c_long");
+
+    let new_array = asm.ccall(
+        rb_ary_new_capa as *const u8,
+        vec![lir::Opnd::Imm(length)],
+    );
+
+    for i in 0..elements.len() {
+        let insn_id = elements.get(i as usize).expect("Element should exist at index");
+        let val = jit.get_opnd(*insn_id).unwrap();
+        asm.ccall(
+            rb_ary_push as *const u8,
+            vec![new_array, val]
+        );
+    }
+
+    new_array
+}
+
 /// Compile code that exits from JIT code with a return value
 fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
     // Pop the current frame (ec->cfp++)

zjit/src/hir.rs

@@ -298,7 +298,7 @@ pub enum Insn {
     StringCopy { val: InsnId },
     StringIntern { val: InsnId },
 
-    NewArray { elements: Vec<InsnId> },
+    NewArray { elements: Vec<InsnId>, state: InsnId },
     ArraySet { array: InsnId, idx: usize, val: InsnId },
     ArrayDup { val: InsnId, state: InsnId },
 
@@ -423,7 +423,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
         match &self.inner {
             Insn::Const { val } => { write!(f, "Const {}", val.print(self.ptr_map)) }
             Insn::Param { idx } => { write!(f, "Param {idx}") }
-            Insn::NewArray { elements } => {
+            Insn::NewArray { elements, .. } => {
                 write!(f, "NewArray")?;
                 let mut prefix = " ";
                 for element in elements {
@@ -1633,12 +1633,13 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_newarray => {
                     let count = get_arg(pc, 0).as_usize();
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.clone() });
                     let mut elements = vec![];
                     for _ in 0..count {
                         elements.push(state.stack_pop()?);
                     }
                     elements.reverse();
-                    state.stack_push(fun.push_insn(block, Insn::NewArray { elements }));
+                    state.stack_push(fun.push_insn(block, Insn::NewArray { elements, state: exit_id }));
                 }
                 YARVINSN_duparray => {
                     let val = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
@@ -2046,15 +2047,16 @@ mod infer_tests {
     #[test]
     fn newarray() {
         let mut function = Function::new(std::ptr::null());
-        let val = function.push_insn(function.entry_block, Insn::NewArray { elements: vec![] });
+        // Fake FrameState index of 0usize
+        let val = function.push_insn(function.entry_block, Insn::NewArray { elements: vec![], state: InsnId(0usize) });
         assert_bit_equal(function.infer_type(val), types::ArrayExact);
     }
 
     #[test]
     fn arraydup() {
         let mut function = Function::new(std::ptr::null());
-        let arr = function.push_insn(function.entry_block, Insn::NewArray { elements: vec![] });
         // Fake FrameState index of 0usize
+        let arr = function.push_insn(function.entry_block, Insn::NewArray { elements: vec![], state: InsnId(0usize) });
         let val = function.push_insn(function.entry_block, Insn::ArrayDup { val: arr, state: InsnId(0usize) });
         assert_bit_equal(function.infer_type(val), types::ArrayExact);
     }
@@ -2168,8 +2170,8 @@ mod tests {
         assert_method_hir("test", expect![[r#"
             fn test:
             bb0():
-              v1:ArrayExact = NewArray
-              Return v1
+              v2:ArrayExact = NewArray
+              Return v2
         "#]]);
     }
 
@@ -2179,8 +2181,8 @@ mod tests {
         assert_method_hir("test", expect![[r#"
             fn test:
             bb0(v0:BasicObject):
-              v2:ArrayExact = NewArray v0
-              Return v2
+              v3:ArrayExact = NewArray v0
+              Return v3
         "#]]);
     }
 
@@ -2190,8 +2192,8 @@ mod tests {
         assert_method_hir("test", expect![[r#"
             fn test:
             bb0(v0:BasicObject, v1:BasicObject):
-              v3:ArrayExact = NewArray v0, v1
-              Return v3
+              v4:ArrayExact = NewArray v0, v1
+              Return v4
         "#]]);
     }
 
@@ -2989,8 +2991,8 @@ mod opt_tests {
         assert_optimized_method_hir("test", expect![[r#"
             fn test:
             bb0():
-              v3:Fixnum[5] = Const Value(5)
-              Return v3
+              v4:Fixnum[5] = Const Value(5)
+              Return v4
         "#]]);
     }
 
@@ -3006,8 +3008,8 @@ mod opt_tests {
         assert_optimized_method_hir("test", expect![[r#"
             fn test:
             bb0(v0:BasicObject):
-              v4:Fixnum[5] = Const Value(5)
-              Return v4
+              v5:Fixnum[5] = Const Value(5)
+              Return v5
         "#]]);
     }