From b46e9e825a473122bb7e270dfb658bbd9aef9ee4 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@astral.sh>
Date: Wed, 5 Jun 2024 14:10:37 -0600
Subject: [PATCH] [red-knot] arithmetic on int literals (#11760)

## Summary

Add support for inferring int literal types from basic arithmetic on int
literals. Just to begin showing examples of resolving more complex
expression types, and because this will be useful in testing walrus
expressions.

## Test Plan

Added test.
---
 crates/red_knot/src/semantic/types.rs       | 55 ++++++++++++++++++++-
 crates/red_knot/src/semantic/types/infer.rs | 33 ++++++++++++-
 2 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/crates/red_knot/src/semantic/types.rs b/crates/red_knot/src/semantic/types.rs
index 24178165cc..8db0d7da9c 100644
--- a/crates/red_knot/src/semantic/types.rs
+++ b/crates/red_knot/src/semantic/types.rs
@@ -8,6 +8,7 @@ use crate::semantic::{
 };
 use crate::{FxDashMap, FxIndexSet, Name};
 use ruff_index::{newtype_index, IndexVec};
+use ruff_python_ast as ast;
 use rustc_hash::FxHashMap;
 
 pub(crate) mod infer;
@@ -55,9 +56,9 @@ impl Type {
 
     pub fn get_member(&self, db: &dyn SemanticDb, name: &Name) -> QueryResult<Option<Type>> {
         match self {
-            Type::Any => todo!("attribute lookup on Any type"),
+            Type::Any => Ok(Some(Type::Any)),
             Type::Never => todo!("attribute lookup on Never type"),
-            Type::Unknown => todo!("attribute lookup on Unknown type"),
+            Type::Unknown => Ok(Some(Type::Unknown)),
             Type::Unbound => todo!("attribute lookup on Unbound type"),
             Type::Function(_) => todo!("attribute lookup on Function type"),
             Type::Module(module_id) => module_id.get_member(db, name),
@@ -85,6 +86,56 @@ impl Type {
             }
         }
     }
+
+    // when this is fully fleshed out, it will use the db arg and may return QueryError
+    #[allow(clippy::unnecessary_wraps)]
+    pub fn resolve_bin_op(
+        &self,
+        _db: &dyn SemanticDb,
+        op: ast::Operator,
+        right_ty: Type,
+    ) -> QueryResult<Type> {
+        match self {
+            Type::Any => Ok(Type::Any),
+            Type::Unknown => Ok(Type::Unknown),
+            Type::IntLiteral(n) => {
+                match right_ty {
+                    Type::IntLiteral(m) => {
+                        match op {
+                            ast::Operator::Add => Ok(n
+                                .checked_add(m)
+                                .map(Type::IntLiteral)
+                                // TODO builtins.int
+                                .unwrap_or(Type::Unknown)),
+                            ast::Operator::Sub => Ok(n
+                                .checked_sub(m)
+                                .map(Type::IntLiteral)
+                                // TODO builtins.int
+                                .unwrap_or(Type::Unknown)),
+                            ast::Operator::Mult => Ok(n
+                                .checked_mul(m)
+                                .map(Type::IntLiteral)
+                                // TODO builtins.int
+                                .unwrap_or(Type::Unknown)),
+                            ast::Operator::Div => Ok(n
+                                .checked_div(m)
+                                .map(Type::IntLiteral)
+                                // TODO builtins.int
+                                .unwrap_or(Type::Unknown)),
+                            ast::Operator::Mod => Ok(n
+                                .checked_rem(m)
+                                .map(Type::IntLiteral)
+                                // TODO division by zero error
+                                .unwrap_or(Type::Unknown)),
+                            _ => todo!("complete binop op support for IntLiteral"),
+                        }
+                    }
+                    _ => todo!("complete binop right_ty support for IntLiteral"),
+                }
+            }
+            _ => todo!("complete binop support"),
+        }
+    }
 }
 
 impl From<FunctionTypeId> for Type {
diff --git a/crates/red_knot/src/semantic/types/infer.rs b/crates/red_knot/src/semantic/types/infer.rs
index 20d25c06db..7e8c9c4139 100644
--- a/crates/red_knot/src/semantic/types/infer.rs
+++ b/crates/red_knot/src/semantic/types/infer.rs
@@ -212,7 +212,15 @@ fn infer_expr_type(db: &dyn SemanticDb, file_id: FileId, expr: &ast::Expr) -> Qu
                 .get_member(db, attr_name)
                 .map(|ty| ty.unwrap_or(Type::Unknown))
         }
-        _ => todo!("full expression type resolution"),
+        ast::Expr::BinOp(ast::ExprBinOp {
+            left, op, right, ..
+        }) => {
+            let left_ty = infer_expr_type(db, file_id, left)?;
+            let right_ty = infer_expr_type(db, file_id, right)?;
+            // TODO add reverse bin op support if right <: left
+            left_ty.resolve_bin_op(db, *op, right_ty)
+        }
+        _ => todo!("expression type resolution for {:?}", expr),
     }
 }
 
@@ -488,4 +496,27 @@ mod tests {
 
         assert_public_type(&case, "a", "x", "(Literal[2] | Literal[3] | Literal[4])")
     }
+
+    #[test]
+    fn literal_int_arithmetic() -> anyhow::Result<()> {
+        let case = create_test()?;
+
+        write_to_path(
+            &case,
+            "a.py",
+            "
+                a = 2 + 1
+                b = a - 4
+                c = a * b
+                d = c / 3
+                e = 5 % 3
+            ",
+        )?;
+
+        assert_public_type(&case, "a", "a", "Literal[3]")?;
+        assert_public_type(&case, "a", "b", "Literal[-1]")?;
+        assert_public_type(&case, "a", "c", "Literal[-3]")?;
+        assert_public_type(&case, "a", "d", "Literal[-1]")?;
+        assert_public_type(&case, "a", "e", "Literal[2]")
+    }
 }