Merge pull request #1502 from youknowone/small-int-cache

Small int cache like CPython

tests/snippets/bools.py

@@ -14,12 +14,12 @@ assert not 0.0
 
 assert not None
 
-assert bool() == False
-assert bool(1) == True
-assert bool({}) == False
+assert bool() is False
+assert bool(1) is True
+assert bool({}) is False
 
-assert bool(NotImplemented) == True
-assert bool(...) == True
+assert bool(NotImplemented) is True
+assert bool(...) is True
 
 if not 1:
     raise BaseException
@@ -105,8 +105,8 @@ class TestMagicMethodLenOne:
         return 1
 
 
-assert bool(TestMagicMethodLenZero()) == False
-assert bool(TestMagicMethodLenOne()) == True
+assert bool(TestMagicMethodLenZero()) is False
+assert bool(TestMagicMethodLenOne()) is True
 
 
 # check __len__ and __bool__

tests/snippets/ints.py

@@ -23,12 +23,12 @@ assert int("101", base=2) == 5
 
 # magic methods should only be implemented for other ints
 
-assert (1).__eq__(1) == True
-assert (1).__ne__(1) == False
-assert (1).__gt__(1) == False
-assert (1).__ge__(1) == True
-assert (1).__lt__(1) == False
-assert (1).__le__(1) == True
+assert (1).__eq__(1) is True
+assert (1).__ne__(1) is False
+assert (1).__gt__(1) is False
+assert (1).__ge__(1) is True
+assert (1).__lt__(1) is False
+assert (1).__le__(1) is True
 assert (1).__add__(1) == 2
 assert (1).__radd__(1) == 2
 assert (2).__sub__(1) == 1
@@ -306,3 +306,8 @@ b = 03 + 2j
 
 with assert_raises(SyntaxError):
     exec(src)
+
+# Small int cache in [-5..256]
+assert 1 is 1  # noqa
+x = 6
+assert 5 is (x-1)  # noqa

vm/src/pyobject.rs

@@ -8,7 +8,7 @@ use std::rc::Rc;
 
 use num_bigint::BigInt;
 use num_complex::Complex64;
-use num_traits::{One, Zero};
+use num_traits::{One, ToPrimitive, Zero};
 
 use crate::bytecode;
 use crate::dictdatatype::DictKey;
@@ -92,6 +92,9 @@ impl fmt::Display for PyObject<dyn PyObjectPayload> {
     }
 }
 
+const INT_CACHE_POOL_MIN: i32 = -5;
+const INT_CACHE_POOL_MAX: i32 = 256;
+
 #[derive(Debug)]
 pub struct PyContext {
     pub true_value: PyIntRef,
@@ -104,6 +107,7 @@ pub struct PyContext {
 
     pub types: TypeZoo,
     pub exceptions: exceptions::ExceptionZoo,
+    pub int_cache_pool: Vec<PyObjectRef>,
 }
 
 pub type PyNotImplementedRef = PyRef<PyNotImplemented>;
@@ -149,6 +153,10 @@ impl PyContext {
             create_type("NotImplementedType", &types.type_type, &types.object_type);
         let not_implemented = create_object(PyNotImplemented, &not_implemented_type);
 
+        let int_cache_pool = (INT_CACHE_POOL_MIN..=INT_CACHE_POOL_MAX)
+            .map(|v| create_object(PyInt::new(BigInt::from(v)), &types.int_type).into_object())
+            .collect();
+
         let true_value = create_object(PyInt::new(BigInt::one()), &types.bool_type);
         let false_value = create_object(PyInt::new(BigInt::zero()), &types.bool_type);
 
@@ -159,12 +167,13 @@ impl PyContext {
             false_value,
             not_implemented,
             none,
+            empty_tuple,
             ellipsis,
             ellipsis_type,
 
             types,
             exceptions,
-            empty_tuple,
+            int_cache_pool,
         };
         initialize_types(&context);
 
@@ -364,7 +373,14 @@ impl PyContext {
         self.types.object_type.clone()
     }
 
-    pub fn new_int<T: Into<BigInt>>(&self, i: T) -> PyObjectRef {
+    #[inline]
+    pub fn new_int<T: Into<BigInt> + ToPrimitive>(&self, i: T) -> PyObjectRef {
+        if let Some(i) = i.to_i32() {
+            if i >= INT_CACHE_POOL_MIN && i <= INT_CACHE_POOL_MAX {
+                let inner_idx = (i - INT_CACHE_POOL_MIN) as usize;
+                return self.int_cache_pool[inner_idx].clone();
+            }
+        }
         PyObject::new(PyInt::new(i), self.int_type(), None)
     }
 

vm/src/vm.rs

@@ -41,6 +41,7 @@ use crate::stdlib;
 use crate::sysmodule;
 use arr_macro::arr;
 use num_bigint::BigInt;
+use num_traits::ToPrimitive;
 #[cfg(feature = "rustpython-compiler")]
 use rustpython_compiler::{compile, error::CompileError};
 
@@ -293,11 +294,13 @@ impl VirtualMachine {
     }
 
     /// Create a new python int object.
-    pub fn new_int<T: Into<BigInt>>(&self, i: T) -> PyObjectRef {
+    #[inline]
+    pub fn new_int<T: Into<BigInt> + ToPrimitive>(&self, i: T) -> PyObjectRef {
         self.ctx.new_int(i)
     }
 
     /// Create a new python bool object.
+    #[inline]
     pub fn new_bool(&self, b: bool) -> PyObjectRef {
         self.ctx.new_bool(b)
     }