From eee6e78ec1f0f5e0210d67a7d618e5487fabc032 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Tue, 21 Mar 2023 00:40:57 +0900
Subject: [PATCH] PyRef::into_exact_or

---
 vm/src/builtins/int.rs | 57 ++++++++++++++++++++----------------------
 vm/src/object/ext.rs   | 15 +++++++++++
 2 files changed, 42 insertions(+), 30 deletions(-)

diff --git a/vm/src/builtins/int.rs b/vm/src/builtins/int.rs
index 6a169cc2f..b3e6b1c98 100644
--- a/vm/src/builtins/int.rs
+++ b/vm/src/builtins/int.rs
@@ -15,7 +15,7 @@ use crate::{
     },
     protocol::{PyNumber, PyNumberMethods},
     types::{AsNumber, Comparable, Constructor, Hashable, PyComparisonOp, Representable},
-    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult,
+    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
     TryFromBorrowedObject, VirtualMachine,
 };
 use num_bigint::{BigInt, Sign};
@@ -321,7 +321,7 @@ impl PyInt {
 
 #[pyclass(
     flags(BASETYPE),
-    with(Comparable, Hashable, Constructor, AsNumber, Representable)
+    with(PyRef, Comparable, Hashable, Constructor, AsNumber, Representable)
 )]
 impl PyInt {
     #[pymethod(name = "__radd__")]
@@ -521,11 +521,6 @@ impl PyInt {
         Ok(zelf)
     }
 
-    #[pymethod(magic)]
-    fn int(zelf: PyRef<Self>) -> PyRef<Self> {
-        zelf
-    }
-
     #[pymethod(magic)]
     fn pos(&self) -> BigInt {
         self.value.clone()
@@ -537,23 +532,23 @@ impl PyInt {
     }
 
     #[pymethod(magic)]
-    fn trunc(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn trunc(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn floor(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn floor(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn ceil(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn ceil(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
-    fn index(zelf: PyRef<Self>) -> PyRef<Self> {
-        zelf
+    fn index(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pymethod(magic)]
@@ -589,8 +584,8 @@ impl PyInt {
     }
 
     #[pymethod]
-    fn conjugate(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn conjugate(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pyclassmethod]
@@ -659,18 +654,9 @@ impl PyInt {
         Ok(bytes.into())
     }
 
-    #[inline]
-    fn clone_if_subclass(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        if zelf.class().is(vm.ctx.types.int_type) {
-            return zelf;
-        }
-
-        vm.ctx.new_bigint(&zelf.value)
-    }
-
     #[pygetset]
-    fn real(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn real(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pygetset]
@@ -679,8 +665,8 @@ impl PyInt {
     }
 
     #[pygetset]
-    fn numerator(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRef<Self> {
-        Self::clone_if_subclass(zelf, vm)
+    fn numerator(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyRefExact<Self> {
+        zelf.int(vm)
     }
 
     #[pygetset]
@@ -701,6 +687,17 @@ impl PyInt {
     }
 }
 
+#[pyclass]
+impl PyRef<PyInt> {
+    #[pymethod(magic)]
+    fn int(self, vm: &VirtualMachine) -> PyRefExact<PyInt> {
+        self.into_exact_or(&vm.ctx, |zelf| unsafe {
+            // TODO: this is actually safe. we need better interface
+            PyRefExact::new_unchecked(vm.ctx.new_bigint(&zelf.value))
+        })
+    }
+}
+
 impl Comparable for PyInt {
     fn cmp(
         zelf: &Py<Self>,
diff --git a/vm/src/object/ext.rs b/vm/src/object/ext.rs
index 864ec7699..b8aa544fd 100644
--- a/vm/src/object/ext.rs
+++ b/vm/src/object/ext.rs
@@ -9,6 +9,7 @@ use crate::common::{
 use crate::{
     builtins::{PyBaseExceptionRef, PyStrInterned, PyType},
     convert::{IntoPyException, ToPyObject, ToPyResult, TryFromObject},
+    vm::Context,
     VirtualMachine,
 };
 use std::{borrow::Borrow, fmt, marker::PhantomData, ops::Deref, ptr::null_mut};
@@ -107,6 +108,20 @@ impl<T: PyPayload> std::borrow::ToOwned for PyExact<T> {
     }
 }
 
+impl<T: PyPayload> PyRef<T> {
+    pub fn into_exact_or(
+        self,
+        ctx: &Context,
+        f: impl FnOnce(Self) -> PyRefExact<T>,
+    ) -> PyRefExact<T> {
+        if self.class().is(T::class(ctx)) {
+            unsafe { PyRefExact::new_unchecked(self) }
+        } else {
+            f(self)
+        }
+    }
+}
+
 /// PyRef but guaranteed not to be a subtype instance
 #[derive(Debug)]
 #[repr(transparent)]