Be a little more precise about how the temporary reference tracer follows line numbers.

Author: braxtonmckee

typed_python/PyTemporaryReferenceTracer.cpp

@@ -50,33 +50,47 @@ bool PyTemporaryReferenceTracer::isLineNewStatement(PyObject* code, int line) {
 
 
 int PyTemporaryReferenceTracer::globalTraceFun(PyObject* dummyObj, PyFrameObject* frame, int what, PyObject* arg) {
-    if (frame != globalTracer.mostRecentEmptyFrame) {
-        bool shouldProcess = true;
-
-        if (what == PyTrace_LINE) {
-            shouldProcess = globalTracer.isLineNewStatement(
-                (PyObject*)frame->f_code,
-                frame->f_lineno
-            );
-        }
+    if (frame != globalTracer.mostRecentEmptyFrame &&
+        globalTracer.frameToActions.find(frame) != globalTracer.frameToActions.end()) {
+        // always process exception and return statements
+        bool forceProcess = (
+            what == PyTrace_EXCEPTION ||
+            what == PyTrace_RETURN
+        );
 
-        if (shouldProcess) {
-            auto it = globalTracer.frameToHandles.find(frame);
+        // we process any statement on a line that's a new statement
+        bool shouldProcess = globalTracer.isLineNewStatement(
+            (PyObject*)frame->f_code,
+            PyFrame_GetLineNumber(frame)
+        );
 
-            if (it == globalTracer.frameToHandles.end()) {
+        if (shouldProcess || forceProcess) {
+            auto it = globalTracer.frameToActions.find(frame);
+
+            if (it == globalTracer.frameToActions.end()) {
                 globalTracer.mostRecentEmptyFrame = frame;
             } else {
                 globalTracer.mostRecentEmptyFrame = nullptr;
 
-                for (auto objAndAction: it->second) {
-                    if (objAndAction.second == TraceAction::ConvertTemporaryReference) {
-                        ((PyInstance*)objAndAction.first)->resolveTemporaryReference();
-                    }
+                std::vector<FrameAction> persistingActions;
 
-                    decref(objAndAction.first);
+                for (auto& frameAction: it->second) {
+                    if (frameAction.lineNumber != PyFrame_GetLineNumber(frame) || forceProcess) {
+                        if (frameAction.action == TraceAction::ConvertTemporaryReference) {
+                            ((PyInstance*)frameAction.obj)->resolveTemporaryReference();
+                        }
+
+                        decref(frameAction.obj);
+                    } else {
+                        persistingActions.push_back(frameAction);
+                    }
                 }
 
-                globalTracer.frameToHandles.erase(it);
+                if (persistingActions.size()) {
+                    it->second = persistingActions;
+                } else {
+                    globalTracer.frameToActions.erase(it);
+                }
             }
         }
     }
@@ -88,7 +102,7 @@ int PyTemporaryReferenceTracer::globalTraceFun(PyObject* dummyObj, PyFrameObject
         );
     }
 
-    if (globalTracer.frameToHandles.size() == 0) {
+    if (globalTracer.frameToActions.size() == 0) {
         // uninstall ourself
         PyEval_SetTrace(globalTracer.priorTraceFunc, globalTracer.priorTraceFuncArg);
         decref(globalTracer.priorTraceFuncArg);
@@ -115,7 +129,13 @@ void PyTemporaryReferenceTracer::installGlobalTraceHandlerIfNecessary() {
 
 void PyTemporaryReferenceTracer::traceObject(PyObject* o, PyFrameObject* f) {
     // mark that we're going to trace
-    globalTracer.frameToHandles[f].push_back(std::make_pair(incref(o), TraceAction::ConvertTemporaryReference));
+    globalTracer.frameToActions[f].push_back(
+        FrameAction(
+            incref(o),
+            TraceAction::ConvertTemporaryReference,
+            PyFrame_GetLineNumber(f)
+        )
+    );
 
     if (globalTracer.mostRecentEmptyFrame == f) {
         globalTracer.mostRecentEmptyFrame = nullptr;
@@ -126,7 +146,13 @@ void PyTemporaryReferenceTracer::traceObject(PyObject* o, PyFrameObject* f) {
 
 void PyTemporaryReferenceTracer::keepaliveForCurrentInstruction(PyObject* o, PyFrameObject* f) {
     // mark that we're going to trace
-    globalTracer.frameToHandles[f].push_back(std::make_pair(incref(o), TraceAction::Decref));
+    globalTracer.frameToActions[f].push_back(
+        FrameAction(
+            incref(o),
+            TraceAction::Decref,
+            PyFrame_GetLineNumber(f)
+        )
+    );
 
     if (globalTracer.mostRecentEmptyFrame == f) {
         globalTracer.mostRecentEmptyFrame = nullptr;

typed_python/PyTemporaryReferenceTracer.hpp

@@ -32,7 +32,24 @@ class PyTemporaryReferenceTracer {
         priorTraceFuncArg(nullptr)
     {}
 
-    std::unordered_map<PyFrameObject*, std::vector<std::pair<PyObject*, TraceAction> > > frameToHandles;
+    // perform an action on the first instruction where a
+    // frame goes out of scope or where it is no longer on the
+    // given line number
+    class FrameAction {
+    public:
+        FrameAction(PyObject* inO, TraceAction inA, int inLine) :
+            obj(inO),
+            action(inA),
+            lineNumber(inLine)
+        {
+        }
+
+        PyObject* obj;
+        TraceAction action;
+        int lineNumber;
+    };
+
+    std::unordered_map<PyFrameObject*, std::vector<FrameAction> > frameToActions;
 
     std::unordered_map<PyObject*, std::set<int> > codeObjectToExpressionLines;
 

typed_python/_types.cpp

@@ -3204,6 +3204,12 @@ PyObject *getTypePointer(PyObject* nullValue, PyObject* args) {
 
     return PyLong_FromLong((uint64_t)type);
 }
+PyObject* _temporaryReferenceTracerActive(PyObject* null, PyObject* args, PyObject* kwargs) {
+    return incref(
+        PyTemporaryReferenceTracer::globalTracer.frameToActions.size() ?
+        Py_True : Py_False
+    );
+}
 
 PyObject* gilReleaseThreadLoop(PyObject* null, PyObject* args, PyObject* kwargs) {
     PyEnsureGilReleased releaseTheGil;
@@ -3283,6 +3289,7 @@ static PyMethodDef module_methods[] = {
     {"couldConvertObjectToTypeAtLevel", (PyCFunction)couldConvertObjectToTypeAtLevel, METH_VARARGS | METH_KEYWORDS, NULL},
     {"isValidArithmeticUpcast", (PyCFunction)isValidArithmeticUpcast, METH_VARARGS | METH_KEYWORDS, NULL},
     {"isValidArithmeticConversion", (PyCFunction)isValidArithmeticConversion, METH_VARARGS | METH_KEYWORDS, NULL},
+    {"_temporaryReferenceTracerActive", (PyCFunction)_temporaryReferenceTracerActive, METH_VARARGS | METH_KEYWORDS, NULL},
     {"gilReleaseThreadLoop", (PyCFunction)gilReleaseThreadLoop, METH_VARARGS | METH_KEYWORDS, NULL},
     {"setModuleDict", (PyCFunction)setModuleDict, METH_VARARGS | METH_KEYWORDS, NULL},
     {NULL, NULL}

typed_python/compiler/python_ast_analysis.py

@@ -200,6 +200,33 @@ def visit(x):
     return closureVars
 
 
+def extractLineNumbersWithStatements(astNode):
+    res = set()
+
+    def visit(x):
+        if isinstance(x, Statement):
+            if x.matches.FunctionDef or x.matches.ClassDef or x.matches.AsyncFunctionDef:
+                return False
+
+            res.add(x.line_number)
+
+        if isinstance(x, Expr):
+            if (
+                x.matches.Lambda
+                or x.matches.ListComp
+                or x.matches.SetComp
+                or x.matches.DictComp
+                or x.matches.GeneratorExp
+            ):
+                return False
+
+        return True
+
+    visitPyAstChildren(astNode, visit)
+
+    return res
+
+
 def extractFunctionDefsInOrder(astNode):
     res = []
 

typed_python/compiler/tests/held_class_compilation_test.py

@@ -1,6 +1,7 @@
 import unittest
 import time
-
+import sys
+from typed_python import _types
 from typed_python import Class, Final, ListOf, Held, Member, Entrypoint, Forward
 
 
@@ -88,6 +89,8 @@ def move(c):
         self.checkCompiler(lambda c: c.h1.typeOfSelf(), c)
 
     def test_compile_list_of_held_class(self):
+        assert not _types._temporaryReferenceTracerActive()
+
         class H(Class, Final):
             x = Member(int, nonempty=True)
             y = Member(float, nonempty=True)
@@ -123,6 +126,14 @@ def incrementAllRange(l):
 
         incrementAllRange(aList)
 
+        def testIt(x):
+            assert _types._temporaryReferenceTracerActive()
+            return x
+
+        testIt(aList[0].increment)
+
+        assert not _types._temporaryReferenceTracerActive()
+
         self.assertEqual(getitem(aList, 0).x, 1)
         self.assertEqual(getitem(aList, 5).x, 1)
 
@@ -138,6 +149,8 @@ def incrementViaIterator(l):
 
         incrementViaIterator(aList)
 
+        assert not _types._temporaryReferenceTracerActive()
+
         self.assertEqual(getitem(aList, 0).x, 2)
         self.assertEqual(getitem(aList, 5).x, 2)
 

typed_python/compiler/tests/held_class_interpreter_semantics_test.py

@@ -50,12 +50,20 @@ def runTest():
 
     def test_construct_and_call_method_multiline(self):
         def runTest():
+            res = H(x=10, y=20).addToX(
+                # this is deliberately on another line
+                10
+            )
+            return res
+
+        def runTest2():
             return H(x=10, y=20).addToX(
                 # this is deliberately on another line
                 10
             )
 
         assert runTest() == 20
+        assert runTest2() == 20
 
     def test_can_assign_to_held_class_in_list(self):
         def runTest():

typed_python/internals.py

@@ -627,15 +627,15 @@ def Held(T):
 def extractCodeObjectNewStatementLineNumbers(codeObject):
     """Return the subset of the line numbers on which codeObject has new statements."""
     from typed_python.python_ast import convertFunctionToAlgebraicPyAst
+    from typed_python.compiler.python_ast_analysis import extractLineNumbersWithStatements
 
     try:
         ast = convertFunctionToAlgebraicPyAst(codeObject)
 
         res = set()
 
         if ast.matches.FunctionDef:
-            for statement in ast.body:
-                res.add(statement.line_number)
+            res = extractLineNumbersWithStatements(ast.body)
 
         return res
     except Exception:

typed_python/python_ast.py

@@ -786,7 +786,7 @@ def ExpressionStr(self):
         "args": TupleOf(Arg),
         "vararg": OneOf(Arg, None),
         "kwonlyargs": TupleOf(Arg),
-        "kw_defaults": TupleOf(Expr),
+        "kw_defaults": TupleOf(OneOf(None, Expr)),
         "kwarg": OneOf(Arg, None),
         "defaults": TupleOf(Expr),
     },