Basic block support in JIT

compiler/src/compile.rs

@@ -1388,14 +1388,18 @@ impl<O: OutputStream> Compiler<O> {
         self.emit(Instruction::CompareOperation {
             op: to_operator(ops.last().unwrap()),
         });
-        self.emit(Instruction::Jump { target: last_label });
 
-        // early exit left us with stack: `rhs, comparison_result`. We need to clean up rhs.
-        self.set_label(break_label);
-        self.emit(Instruction::Rotate { amount: 2 });
-        self.emit(Instruction::Pop);
+        if vals.len() > 2 {
+            self.emit(Instruction::Jump { target: last_label });
+
+            // early exit left us with stack: `rhs, comparison_result`. We need to clean up rhs.
+            self.set_label(break_label);
+            self.emit(Instruction::Rotate { amount: 2 });
+            self.emit(Instruction::Pop);
+
+            self.set_label(last_label);
+        }
 
-        self.set_label(last_label);
         Ok(())
     }
 

jit/src/instructions.rs

@@ -1,6 +1,8 @@
 use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
-use rustpython_bytecode::bytecode::{BinaryOperator, Constant, Instruction, NameScope};
+use rustpython_bytecode::bytecode::{
+    BinaryOperator, CodeObject, ComparisonOperator, Constant, Instruction, Label, NameScope,
+};
 use std::collections::HashMap;
 
 use super::{JitCompileError, JitSig, JitType};
@@ -26,6 +28,7 @@ pub struct FunctionCompiler<'a, 'b> {
     builder: &'a mut FunctionBuilder<'b>,
     stack: Vec<JitValue>,
     variables: HashMap<String, Local>,
+    label_to_block: HashMap<Label, Block>,
     pub(crate) sig: JitSig,
 }
 
@@ -40,6 +43,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             builder,
             stack: Vec::new(),
             variables: HashMap::new(),
+            label_to_block: HashMap::new(),
             sig: JitSig {
                 args: arg_types.to_vec(),
                 ret: None,
@@ -76,8 +80,70 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         }
     }
 
-    pub fn add_instruction(&mut self, instruction: &Instruction) -> Result<(), JitCompileError> {
+    fn boolean_val(&mut self, val: JitValue) -> Result<Value, JitCompileError> {
+        match val.ty {
+            JitType::Float => Err(JitCompileError::NotSupported),
+            JitType::Int => Ok(val.val),
+        }
+    }
+
+    pub fn compile(&mut self, bytecode: &CodeObject) -> Result<(), JitCompileError> {
+        let offset_to_label: HashMap<&usize, &Label> =
+            bytecode.label_map.iter().map(|(k, v)| (v, k)).collect();
+
+        for (offset, instruction) in bytecode.instructions.iter().enumerate() {
+            if let Some(&label) = offset_to_label.get(&offset) {
+                let builder = &mut self.builder;
+                let block = self
+                    .label_to_block
+                    .entry(*label)
+                    .or_insert_with(|| builder.create_block());
+
+                // If the current block is not terminated/filled just jump
+                // into the new block.
+                if !self.builder.is_filled() {
+                    self.builder.ins().jump(*block, &[]);
+                }
+
+                self.builder.switch_to_block(*block);
+            }
+
+            // Sometimes the bytecode contains instructions after a return
+            // just ignore those until we are at the next label
+            if self.builder.is_filled() {
+                continue;
+            }
+
+            self.add_instruction(&instruction)?;
+        }
+
+        Ok(())
+    }
+
+    fn add_instruction(&mut self, instruction: &Instruction) -> Result<(), JitCompileError> {
         match instruction {
+            Instruction::JumpIfFalse { target } => {
+                let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
+
+                let then_block = self.builder.create_block();
+                self.label_to_block.insert(*target, then_block);
+
+                let val = self.boolean_val(cond)?;
+                self.builder.ins().brz(val, then_block, &[]);
+
+                let block = self.builder.create_block();
+                self.builder.ins().fallthrough(block, &[]);
+                self.builder.switch_to_block(block);
+
+                Ok(())
+            }
+            Instruction::Jump { target } => {
+                let target_block = self.builder.create_block();
+                self.label_to_block.insert(*target, target_block);
+                self.builder.ins().jump(target_block, &[]);
+
+                Ok(())
+            }
             Instruction::LoadName {
                 name,
                 scope: NameScope::Local,
@@ -139,6 +205,34 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 self.builder.ins().return_(&[val.val]);
                 Ok(())
             }
+            Instruction::CompareOperation { op, .. } => {
+                // the rhs is popped off first
+                let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
+                let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
+
+                match (a.ty, b.ty) {
+                    (JitType::Int, JitType::Int) => {
+                        let cond = match op {
+                            ComparisonOperator::Equal => IntCC::Equal,
+                            ComparisonOperator::NotEqual => IntCC::NotEqual,
+                            ComparisonOperator::Less => IntCC::SignedLessThan,
+                            ComparisonOperator::LessOrEqual => IntCC::SignedLessThanOrEqual,
+                            ComparisonOperator::Greater => IntCC::SignedGreaterThan,
+                            ComparisonOperator::GreaterOrEqual => IntCC::SignedLessThanOrEqual,
+                            _ => return Err(JitCompileError::NotSupported),
+                        };
+
+                        let val = self.builder.ins().icmp(cond, a.val, b.val);
+                        self.stack.push(JitValue {
+                            val,
+                            ty: JitType::Int, // TODO: Boolean
+                        });
+
+                        Ok(())
+                    }
+                    _ => Err(JitCompileError::NotSupported),
+                }
+            }
             Instruction::BinaryOperation { op, .. } => {
                 // the rhs is popped off first
                 let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

jit/src/lib.rs

@@ -63,20 +63,18 @@ impl Jit {
         let entry_block = builder.create_block();
         builder.append_block_params_for_function_params(entry_block);
         builder.switch_to_block(entry_block);
-        builder.seal_block(entry_block);
 
         let sig = {
             let mut arg_names = bytecode.arg_names.clone();
             arg_names.extend(bytecode.kwonlyarg_names.iter().cloned());
             let mut compiler = FunctionCompiler::new(&mut builder, &arg_names, args, entry_block);
 
-            for instruction in &bytecode.instructions {
-                compiler.add_instruction(instruction)?;
-            }
+            compiler.compile(bytecode)?;
 
             compiler.sig
         };
 
+        builder.seal_all_blocks();
         builder.finalize();
 
         let id = self.module.declare_function(

jit/tests/int_tests.rs

@@ -22,3 +22,35 @@ fn test_sub() {
     assert_eq!(sub(7, 10), Ok(-3));
     assert_eq!(sub(-3, -10), Ok(7));
 }
+
+#[test]
+fn test_eq() {
+    let eq = jit_function! { eq(a:i64, b:i64) -> i64 => r##"
+        def eq(a: int, b: int):
+            if a == b:
+                return 1
+            return 0
+    "## };
+
+    assert_eq!(eq(0, 0), Ok(1));
+    assert_eq!(eq(1, 1), Ok(1));
+    assert_eq!(eq(0, 1), Ok(0));
+    assert_eq!(eq(-200, 200), Ok(0));
+}
+
+#[test]
+fn test_gt() {
+    let gt = jit_function! { gt(a:i64, b:i64) -> i64 => r##"
+        def gt(a: int, b: int):
+            if a > b:
+                return 1
+            return 0
+    "## };
+
+    assert_eq!(gt(5, 2), Ok(1));
+    assert_eq!(gt(2, 5), Ok(0));
+    assert_eq!(gt(2, 2), Ok(0));
+    assert_eq!(gt(5, 5), Ok(0));
+    assert_eq!(gt(-1, -10), Ok(1));
+    assert_eq!(gt(1, -1), Ok(1));
+}

jit/tests/misc_tests.rs

@@ -65,3 +65,22 @@ fn test_args_builder() {
     assert!(args.is_some());
     assert_eq!(args.unwrap().invoke(), Some(AbiValue::Int(1)));
 }
+
+#[test]
+fn test_if_else() {
+    let if_else = jit_function! { if_else(a:i64) -> i64 => r##"
+        def if_else(a: int):
+            if a:
+                return 42
+            else:
+                return 0
+
+            # Prevent type failure from implicit `return None`
+            return 0
+    "## };
+
+    assert_eq!(if_else(0), Ok(0));
+    assert_eq!(if_else(1), Ok(42));
+    assert_eq!(if_else(-1), Ok(42));
+    assert_eq!(if_else(100), Ok(42));
+}