Skip to content

Fix math reduction conversion and numeric protocol handling #9666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Nov 5, 2026
Merged

Fix math reduction conversion and numeric protocol handling #9666

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 77 additions & 104 deletions crates/stdlib/src/math.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,8 @@ mod math {
use crate::vm::{
AsObject, Py, PyObject, PyObjectRef, PyResult, VirtualMachine,
builtins::{
PyFloat, PyInt, PyIntRef, PyModule, PyStrInterned, try_bigint_to_f64, try_f64_to_bigint,
PyFloat, PyInt, PyIntRef, PyModule, PyStrInterned, PyTuple, try_bigint_to_f64,
try_f64_to_bigint,
},
function::{ArgIndex, ArgIntoFloat, ArgIterable, NameCoordinates, OptionalArg, PosArgs},
identifier,
Expand Down Expand Up @@ -279,13 +280,27 @@ mod math {
}

#[pyfunction]
fn dist(p: Vec<ArgIntoFloat>, q: Vec<ArgIntoFloat>, vm: &VirtualMachine) -> PyResult<f64> {
let p = ArgIntoFloat::vec_into_f64(p);
let q = ArgIntoFloat::vec_into_f64(q);
fn dist(p: PyObjectRef, q: PyObjectRef, vm: &VirtualMachine) -> PyResult<f64> {
// Snapshot both points before checking dimensions or converting coordinates.
// Tuple subclasses use their stored elements, just like PyTuple_Check.
let p = match p.downcast_ref::<PyTuple>() {
Some(tuple) => tuple.to_owned(),
None => p.sequence_unchecked().tuple(vm)?,
};
let q = match q.downcast_ref::<PyTuple>() {
Some(tuple) => tuple.to_owned(),
None => q.sequence_unchecked().tuple(vm)?,
};
if p.len() != q.len() {
return Err(vm.new_value_error("both points must have the same number of dimensions"));
}
Ok(pymath::math::dist(&p, &q))
let mut differences = Vec::with_capacity(p.len());
for (p, q) in p.as_slice().iter().zip(q.as_slice()) {
let p: ArgIntoFloat = p.clone().try_into_value(vm)?;
let q: ArgIntoFloat = q.clone().try_into_value(vm)?;
differences.push(p.into_float() - q.into_float());
}
Ok(pymath::math::hypot(&differences))
}

#[pyfunction]
Expand Down Expand Up @@ -541,131 +556,83 @@ mod math {

#[pyfunction]
fn prod(args: ProdArgs, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
use crate::vm::builtins::PyInt;

let iter = args.iterable;
let start = args.start;

// Check if start is provided and what type it is (exact types only, not subclasses)
let (mut obj_result, start_is_int, start_is_float) = match &start {
OptionalArg::Present(s) => {
let is_int = s.class().is(vm.ctx.types.int_type);
let is_float = s.class().is(vm.ctx.types.float_type);
(Some(s.clone()), is_int, is_float)
}
OptionalArg::Missing => (None, true, false), // Default is int 1
let mut items = args.iterable.iter(vm)?;
let mut result = match &args.start {
OptionalArg::Present(start) => start.clone(),
OptionalArg::Missing => vm.ctx.new_int(1).into(),
};

let mut item_iter = iter.iter(vm)?;

// Integer fast path
if start_is_int && !start_is_float {
let mut int_result: i64 = match &start {
OptionalArg::Present(s) => {
if let Some(i) = s.downcast_ref::<PyInt>() {
match i.as_bigint().try_into() {
Ok(v) => v,
Err(_) => {
// Start overflows i64, fall through to generic path
obj_result = Some(s.clone());
i64::MAX // Will be ignored
}
}
} else {
1
}
// Only exact builtins may bypass Python's multiplication methods.
if result.class().is(vm.ctx.types.int_type)
&& let Some(integer) = result.downcast_ref::<PyInt>()
&& let Ok(mut total) = integer.as_bigint().try_into() as Result<i64, _>
{
loop {
let Some(item) = items.next() else {
return Ok(vm.ctx.new_int(total).into());
};
let item = item?;
if item.class().is(vm.ctx.types.int_type)
&& let Some(integer) = item.downcast_ref::<PyInt>()
&& let Ok(value) = integer.as_bigint().try_into() as Result<i64, _>
&& let Some(product) = total.checked_mul(value)
{
total = product;
continue;
}
OptionalArg::Missing => 1,
};

if obj_result.is_none() {
loop {
let item = match item_iter.next() {
Some(r) => r?,
None => return Ok(vm.ctx.new_int(int_result).into()),
};

// Only use fast path for exact int type (not subclasses)
if item.class().is(vm.ctx.types.int_type)
&& let Some(int_item) = item.downcast_ref::<PyInt>()
&& let Ok(b) = int_item.as_bigint().try_into() as Result<i64, _>
&& let Some(product) = int_result.checked_mul(b)
{
int_result = product;
continue;
}

// Overflow or non-int: restore to PyObject and continue
obj_result = Some(vm.ctx.new_int(int_result).into());
let temp = vm._mul(obj_result.as_ref().unwrap(), &item)?;
obj_result = Some(temp);
break;
}
result = vm.ctx.new_int(total).into();
result = vm._mul(&result, &item)?;
break;
}
}

// Float fast path
let obj_float = obj_result
.as_ref()
.and_then(|obj| obj.clone().downcast::<PyFloat>().ok());
if obj_float.is_some() || start_is_float {
let mut flt_result: f64 = if let Some(ref f) = obj_float {
f.to_f64()
} else if start_is_float && let OptionalArg::Present(s) = &start {
s.downcast_ref::<PyFloat>().map_or(1.0, |f| f.to_f64())
} else {
1.0
};

// The integer fallback can return a float subclass. Do not enter this
// path on a payload downcast: its __mul__ must still be called.
if result.class().is(vm.ctx.types.float_type)
&& let Some(float) = result.downcast_ref::<PyFloat>()
{
let mut total = float.to_f64();
loop {
let item = match item_iter.next() {
Some(r) => r?,
None => return Ok(vm.ctx.new_float(flt_result).into()),
let Some(item) = items.next() else {
return Ok(vm.ctx.new_float(total).into());
};

// Only use fast path for exact float/int types (not subclasses)
let item = item?;
if item.class().is(vm.ctx.types.float_type)
&& let Some(f) = item.downcast_ref::<PyFloat>()
&& let Some(float) = item.downcast_ref::<PyFloat>()
{
flt_result *= f.to_f64();
total *= float.to_f64();
continue;
}
if item.class().is(vm.ctx.types.int_type)
&& let Some(i) = item.downcast_ref::<PyInt>()
&& let Ok(v) = i.as_bigint().try_into() as Result<i64, _>
&& let Some(integer) = item.downcast_ref::<PyInt>()
&& let Ok(value) = integer.as_bigint().try_into() as Result<i64, _>
{
flt_result *= v as f64;
total *= value as f64;
continue;
}

// Non-exact-float/int: restore and continue with generic path
obj_result = Some(vm.ctx.new_float(flt_result).into());
let temp = vm._mul(obj_result.as_ref().unwrap(), &item)?;
obj_result = Some(temp);
result = vm.ctx.new_float(total).into();
result = vm._mul(&result, &item)?;
break;
}
}

// Generic path for remaining items
let mut result = obj_result.unwrap_or_else(|| vm.ctx.new_int(1).into());
for item in item_iter {
// Once the specialized paths end, every remaining product uses the
// normal numeric protocol, even if a later result is a builtin again.
for item in items {
let item = item?;
result = vm._mul(&result, &item)?;
}

Ok(result)
}

#[pyfunction]
fn sumprod(
p: ArgIterable<PyObjectRef>,
q: ArgIterable<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<PyObjectRef> {
fn sumprod(p: PyObjectRef, q: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
use crate::vm::builtins::PyInt;

let mut p_iter = p.iter(vm)?;
let mut q_iter = q.iter(vm)?;
let mut p_iter = p.get_iter(vm)?.into_iter::<PyObjectRef>(vm);
let mut q_iter = q.get_iter(vm)?.into_iter::<PyObjectRef>(vm);

// Fast path state
let mut int_path_enabled = true;
Expand All @@ -678,11 +645,13 @@ mod math {
let mut obj_total: Option<PyObjectRef> = None;

loop {
let m_p = p_iter.next();
let m_q = q_iter.next();
// An exception from p must escape before q is advanced. Likewise,
// an exception from q takes precedence over a length mismatch.
let m_p = p_iter.next().transpose()?;
let m_q = q_iter.next().transpose()?;

let (p_i, q_i, finished) = match (m_p, m_q) {
(Some(r_p), Some(r_q)) => (Some(r_p?), Some(r_q?), false),
(Some(p_i), Some(q_i)) => (Some(p_i), Some(q_i), false),
(None, None) => (None, None, true),
_ => return Err(vm.new_value_error("Inputs are not the same length")),
};
Expand Down Expand Up @@ -728,8 +697,12 @@ mod math {

let p_is_exact_float = p_i.class().is(vm.ctx.types.float_type);
let q_is_exact_float = q_i.class().is(vm.ctx.types.float_type);
let p_is_exact_int = p_i.class().is(vm.ctx.types.int_type);
let q_is_exact_int = q_i.class().is(vm.ctx.types.int_type);
// Boolean weights participate in the compensated float
// path, but arbitrary integer subclasses must not.
let p_is_exact_int = p_i.class().is(vm.ctx.types.int_type)
|| p_i.class().is(vm.ctx.types.bool_type);
let q_is_exact_int = q_i.class().is(vm.ctx.types.int_type)
|| q_i.class().is(vm.ctx.types.bool_type);
let p_is_exact_numeric = p_is_exact_float || p_is_exact_int;
let q_is_exact_numeric = q_is_exact_float || q_is_exact_int;
let has_exact_float = p_is_exact_float || q_is_exact_float;
Expand Down
Loading