From 70b05cf620d911a5dcac3db024c8920810badc7a Mon Sep 17 00:00:00 2001
From: Tom Schuster <evilpies@gmail.com>
Date: Sun, 11 Oct 2020 16:59:45 +0200
Subject: [PATCH] bool support

---
 jit/src/instructions.rs             | 30 +++++++++++++++++++++++------
 jit/src/lib.rs                      | 11 +++++++++--
 vm/src/builtins/function/jitfunc.rs |  7 ++++++-
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
index 82e4dec6d..a0b4338cc 100644
--- a/jit/src/instructions.rs
+++ b/jit/src/instructions.rs
@@ -1,7 +1,8 @@
 use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
 use rustpython_bytecode::bytecode::{
-    UnaryOperator, BinaryOperator, CodeObject, ComparisonOperator, Constant, Instruction, Label, NameScope,
+    BinaryOperator, CodeObject, ComparisonOperator, Constant, Instruction, Label, NameScope,
+    UnaryOperator,
 };
 use std::collections::HashMap;
 
@@ -85,9 +86,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             JitType::Float => {
                 let zero = self.builder.ins().f64const(0);
                 let val = self.builder.ins().fcmp(FloatCC::NotEqual, val.val, zero);
-                Ok(self.builder.ins().bint(types::I64, val))
-            },
-            JitType::Int => Ok(val.val),
+                Ok(self.builder.ins().bint(types::I8, val))
+            }
+            JitType::Int => {
+                let zero = self.builder.ins().iconst(types::I64, 0);
+                let val = self.builder.ins().icmp(IntCC::NotEqual, val.val, zero);
+                Ok(self.builder.ins().bint(types::I8, val))
+            }
+            JitType::Bool => Ok(val.val),
         }
     }
 
@@ -228,8 +234,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
 
                         let val = self.builder.ins().icmp(cond, a.val, b.val);
                         self.stack.push(JitValue {
-                            val,
-                            ty: JitType::Int, // TODO: Boolean
+                            val: self.builder.ins().bint(types::I8, val),
+                            ty: JitType::Bool,
                         });
 
                         Ok(())
@@ -264,6 +270,18 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                         }
                         _ => Err(JitCompileError::NotSupported),
                     },
+                    JitType::Bool => match op {
+                        UnaryOperator::Not => {
+                            let val = self.boolean_val(a)?;
+                            let not_val = self.builder.ins().bxor_imm(val, 1);
+                            self.stack.push(JitValue {
+                                val: not_val,
+                                ty: JitType::Bool,
+                            });
+                            Ok(())
+                        }
+                        _ => Err(JitCompileError::NotSupported),
+                    },
                     _ => Err(JitCompileError::NotSupported),
                 }
             }
diff --git a/jit/src/lib.rs b/jit/src/lib.rs
index 1defaf7f9..9103c845f 100644
--- a/jit/src/lib.rs
+++ b/jit/src/lib.rs
@@ -166,6 +166,7 @@ impl JitSig {
 pub enum JitType {
     Int,
     Float,
+    Bool,
 }
 
 impl JitType {
@@ -173,6 +174,7 @@ impl JitType {
         match self {
             Self::Int => types::I64,
             Self::Float => types::F64,
+            Self::Bool => types::I8,
         }
     }
 
@@ -180,6 +182,7 @@ impl JitType {
         match self {
             Self::Int => libffi::middle::Type::i64(),
             Self::Float => libffi::middle::Type::f64(),
+            Self::Bool => libffi::middle::Type::u8(),
         }
     }
 }
@@ -188,6 +191,7 @@ impl JitType {
 pub enum AbiValue {
     Float(f64),
     Int(i64),
+    Bool(bool),
 }
 
 impl AbiValue {
@@ -195,6 +199,7 @@ impl AbiValue {
         match self {
             AbiValue::Int(ref i) => libffi::middle::Arg::new(i),
             AbiValue::Float(ref f) => libffi::middle::Arg::new(f),
+            AbiValue::Bool(ref b) => libffi::middle::Arg::new(b),
         }
     }
 }
@@ -217,7 +222,7 @@ impl TryFrom<AbiValue> for i64 {
     fn try_from(value: AbiValue) -> Result<Self, Self::Error> {
         match value {
             AbiValue::Int(i) => Ok(i),
-            AbiValue::Float(_) => Err(()),
+            _ => Err(()),
         }
     }
 }
@@ -227,8 +232,8 @@ impl TryFrom<AbiValue> for f64 {
 
     fn try_from(value: AbiValue) -> Result<Self, Self::Error> {
         match value {
-            AbiValue::Int(_) => Err(()),
             AbiValue::Float(f) => Ok(f),
+            _ => Err(()),
         }
     }
 }
@@ -244,6 +249,7 @@ fn type_check(ty: &JitType, val: &AbiValue) -> Result<(), JitArgumentError> {
 union UnTypedAbiValue {
     float: f64,
     int: i64,
+    boolean: u8,
     _void: (),
 }
 
@@ -252,6 +258,7 @@ impl UnTypedAbiValue {
         match ty {
             JitType::Int => AbiValue::Int(self.int),
             JitType::Float => AbiValue::Float(self.float),
+            JitType::Bool => AbiValue::Bool(self.boolean != 0),
         }
     }
 }
diff --git a/vm/src/builtins/function/jitfunc.rs b/vm/src/builtins/function/jitfunc.rs
index 5c315cb0b..5c9a5854b 100644
--- a/vm/src/builtins/function/jitfunc.rs
+++ b/vm/src/builtins/function/jitfunc.rs
@@ -1,6 +1,6 @@
 use crate::builtins::dict::PyDictRef;
 use crate::builtins::function::{PyFunction, PyFunctionRef};
-use crate::builtins::{float, int};
+use crate::builtins::{float, int, pybool};
 use crate::exceptions::PyBaseExceptionRef;
 use crate::function::FuncArgs;
 use crate::pyobject::{
@@ -35,6 +35,7 @@ impl IntoPyObject for AbiValue {
         match self {
             AbiValue::Int(i) => i.into_pyobject(vm),
             AbiValue::Float(f) => f.into_pyobject(vm),
+            AbiValue::Bool(b) => b.into_pyobject(vm),
         }
     }
 }
@@ -50,6 +51,8 @@ fn get_jit_arg_type(dict: &PyDictRef, name: &str, vm: &VirtualMachine) -> PyResu
             Ok(JitType::Int)
         } else if value.is(&vm.ctx.types.float_type) {
             Ok(JitType::Float)
+        } else if value.is(&vm.ctx.types.bool_type) {
+            Ok(JitType::Bool)
         } else {
             Err(new_jit_error(
                 "Jit requires argument to be either int or float".to_owned(),
@@ -113,6 +116,8 @@ fn get_jit_value(vm: &VirtualMachine, obj: &PyObjectRef) -> Result<AbiValue, Arg
             .ok_or(ArgsError::IntOverflow)
     } else if cls.is(&vm.ctx.types.float_type) {
         Ok(AbiValue::Float(float::get_value(&obj)))
+    } else if cls.is(&vm.ctx.types.bool_type) {
+        Ok(AbiValue::Bool(pybool::get_value(&obj)))
     } else {
         Err(ArgsError::NonJitType)
     }