Recover from failed cells and add cppjit.undo

docs/ReleaseNotes.md

@@ -16,7 +16,9 @@ This release supports Python 3.12-3.14 and LLVM 21-22.
 
 ## New features
 
-- ...
+- A failed `cppdef` or `cppexec` leaves the interpreter as it was, also
+  when it fails to link or its initializer throws. `cppjit.undo()` withdraws
+  the last successful `cppdef` or `cppexec`. See [Undo](Undo.md).
 
 ## Deprecated features
 

docs/Undo.md

@@ -0,0 +1,55 @@
+# Undo and failed inputs
+
+In a notebook, a cell that fails to compile must not force a restart.
+cppjit gives two guarantees:
+
+- A failed `cppdef` or `cppexec` leaves the interpreter as it was before
+  the call. Fix the code and run it again with the same names.
+- `cppjit.undo()` withdraws the last successful `cppdef` or `cppexec`. Call
+  it again to go back further.
+
+```python
+import cppjit
+
+cppjit.cppdef("int answer() { return 41; }")
+cppjit.undo()
+cppjit.cppdef("int answer() { return 42; }")
+assert cppjit.gbl.answer() == 42
+```
+
+`undo()` also drops the Python proxies that cppjit cached since that call,
+so a name that you define again resolves to the new declaration. An object
+that you made from a withdrawn class is not usable after the undo.
+
+Use `cppjit.undo()`, not `cppjit.gbl.Cpp.Undo`. A raw `Cpp::Undo` counts
+interpreter inputs, and cppjit adds inputs of its own between your cells.
+cppjit routes `gbl.Cpp.Undo` around its JIT call wrapper, so that the
+wrapper cannot withdraw itself.
+
+## What is covered
+
+The tests in `test/test_failed_input.py` show each case. `TestFAILED_INPUT`
+covers failed inputs: parse and Sema errors, link failures, templates first
+used by a failed input, an initializer that throws, and macros. `TestUNDO`
+covers undo: variables, functions, classes, namespaces (new, reopened and
+nested), redeclarations, enumerators, macros, include guards and
+`#pragma once`.
+
+Most of the rollback is in CppInterOp. See its
+[Undo and failed inputs](https://github.com/compiler-research/CppInterOp/blob/main/docs/UndoAndFailedInputs.rst)
+page for the mechanism and the upstream LLVM work.
+
+## What is not covered
+
+These tests are marked as expected failures, with the upstream work in the
+reason:
+
+| Case | Why |
+|---|---|
+| A CodeGen error, for example a duplicate mangled name, can crash the session | State in clang CodeGen; llvm#226285, #234493, #204152, #223305 |
+| A function that fails to link on its first call stays broken after you define the missing symbol | ORC keeps the error; undo the cell and define the function again |
+| Undo of a class definition after a forward declaration in an earlier cell | Private definition data in clang; fixed by llvm#218149 in LLVM 23.1.3 |
+| A header that changed on disk after undo | clang caches file contents; llvm#182750 |
+
+Side effects that already ran, such as memory, files and output, are not
+undone.

python/cppjit/__init__.py

@@ -35,6 +35,7 @@
 
 __all__ = [
     "cppdef",  # declare C++ source to Cling
+    "undo",  # undo the last successful cppdef or cppexec
     "cppexec",  # execute a C++ statement
     "macro",  # attempt to evaluate a cpp macro
     "include",  # load and jit a header file
@@ -218,15 +219,65 @@ def __exit__(self, tp, val, trace):
             self.err = _end_capture_stderr()
 
 
+# Checkpoint per successful cppdef/cppexec: (input count, scope cache size).
+_cells = []
+
+
+def _begin_cell():
+    return (_backend._input_count(), _backend._scope_cache_size())
+
+
+def _end_cell(mark, ok):
+    # CppInterOp withdraws a failed input itself.
+    if ok:
+        _cells.append(mark)
+
+
+def _abort_cell(mark):
+    # A C++ exception unwinds past CppInterOp's own withdraw.
+    _backend._undo_inputs_to(mark[0])
+    _backend._drop_scope_cache_since(mark[1])
+
+
+def undo():
+    """Undo the last successful cppdef or cppexec, with the JIT code made since."""
+    if not _cells:
+        raise RuntimeError("nothing to undo")
+    count, cached = _cells.pop()
+    if not _backend._undo_inputs_to(count):
+        raise RuntimeError("the interpreter refused to undo")
+    # Cached proxies may point into the withdrawn code.
+    _backend._drop_scope_cache_since(cached)
+
+
+def _cpp_undo(N=1):
+    """Cpp::Undo, called from C so it cannot withdraw its own JIT call wrapper."""
+    count = _backend._input_count()
+    if N > count or not _backend._undo_inputs_to(count - N):
+        return 1
+    return 0
+
+
+gbl.Cpp.Undo = _cpp_undo
+
+
 def cppdef(src, verbose=True):
     """Declare C++ source <src> to Cling."""
+    mark = _begin_cell()
     with _stderr_capture() as err:
-        errcode = gbl.Cpp.Declare(src, not verbose)
+        try:
+            errcode = gbl.Cpp.Declare(src, not verbose)
+        except Exception:
+            _abort_cell(mark)
+            raise
     if not errcode == 0 or err.err:
         if "warning" in err.err.lower() and "error" not in err.err.lower():
+            _end_cell(mark, True)
             warnings.warn(err.err, SyntaxWarning)
             return True
+        _end_cell(mark, False)
         raise SyntaxError("Failed to parse the given C++ code%s" % err.err)
+    _end_cell(mark, True)
     return True
 
 
@@ -237,18 +288,22 @@ def cppexec(stmt):
 
     # capture stderr, but note that Process could legitimately be writing to
     # std::cerr, in which case the captured output needs to be printed as normal
+    mark = _begin_cell()
     with _stderr_capture() as err:
         errcode = ctypes.c_int(0)
         try:
             errcode = gbl.Cpp.Process(stmt)
         except Exception as e:
+            _abort_cell(mark)
             sys.stderr.write("%s\n\n" % str(e))
             if not errcode.value:
                 errcode.value = 1
 
     if not errcode == 0:
+        _end_cell(mark, False)
         raise SyntaxError("Failed to parse the given C++ code%s" % err.err)
-    elif err.err and err.err[1:] != "\n":
+    _end_cell(mark, True)
+    if err.err and err.err[1:] != "\n":
         sys.stderr.write(err.err[1:])
 
     return True

src/cpyrt/CPPScope.cxx

@@ -27,6 +27,43 @@ namespace cppjit::cpyrt {
 
 extern PyTypeObject CPPInstance_Type;
 
+//- undo support -------------------------------------------------------------
+// (weak ref to the scope or metaclass, attribute name), oldest first
+static std::vector<std::pair<PyObject*, PyObject*>> gScopeCacheLog;
+
+void LogScopeCache(PyObject* scope, PyObject* name) {
+  PyObject* ref = PyWeakref_NewRef(scope, nullptr);
+  if (!ref) {
+    PyErr_Clear();
+    return;
+  }
+  Py_INCREF(name);
+  gScopeCacheLog.emplace_back(ref, name);
+}
+
+size_t ScopeCacheLogSize() { return gScopeCacheLog.size(); }
+
+void DropScopeCacheSince(size_t size) {
+  while (gScopeCacheLog.size() > size) {
+    auto [ref, name] = gScopeCacheLog.back();
+    gScopeCacheLog.pop_back();
+    if (PyObject* scope = cpyrt_GetWeakRef(ref)) {
+      if (PyType_Type.tp_setattro(scope, name, nullptr) < 0)
+        PyErr_Clear();
+      Py_DECREF(scope);
+    }
+    Py_DECREF(ref);
+    Py_DECREF(name);
+  }
+}
+
+static int CacheOnScope(PyObject* scope, PyObject* name, PyObject* attr) {
+  int rc = PyType_Type.tp_setattro(scope, name, attr);
+  if (rc == 0)
+    LogScopeCache(scope, name);
+  return rc;
+}
+
 //- helpers ------------------------------------------------------------------
 static inline PyObject*
 add_template(PyObject* pyclass, const std::string& name,
@@ -42,7 +79,7 @@ add_template(PyObject* pyclass, const std::string& name,
     PyErr_Clear();
     pytmpl = TemplateProxy_New(ncl, ncl, pyclass);
     // cache the template on its clean name
-    PyType_Type.tp_setattro(pyclass, pyncl, (PyObject*)pytmpl);
+    CacheOnScope(pyclass, pyncl, (PyObject*)pytmpl);
     Py_DECREF(pyncl);
   } else if (!TemplateProxy_CheckExact((PyObject*)pytmpl)) {
     Py_DECREF(pytmpl);
@@ -460,15 +497,15 @@ static PyObject* meta_getattro(PyObject* pyclass, PyObject* pyname) {
     if (attr) {
       // cache the result
       if (CPPDataMember_Check(attr)) {
-        PyType_Type.tp_setattro((PyObject*)Py_TYPE(pyclass), pyname, attr);
+        CacheOnScope((PyObject*)Py_TYPE(pyclass), pyname, attr);
 
         Py_DECREF(attr);
         // The call below goes through "dm_get"
         attr = PyType_Type.tp_getattro(pyclass, pyname);
         if (!attr && PyErr_Occurred())
           Utility::FetchError(errors);
       } else
-        PyType_Type.tp_setattro(pyclass, pyname, attr);
+        CacheOnScope(pyclass, pyname, attr);
 
     } else {
       Utility::FetchError(errors);
@@ -529,7 +566,7 @@ static PyObject* meta_getattro(PyObject* pyclass, PyObject* pyname) {
       Py_DECREF(possibly_shadowed);
     } else {
       attr = possibly_shadowed;
-      PyType_Type.tp_setattro(pyclass, pyname, attr);
+      CacheOnScope(pyclass, pyname, attr);
     }
   }
 

src/cpyrt/CPPScope.h

@@ -112,6 +112,12 @@ inline CPPScope* CPPScopeMeta_New(interop::TCppScope_t klass, PyObject* args) {
   return pymeta;
 }
 
+//- undo support -------------------------------------------------------------
+// Lookups cached on scope proxies, so that undo can drop outdated ones.
+void LogScopeCache(PyObject* scope, PyObject* name);
+size_t ScopeCacheLogSize();
+void DropScopeCacheSince(size_t size);
+
 } // namespace cppjit::cpyrt
 
 #endif // !CPYRT_CPPSCOPE_H

src/cpyrt/ProxyWrappers.cxx

@@ -122,9 +122,10 @@ static inline void AddPropertyToClass(PyObject* pyclass,
 static inline void AddScopeToParent(PyObject* parent, const std::string& name,
                                     PyObject* newscope) {
   PyObject* pyname = PyUnicode_InternFromString((char*)name.c_str());
-  if (CPPScope_Check(parent))
-    PyType_Type.tp_setattro(parent, pyname, newscope);
-  else
+  if (CPPScope_Check(parent)) {
+    if (PyType_Type.tp_setattro(parent, pyname, newscope) == 0)
+      LogScopeCache(parent, pyname);
+  } else
     PyObject_SetAttr(parent, pyname, newscope);
   Py_DECREF(pyname);
 }

src/cpyrt/TemplateProxy.cxx

@@ -8,6 +8,7 @@ using namespace cppjit;
 #include "CPPInstance.h"
 #include "CPPMethod.h"
 #include "CPPOverload.h"
+#include "CPPScope.h"
 #include "PyCallable.h"
 #include "PyStrings.h"
 #include "TemplateProxy.h"
@@ -289,7 +290,8 @@ PyObject* TemplateProxy::Instantiate(const std::string& fname,
       }
 
       // add to class dictionary
-      PyType_Type.tp_setattro(fTI->fPyClass, pycachename, pyol);
+      if (PyType_Type.tp_setattro(fTI->fPyClass, pycachename, pyol) == 0)
+        LogScopeCache(fTI->fPyClass, pycachename);
     }
 
     // Case 3/4: pre-existing method that was either not found b/c the full
@@ -317,7 +319,8 @@ PyObject* TemplateProxy::Instantiate(const std::string& fname,
 
     // Special Case if name was aliased (e.g. typedef in template instantiation)
     if (!exact && !bExactMatch) {
-      PyType_Type.tp_setattro(fTI->fPyClass, pyresname, pyol);
+      if (PyType_Type.tp_setattro(fTI->fPyClass, pyresname, pyol) == 0)
+        LogScopeCache(fTI->fPyClass, pyresname);
     }
 
     // cleanup

src/cpyrt/cpyrtModule.cxx

@@ -894,6 +894,43 @@ static PyObject* AddSmartPtrType(PyObject*, PyObject* args) {
   Py_RETURN_NONE;
 }
 
+//----------------------------------------------------------------------------
+static PyObject* InputCount(PyObject*, PyObject*) {
+  return PyLong_FromSize_t(Cpp::GetInputCount());
+}
+
+//----------------------------------------------------------------------------
+// Not through a JIT wrapper: undo would withdraw the wrapper.
+static PyObject* UndoInputsTo(PyObject*, PyObject* pykeep) {
+  size_t keep = PyLong_AsSize_t(pykeep);
+  if (keep == (size_t)-1 && PyErr_Occurred())
+    return nullptr;
+  size_t count = Cpp::GetInputCount();
+  if (keep > count) {
+    PyErr_SetString(PyExc_ValueError, "checkpoint is ahead of the interpreter");
+    return nullptr;
+  }
+  if (keep == count)
+    Py_RETURN_TRUE;
+  if (Cpp::Undo((unsigned)(count - keep)) != 0)
+    Py_RETURN_FALSE;
+  Py_RETURN_TRUE;
+}
+
+//----------------------------------------------------------------------------
+static PyObject* ScopeCacheSize(PyObject*, PyObject*) {
+  return PyLong_FromSize_t(ScopeCacheLogSize());
+}
+
+//----------------------------------------------------------------------------
+static PyObject* DropScopeCache(PyObject*, PyObject* pysize) {
+  size_t size = PyLong_AsSize_t(pysize);
+  if (size == (size_t)-1 && PyErr_Occurred())
+    return nullptr;
+  DropScopeCacheSince(size);
+  Py_RETURN_NONE;
+}
+
 //----------------------------------------------------------------------------
 static PyObject* BeginCaptureStderr(PyObject*, PyObject*) {
   gOldErrorBuffer = std::cerr.rdbuf();
@@ -969,6 +1006,14 @@ static PyMethodDef gcpyrtMethods[] = {
      (char*)"Modify held C++ object ownership."},
     {(char*)"AddSmartPtrType", (PyCFunction)AddSmartPtrType, METH_VARARGS,
      (char*)"Add a smart pointer to the list of known smart pointer types."},
+    {(char*)"_input_count", (PyCFunction)InputCount, METH_NOARGS,
+     (char*)"Number of inputs the interpreter holds."},
+    {(char*)"_undo_inputs_to", (PyCFunction)UndoInputsTo, METH_O,
+     (char*)"Undo the inputs after the given input count."},
+    {(char*)"_scope_cache_size", (PyCFunction)ScopeCacheSize, METH_NOARGS,
+     (char*)"Number of lookups cached on scope proxies."},
+    {(char*)"_drop_scope_cache_since", (PyCFunction)DropScopeCache, METH_O,
+     (char*)"Drop the lookups cached after the given count."},
     {(char*)"_begin_capture_stderr", (PyCFunction)BeginCaptureStderr,
      METH_NOARGS, (char*)"Begin capturing stderr to a in memory buffer."},
     {(char*)"_end_capture_stderr", (PyCFunction)EndCaptureStderr, METH_NOARGS,