From c6127af8685c2a9b416207e46089cee79d028b85 Mon Sep 17 00:00:00 2001
From: Michael Droettboom <mdboom@gmail.com>
Date: Tue, 8 Oct 2024 13:18:39 -0400
Subject: gh-125063: Emit slices as constants in the bytecode compiler
 (#125064)

* Make slices marshallable

* Emit slices as constants

* Update Python/marshal.c

Co-authored-by: Peter Bierma <zintensitydev@gmail.com>

* Refactor codegen_slice into two functions so it
always has the same net effect

* Fix for free-threaded builds

* Simplify marshal loading of slices

* Only return SUCCESS/ERROR from codegen_slice

---------

Co-authored-by: Mark Shannon <mark@hotpy.org>
Co-authored-by: Peter Bierma <zintensitydev@gmail.com>
---
 Python/codegen.c | 78 ++++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 59 insertions(+), 19 deletions(-)

(limited to 'Python/codegen.c')

diff --git a/Python/codegen.c b/Python/codegen.c
index 896c30cc149..689d2b5124e 100644
--- a/Python/codegen.c
+++ b/Python/codegen.c
@@ -194,6 +194,7 @@ static int codegen_visit_expr(compiler *, expr_ty);
 static int codegen_augassign(compiler *, stmt_ty);
 static int codegen_annassign(compiler *, stmt_ty);
 static int codegen_subscript(compiler *, expr_ty);
+static int codegen_slice_two_parts(compiler *, expr_ty);
 static int codegen_slice(compiler *, expr_ty);
 
 static bool are_all_items_const(asdl_expr_seq *, Py_ssize_t, Py_ssize_t);
@@ -5005,12 +5006,8 @@ codegen_visit_expr(compiler *c, expr_ty e)
         }
         break;
     case Slice_kind:
-    {
-        int n = codegen_slice(c, e);
-        RETURN_IF_ERROR(n);
-        ADDOP_I(c, loc, BUILD_SLICE, n);
+        RETURN_IF_ERROR(codegen_slice(c, e));
         break;
-    }
     case Name_kind:
         return codegen_nameop(c, loc, e->v.Name.id, e->v.Name.ctx);
     /* child nodes of List and Tuple will have expr_context set */
@@ -5023,9 +5020,22 @@ codegen_visit_expr(compiler *c, expr_ty e)
 }
 
 static bool
-is_two_element_slice(expr_ty s)
+is_constant_slice(expr_ty s)
 {
     return s->kind == Slice_kind &&
+        (s->v.Slice.lower == NULL ||
+         s->v.Slice.lower->kind == Constant_kind) &&
+        (s->v.Slice.upper == NULL ||
+         s->v.Slice.upper->kind == Constant_kind) &&
+        (s->v.Slice.step == NULL ||
+         s->v.Slice.step->kind == Constant_kind);
+}
+
+static bool
+should_apply_two_element_slice_optimization(expr_ty s)
+{
+    return !is_constant_slice(s) &&
+           s->kind == Slice_kind &&
            s->v.Slice.step == NULL;
 }
 
@@ -5046,8 +5056,8 @@ codegen_augassign(compiler *c, stmt_ty s)
         break;
     case Subscript_kind:
         VISIT(c, expr, e->v.Subscript.value);
-        if (is_two_element_slice(e->v.Subscript.slice)) {
-            RETURN_IF_ERROR(codegen_slice(c, e->v.Subscript.slice));
+        if (should_apply_two_element_slice_optimization(e->v.Subscript.slice)) {
+            RETURN_IF_ERROR(codegen_slice_two_parts(c, e->v.Subscript.slice));
             ADDOP_I(c, loc, COPY, 3);
             ADDOP_I(c, loc, COPY, 3);
             ADDOP_I(c, loc, COPY, 3);
@@ -5084,7 +5094,7 @@ codegen_augassign(compiler *c, stmt_ty s)
         ADDOP_NAME(c, loc, STORE_ATTR, e->v.Attribute.attr, names);
         break;
     case Subscript_kind:
-        if (is_two_element_slice(e->v.Subscript.slice)) {
+        if (should_apply_two_element_slice_optimization(e->v.Subscript.slice)) {
             ADDOP_I(c, loc, SWAP, 4);
             ADDOP_I(c, loc, SWAP, 3);
             ADDOP_I(c, loc, SWAP, 2);
@@ -5231,8 +5241,10 @@ codegen_subscript(compiler *c, expr_ty e)
     }
 
     VISIT(c, expr, e->v.Subscript.value);
-    if (is_two_element_slice(e->v.Subscript.slice) && ctx != Del) {
-        RETURN_IF_ERROR(codegen_slice(c, e->v.Subscript.slice));
+    if (should_apply_two_element_slice_optimization(e->v.Subscript.slice) &&
+        ctx != Del
+    ) {
+        RETURN_IF_ERROR(codegen_slice_two_parts(c, e->v.Subscript.slice));
         if (ctx == Load) {
             ADDOP(c, loc, BINARY_SLICE);
         }
@@ -5254,15 +5266,9 @@ codegen_subscript(compiler *c, expr_ty e)
     return SUCCESS;
 }
 
-/* Returns the number of the values emitted,
- * thus are needed to build the slice, or -1 if there is an error. */
 static int
-codegen_slice(compiler *c, expr_ty s)
+codegen_slice_two_parts(compiler *c, expr_ty s)
 {
-    int n = 2;
-    assert(s->kind == Slice_kind);
-
-    /* only handles the cases where BUILD_SLICE is emitted */
     if (s->v.Slice.lower) {
         VISIT(c, expr, s->v.Slice.lower);
     }
@@ -5277,11 +5283,45 @@ codegen_slice(compiler *c, expr_ty s)
         ADDOP_LOAD_CONST(c, LOC(s), Py_None);
     }
 
+    return 0;
+}
+
+static int
+codegen_slice(compiler *c, expr_ty s)
+{
+    int n = 2;
+    assert(s->kind == Slice_kind);
+
+    if (is_constant_slice(s)) {
+        PyObject *start = NULL;
+        if (s->v.Slice.lower) {
+            start = s->v.Slice.lower->v.Constant.value;
+        }
+        PyObject *stop = NULL;
+        if (s->v.Slice.upper) {
+            stop = s->v.Slice.upper->v.Constant.value;
+        }
+        PyObject *step = NULL;
+        if (s->v.Slice.step) {
+            step = s->v.Slice.step->v.Constant.value;
+        }
+        PyObject *slice = PySlice_New(start, stop, step);
+        if (slice == NULL) {
+            return ERROR;
+        }
+        ADDOP_LOAD_CONST_NEW(c, LOC(s), slice);
+        return SUCCESS;
+    }
+
+    RETURN_IF_ERROR(codegen_slice_two_parts(c, s));
+
     if (s->v.Slice.step) {
         n++;
         VISIT(c, expr, s->v.Slice.step);
     }
-    return n;
+
+    ADDOP_I(c, LOC(s), BUILD_SLICE, n);
+    return SUCCESS;
 }
 
 
-- 
cgit v1.2.3