[ty] Memoize binary operator return types (#24700)

## Summary

Especially for cases like astral-sh/ty#3039,
we were running binary operator inference over and over, and throwing
away everything except the return type. This PR adds a cached query for
_just_ the return type, which is more lightweight than storing the
entire `Bindings` but seemingly still very effective.

For:

```python
import pandas as pd

df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
df["d"] = df["a"] + df["b"] + df["c"] + 1 + (df["a"] ** 2 + df["b"] ** 2 + df["c"] ** 2)
```

Codex reports a 3.32x speedup. Repeating that expression 20 times, Codex
reports a 50.79x speedup (from 52.471s down to 1.033s).

Closes astral-sh/ty#3039.

Author: charliermarsh

crates/ty_python_semantic/src/types/call.rs

@@ -11,6 +11,29 @@ pub(super) use arguments::{Argument, CallArguments};
 pub(super) use bind::{Binding, Bindings, CallableBinding, MatchedArgument};
 
 impl<'db> Type<'db> {
+    /// Memoize the pure return-type part of binary dunder resolution so repeated identical
+    /// expressions don't re-run overload selection at every call site.
+    pub(crate) fn try_call_bin_op_return_type(
+        db: &'db dyn Db,
+        left_ty: Type<'db>,
+        op: ast::Operator,
+        right_ty: Type<'db>,
+    ) -> Option<Type<'db>> {
+        #[salsa::tracked]
+        fn try_call_bin_op_return_type_impl<'db>(
+            db: &'db dyn Db,
+            left_ty: Type<'db>,
+            op: ast::Operator,
+            right_ty: Type<'db>,
+        ) -> Option<Type<'db>> {
+            Type::try_call_bin_op(db, left_ty, op, right_ty)
+                .ok()
+                .map(|bindings| bindings.return_type(db))
+        }
+
+        try_call_bin_op_return_type_impl(db, left_ty, op, right_ty)
+    }
+
     pub(crate) fn try_call_bin_op(
         db: &'db dyn Db,
         left_ty: Type<'db>,

crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs

@@ -450,9 +450,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                         )
                     }
                     // For bounded TypeVars or unconstrained TypeVars, fall through to the default handling.
-                    _ => Type::try_call_bin_op(db, left_ty, op, right_ty)
-                        .map(|outcome| outcome.return_type(db))
-                        .ok(),
+                    _ => Type::try_call_bin_op_return_type(db, left_ty, op, right_ty),
                 }
             }
 
@@ -483,9 +481,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                         )
                     }
                     // For bounded TypeVars or unconstrained TypeVars, fall through to the default handling.
-                    _ => Type::try_call_bin_op(db, left_ty, op, right_ty)
-                        .map(|outcome| outcome.return_type(db))
-                        .ok(),
+                    _ => Type::try_call_bin_op_return_type(db, left_ty, op, right_ty),
                 }
             }
 
@@ -511,9 +507,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                         )
                     }
                     // For bounded TypeVars or unconstrained TypeVars, fall through to the default handling.
-                    _ => Type::try_call_bin_op(db, left_ty, op, right_ty)
-                        .map(|outcome| outcome.return_type(db))
-                        .ok(),
+                    _ => Type::try_call_bin_op_return_type(db, left_ty, op, right_ty),
                 }
             }
 
@@ -524,34 +518,28 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             // positional arguments get. In those cases we need to explicitly delegate to the base
             // type, so that it hits the `Type::Union` branches above.
             (Type::NewTypeInstance(newtype), rhs, _) => {
-                Type::try_call_bin_op(db, left_ty, op, right_ty)
-                    .map(|outcome| outcome.return_type(db))
-                    .ok()
-                    .or_else(|| {
-                        self.infer_binary_expression_type_impl(
-                            node,
-                            emitted_division_by_zero_diagnostic,
-                            newtype.concrete_base_type(db),
-                            rhs,
-                            op,
-                            visitor,
-                        )
-                    })
+                Type::try_call_bin_op_return_type(db, left_ty, op, right_ty).or_else(|| {
+                    self.infer_binary_expression_type_impl(
+                        node,
+                        emitted_division_by_zero_diagnostic,
+                        newtype.concrete_base_type(db),
+                        rhs,
+                        op,
+                        visitor,
+                    )
+                })
             }
             (lhs, Type::NewTypeInstance(newtype), _) => {
-                Type::try_call_bin_op(db, left_ty, op, right_ty)
-                    .map(|outcome| outcome.return_type(db))
-                    .ok()
-                    .or_else(|| {
-                        self.infer_binary_expression_type_impl(
-                            node,
-                            emitted_division_by_zero_diagnostic,
-                            lhs,
-                            newtype.concrete_base_type(db),
-                            op,
-                            visitor,
-                        )
-                    })
+                Type::try_call_bin_op_return_type(db, left_ty, op, right_ty).or_else(|| {
+                    self.infer_binary_expression_type_impl(
+                        node,
+                        emitted_division_by_zero_diagnostic,
+                        lhs,
+                        newtype.concrete_base_type(db),
+                        op,
+                        visitor,
+                    )
+                })
             }
 
             (
@@ -854,9 +842,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                         Some(result)
                     }
 
-                    _ => Type::try_call_bin_op(db, left_ty, op, right_ty)
-                        .map(|outcome| outcome.return_type(db))
-                        .ok(),
+                    _ => Type::try_call_bin_op_return_type(db, left_ty, op, right_ty),
                 }
             }
 
@@ -1039,9 +1025,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                 | Type::TypeGuard(_)
                 | Type::TypedDict(_),
                 op,
-            ) => Type::try_call_bin_op(db, left_ty, op, right_ty)
-                .map(|outcome| outcome.return_type(db))
-                .ok(),
+            ) => Type::try_call_bin_op_return_type(db, left_ty, op, right_ty),
         }
     }