gh-119477: Analyze through BUILD_TUPLE and BINARY_SUBSCR_TUPLE_INT

Include/internal/pycore_optimizer.h

@@ -33,6 +33,8 @@ struct _Py_UopsSymbol {
     int flags;  // 0 bits: Top; 2 or more bits: Bottom
     PyTypeObject *typ;  // Borrowed reference
     PyObject *const_val;  // Owned reference (!)
+    int tuple_count;
+    struct _Py_UopsSymbol **tuple_val;
 };
 
 #define UOP_FORMAT_TARGET 0
@@ -92,6 +94,12 @@ typedef struct ty_arena {
     _Py_UopsSymbol arena[TY_ARENA_SIZE];
 } ty_arena;
 
+typedef struct tup_arena {
+    int tup_curr_number;
+    int tup_max_number;
+    _Py_UopsSymbol *arena[TY_ARENA_SIZE];
+} tup_arena;
+
 struct _Py_UOpsContext {
     char done;
     char out_of_space;
@@ -104,6 +112,9 @@ struct _Py_UOpsContext {
     // Arena for the symbolic types.
     ty_arena t_arena;
 
+    // Arena for the symbolic tuple objects.
+    tup_arena tup_arena;
+
     _Py_UopsSymbol **n_consumed;
     _Py_UopsSymbol **limit;
     _Py_UopsSymbol *locals_and_stack[MAX_ABSTRACT_INTERP_SIZE];
@@ -119,6 +130,8 @@ extern _Py_UopsSymbol *_Py_uop_sym_new_unknown(_Py_UOpsContext *ctx);
 extern _Py_UopsSymbol *_Py_uop_sym_new_not_null(_Py_UOpsContext *ctx);
 extern _Py_UopsSymbol *_Py_uop_sym_new_type(
     _Py_UOpsContext *ctx, PyTypeObject *typ);
+extern _Py_UopsSymbol *_Py_uop_sym_new_tuple(_Py_UOpsContext *ctx, int count);
+extern _Py_UopsSymbol *_Py_uop_sym_tuple_at(_Py_UopsSymbol *sym, int idx);
 extern _Py_UopsSymbol *_Py_uop_sym_new_const(_Py_UOpsContext *ctx, PyObject *const_val);
 extern _Py_UopsSymbol *_Py_uop_sym_new_null(_Py_UOpsContext *ctx);
 extern bool _Py_uop_sym_has_type(_Py_UopsSymbol *sym);

Lib/test/test_capi/test_opt.py

@@ -599,6 +599,24 @@ def _run_with_optimizer(self, testfunc, arg):
         ex = get_first_executor(testfunc)
         return res, ex
 
+    def test_tuple_type_propagation(self):
+        def testfunc(loops):
+            a = 0
+            b = 1
+            for i in range(loops):
+                x = (a, b)
+                x[0] + x[1]
+                a + b
+            return
+
+        res, ex = self._run_with_optimizer(testfunc, 32)
+        self.assertIsNotNone(ex)
+        self.assertEqual(res, None)
+        print("\n".join(list(iter_opnames(ex))))
+        binop_count = [opname for opname in iter_opnames(ex) if opname == "_BINARY_OP_ADD_INT"]
+        guard_both_int_count = [opname for opname in iter_opnames(ex) if opname == "_GUARD_BOTH_INT"]
+        self.assertGreaterEqual(len(binop_count), 2)
+        self.assertLessEqual(len(guard_both_int_count), 1)
 
     def test_int_type_propagation(self):
         def testfunc(loops):

Python/optimizer_analysis.c

@@ -307,6 +307,8 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
 #define sym_is_null _Py_uop_sym_is_null
 #define sym_new_const _Py_uop_sym_new_const
 #define sym_new_null _Py_uop_sym_new_null
+#define sym_new_tuple _Py_uop_sym_new_tuple
+#define sym_tuple_at _Py_uop_sym_tuple_at
 #define sym_has_type _Py_uop_sym_has_type
 #define sym_get_type _Py_uop_sym_get_type
 #define sym_matches_type _Py_uop_sym_matches_type

Python/optimizer_bytecodes.c

@@ -20,6 +20,8 @@ typedef struct _Py_UOpsAbstractFrame _Py_UOpsAbstractFrame;
 #define sym_is_null _Py_uop_sym_is_null
 #define sym_new_const _Py_uop_sym_new_const
 #define sym_new_null _Py_uop_sym_new_null
+#define sym_new_tuple _Py_uop_sym_new_tuple
+#define sym_tuple_at _Py_uop_sym_tuple_at
 #define sym_matches_type _Py_uop_sym_matches_type
 #define sym_get_type _Py_uop_sym_get_type
 #define sym_has_type _Py_uop_sym_has_type
@@ -762,6 +764,25 @@ dummy_func(void) {
         ctx->done = true;
     }
 
+    op(_BUILD_TUPLE, (values[oparg] -- tup)) {
+        tup = sym_new_tuple(ctx, oparg);
+        assert(tup != NULL);
+        assert(!ctx->out_of_space);
+        for (int i = 0; i < oparg; i++) {
+            tup->tuple_val[i] = values[i];
+        }
+    }
+
+    op(_BINARY_SUBSCR_TUPLE_INT, (tuple, sub -- res)) {
+        if (sym_has_type(tuple) && sym_matches_type(tuple, &PyTuple_Type) && sym_is_const(sub) ) {
+            PyObject* value = sym_get_const(sub);
+            Py_ssize_t index = ((PyLongObject*)value)->long_value.ob_digit[0];
+            res = sym_tuple_at(tuple, index);
+        } else {
+            res = sym_new_unknown(ctx);
+        }
+    }
+
 
 // END BYTECODES //
 

Python/optimizer_symbols.c

@@ -76,10 +76,45 @@ sym_new(_Py_UOpsContext *ctx)
     self->flags = 0;
     self->typ = NULL;
     self->const_val = NULL;
+    self->tuple_val = NULL;
 
     return self;
 }
 
+static _Py_UopsSymbol **
+sym_new_array(_Py_UOpsContext *ctx, int count)
+{
+    _Py_UopsSymbol **start = &ctx->tup_arena.arena[ctx->tup_arena.tup_curr_number];
+    if (ctx->tup_arena.tup_curr_number >= ctx->tup_arena.tup_max_number - count) {
+        OPT_STAT_INC(optimizer_failure_reason_no_memory);
+        DPRINTF(1, "out of space for symbolic expression type\n");
+        return NULL;
+    }
+    ctx->tup_arena.tup_curr_number += count;
+    return start;
+}
+
+_Py_UopsSymbol *
+_Py_uop_sym_new_tuple(_Py_UOpsContext *ctx, int count)
+{
+    _Py_UopsSymbol *self = _Py_uop_sym_new_unknown(ctx);
+    if (self == NULL) {
+        return out_of_space(ctx);
+    }
+    _Py_uop_sym_set_type(ctx, self, &PyTuple_Type);
+    _Py_UopsSymbol **items = sym_new_array(ctx, count);
+    if (items == NULL) {
+        return out_of_space(ctx);
+    }
+    self->tuple_val = items;
+    return self;
+}
+
+_Py_UopsSymbol *
+_Py_uop_sym_tuple_at(_Py_UopsSymbol *sym, int idx) {
+    return sym->tuple_val[idx];
+}
+
 static inline void
 sym_set_flag(_Py_UopsSymbol *sym, int flag)
 {
@@ -376,6 +411,10 @@ _Py_uop_abstractcontext_init(_Py_UOpsContext *ctx)
     ctx->t_arena.ty_curr_number = 0;
     ctx->t_arena.ty_max_number = TY_ARENA_SIZE;
 
+    // Setup the arena for sym expressions.
+    ctx->tup_arena.tup_curr_number = 0;
+    ctx->tup_arena.tup_max_number = TY_ARENA_SIZE;
+
     // Frame setup
     ctx->curr_frame_depth = 0;
 }