Skip to content

Commit 2a3836d

Browse files
youknowoneyouknowone
committed
int/float/complex
1 parent 0184a1f commit 2a3836d

File tree

3 files changed

+43
-46
lines changed

3 files changed

+43
-46
lines changed

‎crates/vm/src/builtins/complex.rs‎

Lines changed: 16 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -152,22 +152,25 @@ impl Constructor for PyComplex {
152152
type Args = ComplexArgs;
153153

154154
fn slot_new(cls: PyTypeRef, func_args: FuncArgs, vm: &VirtualMachine) -> PyResult {
155+
// Optimization: return exact complex as-is (only when imag is not provided)
156+
if cls.is(vm.ctx.types.complex_type)
157+
&& func_args.args.len() == 1
158+
&& func_args.kwargs.is_empty()
159+
&& func_args.args[0].class().is(vm.ctx.types.complex_type)
160+
{
161+
return Ok(func_args.args[0].clone());
162+
}
163+
155164
let args: Self::Args = func_args.bind(vm)?;
165+
let payload = Self::py_new(&cls, args, vm)?;
166+
payload.into_ref_with_type(vm, cls).map(Into::into)
167+
}
168+
169+
fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
156170
let imag_missing = args.imag.is_missing();
157171
let (real, real_was_complex) = match args.real {
158172
OptionalArg::Missing => (Complex64::new(0.0, 0.0), false),
159173
OptionalArg::Present(val) => {
160-
let val = if cls.is(vm.ctx.types.complex_type) && imag_missing {
161-
match val.downcast_exact::<Self>(vm) {
162-
Ok(c) => {
163-
return Ok(c.into_pyref().into());
164-
}
165-
Err(val) => val,
166-
}
167-
} else {
168-
val
169-
};
170-
171174
if let Some(c) = val.try_complex(vm)? {
172175
c
173176
} else if let Some(s) = val.downcast_ref::<PyStr>() {
@@ -180,9 +183,7 @@ impl Constructor for PyComplex {
180183
.to_str()
181184
.and_then(rustpython_literal::complex::parse_str)
182185
.ok_or_else(|| vm.new_value_error("complex() arg is a malformed string"))?;
183-
return Self::from(Complex64 { re, im })
184-
.into_ref_with_type(vm, cls)
185-
.map(Into::into);
186+
return Ok(Self::from(Complex64 { re, im }));
186187
} else {
187188
return Err(vm.new_type_error(format!(
188189
"complex() first argument must be a string or a number, not '{}'",
@@ -222,13 +223,7 @@ impl Constructor for PyComplex {
222223
imag.re
223224
};
224225
let value = Complex64::new(final_real, final_imag);
225-
Self::from(value)
226-
.into_ref_with_type(vm, cls)
227-
.map(Into::into)
228-
}
229-
230-
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
231-
unreachable!("use slot_new")
226+
Ok(Self::from(value))
232227
}
233228
}
234229

‎crates/vm/src/builtins/float.rs‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -132,28 +132,32 @@ impl Constructor for PyFloat {
132132
type Args = OptionalArg<PyObjectRef>;
133133

134134
fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
135+
// Optimization: return exact float as-is
136+
if cls.is(vm.ctx.types.float_type)
137+
&& args.kwargs.is_empty()
138+
&& let Some(first) = args.args.first()
139+
&& first.class().is(vm.ctx.types.float_type)
140+
{
141+
return Ok(first.clone());
142+
}
143+
135144
let arg: Self::Args = args.bind(vm)?;
145+
let payload = Self::py_new(&cls, arg, vm)?;
146+
payload.into_ref_with_type(vm, cls).map(Into::into)
147+
}
148+
149+
fn py_new(_cls: &Py<PyType>, arg: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
136150
let float_val = match arg {
137151
OptionalArg::Missing => 0.0,
138152
OptionalArg::Present(val) => {
139-
if cls.is(vm.ctx.types.float_type) && val.class().is(vm.ctx.types.float_type) {
140-
return Ok(val);
141-
}
142-
143153
if let Some(f) = val.try_float_opt(vm) {
144154
f?.value
145155
} else {
146156
float_from_string(val, vm)?
147157
}
148158
}
149159
};
150-
Self::from(float_val)
151-
.into_ref_with_type(vm, cls)
152-
.map(Into::into)
153-
}
154-
155-
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
156-
unreachable!("use slot_new")
160+
Ok(Self::from(float_val))
157161
}
158162
}
159163

‎crates/vm/src/builtins/int.rs‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -205,14 +205,23 @@ fn inner_truediv(i1: &BigInt, i2: &BigInt, vm: &VirtualMachine) -> PyResult {
205205
}
206206

207207
impl Constructor for PyInt {
208-
type Args = IntOptions;
208+
type Args = FuncArgs;
209209

210210
fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
211-
let options: Self::Args = args.bind(vm)?;
212211
if cls.is(vm.ctx.types.bool_type) {
213212
return Err(vm.new_type_error("int.__new__(bool) is not safe, use bool.__new__()"));
214213
}
215214

215+
// Optimization: return exact int as-is (only for exact int type, not subclasses)
216+
if cls.is(vm.ctx.types.int_type)
217+
&& args.args.len() == 1
218+
&& args.kwargs.is_empty()
219+
&& args.args[0].class().is(vm.ctx.types.int_type)
220+
{
221+
return Ok(args.args[0].clone());
222+
}
223+
224+
let options: IntOptions = args.bind(vm)?;
216225
let value = if let OptionalArg::Present(val) = options.val_options {
217226
if let OptionalArg::Present(base) = options.base {
218227
let base = base
@@ -223,17 +232,6 @@ impl Constructor for PyInt {
223232
.ok_or_else(|| vm.new_value_error("int() base must be >= 2 and <= 36, or 0"))?;
224233
try_int_radix(&val, base, vm)
225234
} else {
226-
let val = if cls.is(vm.ctx.types.int_type) {
227-
match val.downcast_exact::<Self>(vm) {
228-
Ok(i) => {
229-
return Ok(i.into_pyref().into());
230-
}
231-
Err(val) => val,
232-
}
233-
} else {
234-
val
235-
};
236-
237235
val.try_int(vm).map(|x| x.as_bigint().clone())
238236
}
239237
} else if let OptionalArg::Present(_) = options.base {
@@ -242,7 +240,7 @@ impl Constructor for PyInt {
242240
Ok(Zero::zero())
243241
}?;
244242

245-
Self::with_value(cls, value, vm).to_pyresult(vm)
243+
Self::with_value(cls, value, vm).map(Into::into)
246244
}
247245

248246
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {

0 commit comments

Comments
 (0)