diff --git a/pyrefly/lib/alt/expr.rs b/pyrefly/lib/alt/expr.rs index 48cb021c05..7772a7f84b 100644 --- a/pyrefly/lib/alt/expr.rs +++ b/pyrefly/lib/alt/expr.rs @@ -219,7 +219,7 @@ impl Display for ConditionRedundantReason { } } -static MAX_TUPLE_LENGTH: usize = 256; +pub(crate) const MAX_TUPLE_LENGTH: usize = 256; impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { fn synthesized_functional_class_type(&self, call: &ExprCall) -> Option { diff --git a/pyrefly/lib/alt/operators.rs b/pyrefly/lib/alt/operators.rs index 4c6267a755..3a9d2cb44b 100644 --- a/pyrefly/lib/alt/operators.rs +++ b/pyrefly/lib/alt/operators.rs @@ -35,6 +35,7 @@ use crate::alt::answers::LookupAnswer; use crate::alt::answers_solver::AnswersSolver; use crate::alt::call::CallStyle; use crate::alt::callable::CallArg; +use crate::alt::expr::MAX_TUPLE_LENGTH; use crate::alt::unwrap::HintRef; use crate::binding::binding::KeyAnnotation; use crate::config::error_kind::ErrorKind; @@ -315,6 +316,42 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { } } + fn try_tuple_repeat(&self, lhs: &Type, rhs: &Type) -> Option { + let repeat_count = |ty: &Type| match ty { + Type::Literal(box Literal { + value: Lit::Int(n), .. + }) => n.as_i64(), + _ => None, + }; + let (tuple, repeats) = match (lhs, rhs) { + (Type::Tuple(tuple), rhs) => (tuple, repeat_count(rhs)?), + (lhs, Type::Tuple(tuple)) => (tuple, repeat_count(lhs)?), + _ => return None, + }; + if repeats <= 0 { + return Some(self.heap.mk_concrete_tuple(Vec::new())); + } + if repeats == 1 { + return Some(self.heap.mk_tuple(tuple.clone())); + } + let repeats = usize::try_from(repeats).ok()?; + match tuple { + Tuple::Concrete(elements) => { + let total_len = elements.len().checked_mul(repeats)?; + if total_len > MAX_TUPLE_LENGTH { + return None; + } + let mut repeated = Vec::with_capacity(total_len); + for _ in 0..repeats { + repeated.extend(elements.iter().cloned()); + } + Some(self.heap.mk_concrete_tuple(repeated)) + } + Tuple::Unbounded(element) => Some(self.heap.mk_unbounded_tuple((**element).clone())), + Tuple::Unpacked(_) => None, + } + } + pub fn binop_infer( &self, x: &ExprBinOp, @@ -405,6 +442,10 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> { && let Type::Tuple(r) = rhs { self.tuple_concat(l, r) + } else if x.op == Operator::Mult + && let Some(result) = self.try_tuple_repeat(lhs, rhs) + { + result } else if matches!( x.op, Operator::Add diff --git a/pyrefly/lib/test/lsp/lsp_interaction/pytorch_benchmark.rs b/pyrefly/lib/test/lsp/lsp_interaction/pytorch_benchmark.rs index cf56f84412..62429ffd53 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/pytorch_benchmark.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/pytorch_benchmark.rs @@ -48,7 +48,7 @@ fn test_pytorch_error_propagation_latency() { }; // Use all available cores for realistic benchmarking let mut interaction = - LspInteraction::new_with_args(args, NoTelemetry, Some(ThreadCount::AllThreads)); + LspInteraction::new_with_args(args, NoTelemetry, Some(ThreadCount::AllThreads), None); interaction.set_root(pytorch_root.clone()); interaction diff --git a/pyrefly/lib/test/tuple.rs b/pyrefly/lib/test/tuple.rs index 33aea38c58..a971f54f4c 100644 --- a/pyrefly/lib/test/tuple.rs +++ b/pyrefly/lib/test/tuple.rs @@ -485,6 +485,20 @@ def test(x: tuple[int] | tuple[str]) -> None: "#, ); +testcase!( + test_tuple_repeat, + r#" +from typing import assert_type, Literal + +assert_type((42,) * 2, tuple[Literal[42], Literal[42]]) +assert_type(2 * (42,), tuple[Literal[42], Literal[42]]) +assert_type((1, "x") * 2, tuple[Literal[1], Literal["x"], Literal[1], Literal["x"]]) +assert_type((1,) * 0, tuple[()]) +assert_type((1,) * -1, tuple[()]) +assert_type((1,) * 257, tuple[Literal[1], ...]) +"#, +); + testcase!( test_unpack_tuple_with_double_def, r#"