From e70b2db78d7fbc335d437e970d0582b6c309ffb1 Mon Sep 17 00:00:00 2001
From: Ben Lewis <benlewisj@gmail.com>
Date: Sun, 23 Aug 2020 14:28:08 +1200
Subject: [PATCH] Add jit feature when running tests, and run rustfmt.

---
 .github/workflows/ci.yaml |  2 +-
 jit/src/instructions.rs   | 35 ++++++++++++++++++++++-------------
 tests/test_snippets.py    |  4 +++-
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 5efca9a45..f21122e79 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -6,7 +6,7 @@ on:
 name: CI
 
 env:
-  CARGO_ARGS: --features ssl
+  CARGO_ARGS: --features ssl jit
 
 jobs:
   rust_tests:
diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
index 5c832903c..aa21ea7b6 100644
--- a/jit/src/instructions.rs
+++ b/jit/src/instructions.rs
@@ -1,14 +1,14 @@
-use std::collections::HashMap;
 use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
 use rustpython_bytecode::bytecode::{BinaryOperator, Constant, Instruction, NameScope};
+use std::collections::HashMap;
 
 use super::JITCompileError;
 
 pub struct FunctionCompiler<'a, 'b> {
     builder: &'a mut FunctionBuilder<'b>,
     stack: Vec<Value>,
-    variables: HashMap<String, Variable>
+    variables: HashMap<String, Variable>,
 }
 
 impl<'a, 'b> FunctionCompiler<'a, 'b> {
@@ -16,7 +16,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         FunctionCompiler {
             builder,
             stack: Vec::new(),
-            variables: HashMap::new()
+            variables: HashMap::new(),
         }
     }
 
@@ -24,15 +24,18 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         match instruction {
             Instruction::LoadName {
                 name,
-                scope: NameScope::Local
+                scope: NameScope::Local,
             } => {
-                let var = self.variables.get(name).ok_or(JITCompileError::BadBytecode)?;
+                let var = self
+                    .variables
+                    .get(name)
+                    .ok_or(JITCompileError::BadBytecode)?;
                 self.stack.push(self.builder.use_var(*var));
                 Ok(())
             }
             Instruction::StoreName {
                 name,
-                scope: NameScope::Local
+                scope: NameScope::Local,
             } => {
                 let var = match self.variables.get(name) {
                     Some(var) => *var,
@@ -43,7 +46,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                         var
                     }
                 };
-                self.builder.def_var(var, self.stack.pop().ok_or(JITCompileError::BadBytecode)?);
+                self.builder
+                    .def_var(var, self.stack.pop().ok_or(JITCompileError::BadBytecode)?);
                 Ok(())
             }
             Instruction::LoadConst {
@@ -62,22 +66,27 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     .return_(&[self.stack.pop().ok_or(JITCompileError::BadBytecode)?]);
                 Ok(())
             }
-            Instruction::BinaryOperation {
-                op,
-                ..
-            } => {
+            Instruction::BinaryOperation { op, .. } => {
                 let a = self.stack.pop().ok_or(JITCompileError::BadBytecode)?;
                 let b = self.stack.pop().ok_or(JITCompileError::BadBytecode)?;
                 match op {
                     BinaryOperator::Add => {
                         let (out, carry) = self.builder.ins().iadd_ifcout(a, b);
-                        self.builder.ins().trapif(IntCC::Overflow, carry, TrapCode::IntegerOverflow);
+                        self.builder.ins().trapif(
+                            IntCC::Overflow,
+                            carry,
+                            TrapCode::IntegerOverflow,
+                        );
                         self.stack.push(out);
                         Ok(())
                     }
                     BinaryOperator::Subtract => {
                         let (out, carry) = self.builder.ins().isub_ifbout(a, b);
-                        self.builder.ins().trapif(IntCC::Overflow, carry, TrapCode::IntegerOverflow);
+                        self.builder.ins().trapif(
+                            IntCC::Overflow,
+                            carry,
+                            TrapCode::IntegerOverflow,
+                        );
                         self.stack.push(out);
                         Ok(())
                     }
diff --git a/tests/test_snippets.py b/tests/test_snippets.py
index 9c2bb8f4e..7178d9b09 100644
--- a/tests/test_snippets.py
+++ b/tests/test_snippets.py
@@ -25,6 +25,8 @@ TEST_DIRS = {_TestType.functional: os.path.join(TEST_ROOT, "snippets")}
 CPYTHON_RUNNER_DIR = os.path.abspath(os.path.join(ROOT_DIR, "py_code_object"))
 RUSTPYTHON_RUNNER_DIR = os.path.abspath(os.path.join(ROOT_DIR))
 RUSTPYTHON_LIB_DIR = os.path.abspath(os.path.join(ROOT_DIR, "Lib"))
+RUSTPYTHON_FEATURES = ["jit"]
+
 
 @contextlib.contextmanager
 def pushd(path):
@@ -136,7 +138,7 @@ class SampleTestCase(unittest.TestCase):
 
         # cargo stuff
         profile_args = [] if RUST_DEBUG else ["--release"]
-        subprocess.check_call(["cargo", "build", *profile_args])
+        subprocess.check_call(["cargo", "build", "--features", *RUSTPYTHON_FEATURES, *profile_args])
 
     @classmethod
     def tearDownClass(cls):