Address review, add test

Author: Fidget-Spinner

Lib/test/test_generated_cases.py

@@ -2253,5 +2253,40 @@ def test_validate_uop_unused_size_mismatch(self):
                                     "Inputs must have equal sizes"):
             self.run_cases_test(input, input2, output)
 
+    def test_pure_uop_body_copied_in(self):
+        input = """
+        pure op(OP, (foo -- res)) {
+            res = body(foo);
+        }
+        """
+        input2 = """
+        op(OP, (foo -- res)) {
+            res = sym_new_unknown(ctx);
+        }
+        """
+        output = """
+        case OP: {
+            JitOptSymbol *res;
+            if (
+                sym_is_const(ctx, foo)
+            ) {
+                JitOptSymbol *foo_sym = foo;
+                _PyStackRef foo = sym_get_const_as_stackref(ctx, foo_sym);
+                _PyStackRef res_stackref;
+                /* Start of pure uop copied from bytecodes for constant evaluation */
+                res_stackref = body(foo);
+                /* End of pure uop copied from bytecodes for constant evaluation */
+                res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectBorrow(res_stackref));
+                stack_pointer[-1] = res;
+            }
+            else {
+                res = sym_new_unknown(ctx);
+                stack_pointer[-1] = res;
+            }
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
 if __name__ == "__main__":
     unittest.main()

Python/optimizer_cases.c.h

@@ -75,19 +75,6 @@ def write_header(
 """
     )
 
-def skip_to(tkn_iter: TokenIterator, end: str) -> Token:
-    tkn = None
-    parens = 0
-    for tkn in tkn_iter:
-        if tkn.kind == end and parens == 0:
-            return tkn
-        if tkn.kind == "LPAREN":
-            parens += 1
-        if tkn.kind == "RPAREN":
-            parens -= 1
-    assert tkn is not None
-    return tkn
-
 def emit_to(out: CWriter, tkn_iter: TokenIterator, end: str) -> Token:
     parens = 0
     for tkn in tkn_iter:

Tools/cases_generator/generators_common.py

@@ -22,11 +22,9 @@
     write_header,
     Emitter,
     TokenIterator,
-    emit_to,
-    skip_to,
 )
 from cwriter import CWriter
-from typing import TextIO, Callable
+from typing import TextIO
 from lexer import Token
 from stack import Local, Stack, StackError, Storage
 
@@ -87,7 +85,7 @@ def stackref_type_name(var: StackItem) -> str:
         assert False, "Unsafe to convert a symbol to an array-like StackRef."
     if var.type:
         return var.type
-    return f"_PyStackRef "
+    return "_PyStackRef "
 
 def declare_variables(uop: Uop, out: CWriter, skip_inputs: bool) -> None:
     variables = {"unused"}
@@ -210,17 +208,19 @@ def write_uop_pure_evaluation_region_header(
     # No reference management of outputs needed.
     for var in storage.outputs:
         var.in_local = True
-    emitter.emit_tokens(uop, storage, None, False, is_abstract=True)
+    emitter.emit("/* Start of pure uop copied from bytecodes for constant evaluation */\n")
+    emitter.emit_tokens(uop, storage, inst=None, emit_braces=False, is_abstract=True)
     out.start_line()
+    emitter.emit("/* End of pure uop copied from bytecodes for constant evaluation */\n")
     # Finally, assign back the output stackrefs to symbolics.
     for outp in uop.stack.outputs:
         # All new stackrefs are created from new references.
         # That's how the stackref contract works.
         if not outp.peek:
-            out.emit(f"{outp.name} = sym_new_const_steal(ctx, PyStackRef_AsPyObjectBorrow({outp.name}_stackref));\n")
+            emitter.emit(f"{outp.name} = sym_new_const_steal(ctx, PyStackRef_AsPyObjectBorrow({outp.name}_stackref));\n")
         else:
-            out.emit(f"{outp.name} = sym_new_const(ctx, PyStackRef_AsPyObjectBorrow({outp.name}_stackref));\n")
-        storage.flush(out)
+            emitter.emit(f"{outp.name} = sym_new_const(ctx, PyStackRef_AsPyObjectBorrow({outp.name}_stackref));\n")
+    storage.flush(out)
     emitter.emit("}\n")
     emitter.emit("else {\n")
 
@@ -266,7 +266,8 @@ def write_uop(
             out.start_line()
             if override:
                 emitter = OptimizerEmitter(out, {})
-                _, storage = emitter.emit_tokens(override, storage, None, False, is_abstract=True)
+                _, storage = emitter.emit_tokens(override, storage, inst=None,
+                                                 emit_braces=False, is_abstract=True)
                 storage.flush(out)
             else:
                 emit_default(out, uop, stack)