Firmament2: SOURCE_BEGIN/END around compilation and execution

Author: jamesclement1776

Python/compile.c

@@ -29,6 +29,8 @@
 
 #include <stdbool.h>
 
+#include "firmament2.h"
+
 
 #undef SUCCESS
 #undef ERROR
@@ -103,7 +105,6 @@ typedef struct _PyCompiler {
     bool c_save_nested_seqs;     /* if true, construct recursive instruction sequences
                                   * (including instructions for nested code objects)
                                   */
-    int c_disable_warning;
 } compiler;
 
 static int
@@ -136,7 +137,7 @@ compiler_setup(compiler *c, mod_ty mod, PyObject *filename,
     c->c_optimize = (optimize == -1) ? _Py_GetConfig()->optimization_level : optimize;
     c->c_save_nested_seqs = false;
 
-    if (!_PyAST_Preprocess(mod, arena, filename, c->c_optimize, merged, 0, 1)) {
+    if (!_PyAST_Preprocess(mod, arena, filename, c->c_optimize, merged, 0)) {
         return ERROR;
     }
     c->c_st = _PySymtable_Build(mod, filename, &c->c_future);
@@ -766,9 +767,6 @@ _PyCompile_PushFBlock(compiler *c, location loc,
     f->fb_loc = loc;
     f->fb_exit = exit;
     f->fb_datum = datum;
-    if (t == COMPILE_FBLOCK_FINALLY_END) {
-        c->c_disable_warning++;
-    }
     return SUCCESS;
 }
 
@@ -780,9 +778,6 @@ _PyCompile_PopFBlock(compiler *c, fblocktype t, jump_target_label block_label)
     u->u_nfblocks--;
     assert(u->u_fblock[u->u_nfblocks].fb_type == t);
     assert(SAME_JUMP_TARGET_LABEL(u->u_fblock[u->u_nfblocks].fb_block, block_label));
-    if (t == COMPILE_FBLOCK_FINALLY_END) {
-        c->c_disable_warning--;
-    }
 }
 
 fblockinfo *
@@ -1210,9 +1205,6 @@ _PyCompile_Error(compiler *c, location loc, const char *format, ...)
 int
 _PyCompile_Warn(compiler *c, location loc, const char *format, ...)
 {
-    if (c->c_disable_warning) {
-        return 0;
-    }
     va_list vargs;
     va_start(vargs, format);
     PyObject *msg = PyUnicode_FromFormatV(format, vargs);
@@ -1473,7 +1465,7 @@ _PyCompile_OptimizeAndAssemble(compiler *c, int addNone)
 
     return optimize_and_assemble_code_unit(u, const_cache, code_flags, filename);
 }
-
+/* ==================== Firmament2: source scope around public API ==================== */
 PyCodeObject *
 _PyAST_Compile(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
                int optimize, PyArena *arena)
@@ -1489,7 +1481,31 @@ _PyAST_Compile(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
     assert(co || PyErr_Occurred());
     return co;
 }
+/* Public API: place SOURCE_BEGIN/SOURCE_END so tokenizer+AST+codegen can inherit scope
+   when callers go through PyAST_CompileObject(). */
+PyCodeObject *
+PyAST_CompileObject(mod_ty mod, PyObject *filename, PyCompilerFlags *pflags,
+                    int optimize, PyArena *arena)
+{
+    const char *fn_utf8 = NULL;
+    int began_scope = 0;
 
+    if (_firm2_enabled()) {
+        if (filename && PyUnicode_Check(filename)) {
+            fn_utf8 = PyUnicode_AsUTF8(filename);  /* borrowed */
+        }
+        _firm2_source_begin(fn_utf8);
+        began_scope = 1;
+    }
+
+    PyCodeObject *co = _PyAST_Compile(mod, filename, pflags, optimize, arena);
+
+    if (began_scope) {
+        _firm2_source_end();
+    }
+    return co;
+}
+/* ================== end Firmament2: source scope around public API =================== */
 int
 _PyCompile_AstPreprocess(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
                          int optimize, PyArena *arena, int no_const_folding)
@@ -1502,7 +1518,7 @@ _PyCompile_AstPreprocess(mod_ty mod, PyObject *filename, PyCompilerFlags *cf,
     if (optimize == -1) {
         optimize = _Py_GetConfig()->optimization_level;
     }
-    if (!_PyAST_Preprocess(mod, arena, filename, optimize, flags, no_const_folding, 0)) {
+    if (!_PyAST_Preprocess(mod, arena, filename, optimize, flags, no_const_folding)) {
         return -1;
     }
     return 0;

Python/pythonrun.c

@@ -41,6 +41,9 @@
 #  include "windows.h"
 #endif
 
+/* Firmament2 */
+#include "firmament2.h"
+
 /* Forward */
 static void flush_io(void);
 static PyObject *run_mod(mod_ty, PyObject *, PyObject *, PyObject *,
@@ -1234,36 +1237,49 @@ _PyRun_StringFlagsWithName(const char *str, PyObject* name, int start,
                            int generate_new_source)
 {
     PyObject *ret = NULL;
-    mod_ty mod;
-    PyArena *arena;
-
-    arena = _PyArena_New();
+    PyArena *arena = _PyArena_New();
     if (arena == NULL)
         return NULL;
 
     PyObject* source = NULL;
     _Py_DECLARE_STR(anon_string, "<string>");
 
-    if (name) {
+    if (!name) {
+        name = &_Py_STR(anon_string);
+    } else {
         source = PyUnicode_FromString(str);
         if (!source) {
             PyErr_Clear();
         }
-    } else {
-        name = &_Py_STR(anon_string);
     }
 
-    mod = _PyParser_ASTFromString(str, name, start, flags, arena);
+    /* Begin source scope before parsing so tokens are attributed correctly */
+    int began_scope = 0;
+    const char *fn_utf8 = NULL;
+    if (_firm2_enabled()) {
+        if (name && PyUnicode_Check(name)) {
+            fn_utf8 = PyUnicode_AsUTF8(name);
+        }
+        _firm2_source_begin(fn_utf8);
+        began_scope = 1;
+    }
+
+    mod_ty mod = _PyParser_ASTFromString(str, name, start, flags, arena);
 
-   if (mod != NULL) {
-        ret = run_mod(mod, name, globals, locals, flags, arena, source, generate_new_source);
+    if (mod != NULL) {
+        ret = run_mod(mod, name, globals, locals, flags, arena,
+                      source, generate_new_source);
     }
+
+    if (began_scope) {
+        _firm2_source_end();
+    }
+
     Py_XDECREF(source);
     _PyArena_Free(arena);
     return ret;
 }
 
-
 PyObject *
 PyRun_StringFlags(const char *str, int start, PyObject *globals,
                      PyObject *locals, PyCompilerFlags *flags) {
@@ -1280,27 +1296,37 @@ pyrun_file(FILE *fp, PyObject *filename, int start, PyObject *globals,
         return NULL;
     }
 
-    mod_ty mod;
-    mod = _PyParser_ASTFromFile(fp, filename, NULL, start, NULL, NULL,
-                                flags, NULL, arena);
+    /* Begin source scope before we tokenize/parse so tokens get filename/source_id */
+    int began_scope = 0;
+    const char *fn_utf8 = NULL;
+    if (_firm2_enabled()) {
+        if (filename && PyUnicode_Check(filename)) {
+            fn_utf8 = PyUnicode_AsUTF8(filename);
+        }
+        _firm2_source_begin(fn_utf8);
+        began_scope = 1;
+    }
+
+    mod_ty mod = _PyParser_ASTFromFile(fp, filename, NULL, start, NULL, NULL,
+                                       flags, NULL, arena);
 
     if (closeit) {
         fclose(fp);
     }
 
-    PyObject *ret;
+    PyObject *ret = NULL;
     if (mod != NULL) {
         ret = run_mod(mod, filename, globals, locals, flags, arena, NULL, 0);
     }
-    else {
-        ret = NULL;
+
+    if (began_scope) {
+        _firm2_source_end();
     }
-    _PyArena_Free(arena);
 
+    _PyArena_Free(arena);
     return ret;
 }
 
-
 PyObject *
 PyRun_FileExFlags(FILE *fp, const char *filename, int start, PyObject *globals,
                   PyObject *locals, int closeit, PyCompilerFlags *flags)
@@ -1508,30 +1534,60 @@ PyObject *
 Py_CompileStringObject(const char *str, PyObject *filename, int start,
                        PyCompilerFlags *flags, int optimize)
 {
-    PyCodeObject *co;
-    mod_ty mod;
+    PyObject *res = NULL;          /* return value: AST or PyCodeObject* */
+    mod_ty mod = NULL;
     PyArena *arena = _PyArena_New();
-    if (arena == NULL)
+    if (arena == NULL) {
         return NULL;
+    }
+
+    /* ---- Firmament2 source scope: begin before tokenization ---- */
+    const char *fn_utf8 = NULL;
+    int began_scope = 0;
+    if (_firm2_enabled()) {
+        if (filename && PyUnicode_Check(filename)) {
+            fn_utf8 = PyUnicode_AsUTF8(filename);  /* borrowed; may be NULL */
+        }
+        if (fn_utf8 == NULL) {
+            /* Provide a stable label for string inputs if not a PyUnicode */
+            fn_utf8 = "<string>";
+        }
+        _firm2_source_begin(fn_utf8);
+        began_scope = 1;
+    }
+    /* ------------------------------------------------------------ */
 
     mod = _PyParser_ASTFromString(str, filename, start, flags, arena);
     if (mod == NULL) {
-        _PyArena_Free(arena);
-        return NULL;
+        goto done;  /* arena freed + scope ended in done: block */
     }
+
     if (flags && (flags->cf_flags & PyCF_ONLY_AST)) {
-        int syntax_check_only = ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_ONLY_AST); /* unoptiomized AST */
-        if (_PyCompile_AstPreprocess(mod, filename, flags, optimize, arena, syntax_check_only) < 0) {
-            _PyArena_Free(arena);
-            return NULL;
+        /* Return AST object (optionally preprocessed), not bytecode */
+        int syntax_check_only =
+            ((flags->cf_flags & PyCF_OPTIMIZED_AST) == PyCF_ONLY_AST); /* unoptimized AST */
+        if (_PyCompile_AstPreprocess(mod, filename, flags, optimize, arena,
+                                     syntax_check_only) < 0) {
+            goto done;
         }
-        PyObject *result = PyAST_mod2obj(mod);
-        _PyArena_Free(arena);
-        return result;
+        res = PyAST_mod2obj(mod);
+        goto done;
+    }
+
+    /* Compile to code object WITHOUT calling the public wrapper, so we
+       keep exactly one SOURCE_BEGIN/SOURCE_END from this scope. */
+    {
+        PyCodeObject *co = _PyAST_Compile(mod, filename, flags, optimize, arena);
+        res = (PyObject *)co;  /* may be NULL on error */
     }
-    co = _PyAST_Compile(mod, filename, flags, optimize, arena);
+
+done:
     _PyArena_Free(arena);
-    return (PyObject *)co;
+
+    if (began_scope) {
+        _firm2_source_end();
+    }
+    return res;
 }
 
 PyObject *