From edb4fcf731f2b9a2a0bb6cff9e27560b1aa80ef9 Mon Sep 17 00:00:00 2001
From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
Date: Thu, 15 Feb 2024 02:25:29 +0800
Subject: [PATCH 1/2] Type and constant propagation for int BINARY_OPs

---
 Python/optimizer_analysis.c                   | 12 ++++
 .../tier2_redundancy_eliminator_bytecodes.c   | 53 +++++++++++++++--
 Python/tier2_redundancy_eliminator_cases.c.h  | 59 +++++++++++++++----
 3 files changed, 108 insertions(+), 16 deletions(-)

diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 49974520de924d..d73bc310345f41 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -341,6 +341,18 @@ sym_new_const(_Py_UOpsAbstractInterpContext *ctx, PyObject *const_val)
     return temp;
 }
 
+static inline bool
+is_const(_Py_UOpsSymType *sym)
+{
+    return sym->const_val != NULL;
+}
+
+static inline PyObject *
+get_const(_Py_UOpsSymType *sym)
+{
+    return sym->const_val;
+}
+
 static _Py_UOpsSymType*
 sym_new_null(_Py_UOpsAbstractInterpContext *ctx)
 {
diff --git a/Python/tier2_redundancy_eliminator_bytecodes.c b/Python/tier2_redundancy_eliminator_bytecodes.c
index 3272b187f20d0e..f53e9d7d9fb56a 100644
--- a/Python/tier2_redundancy_eliminator_bytecodes.c
+++ b/Python/tier2_redundancy_eliminator_bytecodes.c
@@ -81,12 +81,53 @@ dummy_func(void) {
 
 
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {
-        // TODO constant propagation
-        (void)left;
-        (void)right;
-        res = sym_new_known_type(ctx, &PyLong_Type);
-        if (res == NULL) {
-            goto out_of_space;
+        if (is_const(left) && is_const(right)) {
+            PyObject *temp = _PyLong_Add(get_const(left), get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
+        }
+    }
+
+    op(_BINARY_OP_SUBTRACT_INT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            PyObject *temp = _PyLong_Subtract(get_const(left), get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
+        }
+    }
+
+    op(_BINARY_OP_MULTIPLY_INT, (left, right -- res)) {
+        if (is_const(left) && is_const(right)) {
+            PyObject *temp = _PyLong_Multiply(get_const(left), get_const(right));
+            if (temp == NULL) {
+                goto error;
+            }
+            res = sym_new_const(ctx, temp);
+            // TODO replace opcode with constant propagated one and add tests!
+        }
+        else {
+            res = sym_new_known_type(ctx, &PyLong_Type);
+            if (res == NULL) {
+                goto out_of_space;
+            }
         }
     }
 
diff --git a/Python/tier2_redundancy_eliminator_cases.c.h b/Python/tier2_redundancy_eliminator_cases.c.h
index c2b7bbaf1c4481..f56051373674c4 100644
--- a/Python/tier2_redundancy_eliminator_cases.c.h
+++ b/Python/tier2_redundancy_eliminator_cases.c.h
@@ -180,9 +180,25 @@
         }
 
         case _BINARY_OP_MULTIPLY_INT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                PyObject *temp = _PyLong_Multiply(get_const(left), get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;
@@ -194,12 +210,19 @@
             _Py_UOpsSymType *res;
             right = stack_pointer[-1];
             left = stack_pointer[-2];
-            // TODO constant propagation
-            (void)left;
-            (void)right;
-            res = sym_new_known_type(ctx, &PyLong_Type);
-            if (res == NULL) {
-                goto out_of_space;
+            if (is_const(left) && is_const(right)) {
+                PyObject *temp = _PyLong_Add(get_const(left), get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
             }
             stack_pointer[-2] = res;
             stack_pointer += -1;
@@ -207,9 +230,25 @@
         }
 
         case _BINARY_OP_SUBTRACT_INT: {
+            _Py_UOpsSymType *right;
+            _Py_UOpsSymType *left;
             _Py_UOpsSymType *res;
-            res = sym_new_unknown(ctx);
-            if (res == NULL) goto out_of_space;
+            right = stack_pointer[-1];
+            left = stack_pointer[-2];
+            if (is_const(left) && is_const(right)) {
+                PyObject *temp = _PyLong_Subtract(get_const(left), get_const(right));
+                if (temp == NULL) {
+                    goto error;
+                }
+                res = sym_new_const(ctx, temp);
+                // TODO replace opcode with constant propagated one and add tests!
+            }
+            else {
+                res = sym_new_known_type(ctx, &PyLong_Type);
+                if (res == NULL) {
+                    goto out_of_space;
+                }
+            }
             stack_pointer[-2] = res;
             stack_pointer += -1;
             break;

From 6ca3f78a36d3cf6a6aebc719315cb5ddb1ba406a Mon Sep 17 00:00:00 2001
From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
Date: Thu, 15 Feb 2024 03:00:38 +0800
Subject: [PATCH 2/2] add asserts

---
 Python/tier2_redundancy_eliminator_bytecodes.c | 15 ++++++++++++---
 Python/tier2_redundancy_eliminator_cases.c.h   | 15 ++++++++++++---
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/Python/tier2_redundancy_eliminator_bytecodes.c b/Python/tier2_redundancy_eliminator_bytecodes.c
index f53e9d7d9fb56a..39ea0eef627632 100644
--- a/Python/tier2_redundancy_eliminator_bytecodes.c
+++ b/Python/tier2_redundancy_eliminator_bytecodes.c
@@ -82,7 +82,10 @@ dummy_func(void) {
 
     op(_BINARY_OP_ADD_INT, (left, right -- res)) {
         if (is_const(left) && is_const(right)) {
-            PyObject *temp = _PyLong_Add(get_const(left), get_const(right));
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Add((PyLongObject *)get_const(left),
+                                         (PyLongObject *)get_const(right));
             if (temp == NULL) {
                 goto error;
             }
@@ -99,7 +102,10 @@ dummy_func(void) {
 
     op(_BINARY_OP_SUBTRACT_INT, (left, right -- res)) {
         if (is_const(left) && is_const(right)) {
-            PyObject *temp = _PyLong_Subtract(get_const(left), get_const(right));
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Subtract((PyLongObject *)get_const(left),
+                                              (PyLongObject *)get_const(right));
             if (temp == NULL) {
                 goto error;
             }
@@ -116,7 +122,10 @@ dummy_func(void) {
 
     op(_BINARY_OP_MULTIPLY_INT, (left, right -- res)) {
         if (is_const(left) && is_const(right)) {
-            PyObject *temp = _PyLong_Multiply(get_const(left), get_const(right));
+            assert(PyLong_CheckExact(get_const(left)));
+            assert(PyLong_CheckExact(get_const(right)));
+            PyObject *temp = _PyLong_Multiply((PyLongObject *)get_const(left),
+                                              (PyLongObject *)get_const(right));
             if (temp == NULL) {
                 goto error;
             }
diff --git a/Python/tier2_redundancy_eliminator_cases.c.h b/Python/tier2_redundancy_eliminator_cases.c.h
index f56051373674c4..a9617f51ef4615 100644
--- a/Python/tier2_redundancy_eliminator_cases.c.h
+++ b/Python/tier2_redundancy_eliminator_cases.c.h
@@ -186,7 +186,10 @@
             right = stack_pointer[-1];
             left = stack_pointer[-2];
             if (is_const(left) && is_const(right)) {
-                PyObject *temp = _PyLong_Multiply(get_const(left), get_const(right));
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Multiply((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
                 if (temp == NULL) {
                     goto error;
                 }
@@ -211,7 +214,10 @@
             right = stack_pointer[-1];
             left = stack_pointer[-2];
             if (is_const(left) && is_const(right)) {
-                PyObject *temp = _PyLong_Add(get_const(left), get_const(right));
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Add((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
                 if (temp == NULL) {
                     goto error;
                 }
@@ -236,7 +242,10 @@
             right = stack_pointer[-1];
             left = stack_pointer[-2];
             if (is_const(left) && is_const(right)) {
-                PyObject *temp = _PyLong_Subtract(get_const(left), get_const(right));
+                assert(PyLong_CheckExact(get_const(left)));
+                assert(PyLong_CheckExact(get_const(right)));
+                PyObject *temp = _PyLong_Subtract((PyLongObject *)get_const(left),
+                    (PyLongObject *)get_const(right));
                 if (temp == NULL) {
                     goto error;
                 }