Skip to content

Commit 9ea9858

Browse files
youknowoneyouknowone
committed
int/float/complex
1 parent 0184a1f commit 9ea9858

3 files changed

Lines changed: 47 additions & 49 deletions

File tree

‎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.args.len() == 1
138+
&& args.kwargs.is_empty()
139+
&& args.args[0].class().is(vm.ctx.types.float_type)
140+
{
141+
return Ok(args.args[0].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: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -208,11 +208,25 @@ impl Constructor for PyInt {
208208
type Args = IntOptions;
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
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: Self::Args = args.bind(vm)?;
225+
let payload = Self::py_new(&cls, options, vm)?;
226+
payload.into_ref_with_type(vm, cls).map(Into::into)
227+
}
228+
229+
fn py_new(_cls: &Py<PyType>, options: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
216230
let value = if let OptionalArg::Present(val) = options.val_options {
217231
if let OptionalArg::Present(base) = options.base {
218232
let base = base
@@ -223,17 +237,6 @@ impl Constructor for PyInt {
223237
.ok_or_else(|| vm.new_value_error("int() base must be >= 2 and <= 36, or 0"))?;
224238
try_int_radix(&val, base, vm)
225239
} 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-
237240
val.try_int(vm).map(|x| x.as_bigint().clone())
238241
}
239242
} else if let OptionalArg::Present(_) = options.base {
@@ -242,11 +245,7 @@ impl Constructor for PyInt {
242245
Ok(Zero::zero())
243246
}?;
244247

245-
Self::with_value(cls, value, vm).to_pyresult(vm)
246-
}
247-
248-
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
249-
unreachable!("use slot_new")
248+
Ok(Self::from(value))
250249
}
251250
}
252251

0 commit comments

Comments
 (0)