[clang-repl] Implement code undo.

In interactive C++ it is convenient to roll back to a previous state of the
compiler. For example:
clang-repl> int x = 42;
clang-repl> %undo
clang-repl> float x = 24 // not an error

To support this, the patch extends the functionality used to recover from
errors and adds functionality to recover the low-level execution infrastructure.

The current implementation is based on watermarks. It exploits the fact that
at each incremental input the underlying compiler infrastructure is in a valid
state. We can only go N incremental inputs back to a previous valid state. We do
not need and do not do any further dependency tracking.

This patch was co-developed with V. Vassilev, relies on the past work of Purva
Chaudhari in clang-repl and is inspired by the past work on the same feature
in the Cling interpreter.

Co-authored-by: Purva-Chaudhari <purva.chaudhari02@gmail.com>
Co-authored-by: Vassil Vassilev <v.g.vassilev@gmail.com>
Signed-off-by: Jun Zhang <jun@junz.org>

clang/include/clang/Interpreter/Interpreter.h

@@ -69,6 +69,9 @@ public:
     return llvm::Error::success();
   }
 
+  /// Undo N previous incremental inputs.
+  llvm::Error Undo(unsigned N = 1);
+
   /// \returns the \c JITTargetAddress of a \c GlobalDecl. This interface uses
   /// the CodeGenModule's internal mangling cache to avoid recomputing the
   /// mangled name.

clang/lib/Interpreter/IncrementalExecutor.cpp

@@ -12,6 +12,7 @@
 
 #include "IncrementalExecutor.h"
 
+#include "clang/Interpreter/PartialTranslationUnit.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
@@ -58,8 +59,24 @@ llvm::Error IncrementalExecutor::cleanUp() {
   return Jit->deinitialize(Jit->getMainJITDylib());
 }
 
-llvm::Error IncrementalExecutor::addModule(std::unique_ptr<llvm::Module> M) {
-  return Jit->addIRModule(llvm::orc::ThreadSafeModule(std::move(M), TSCtx));
+llvm::Error IncrementalExecutor::addModule(PartialTranslationUnit &PTU) {
+  llvm::orc::ResourceTrackerSP RT =
+      Jit->getMainJITDylib().createResourceTracker();
+  ResourceTrackers[&PTU] = RT;
+
+  return Jit->addIRModule(RT, {std::move(PTU.TheModule), TSCtx});
+}
+
+llvm::Error IncrementalExecutor::removeModule(PartialTranslationUnit &PTU) {
+
+  llvm::orc::ResourceTrackerSP RT = std::move(ResourceTrackers[&PTU]);
+  if (!RT)
+    return llvm::Error::success();
+
+  ResourceTrackers.erase(&PTU);
+  if (llvm::Error Err = RT->remove())
+    return Err;
+  return llvm::Error::success();
 }
 
 llvm::Error IncrementalExecutor::runCtors() const {

clang/lib/Interpreter/IncrementalExecutor.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_CLANG_LIB_INTERPRETER_INCREMENTALEXECUTOR_H
 #define LLVM_CLANG_LIB_INTERPRETER_INCREMENTALEXECUTOR_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
@@ -29,11 +30,17 @@ class ThreadSafeContext;
 } // namespace llvm
 
 namespace clang {
+
+struct PartialTranslationUnit;
+
 class IncrementalExecutor {
   using CtorDtorIterator = llvm::orc::CtorDtorIterator;
   std::unique_ptr<llvm::orc::LLJIT> Jit;
   llvm::orc::ThreadSafeContext &TSCtx;
 
+  llvm::DenseMap<const PartialTranslationUnit *, llvm::orc::ResourceTrackerSP>
+      ResourceTrackers;
+
 public:
   enum SymbolNameKind { IRName, LinkerName };
 
@@ -41,7 +48,8 @@ public:
                       const llvm::Triple &Triple);
   ~IncrementalExecutor();
 
-  llvm::Error addModule(std::unique_ptr<llvm::Module> M);
+  llvm::Error addModule(PartialTranslationUnit &PTU);
+  llvm::Error removeModule(PartialTranslationUnit &PTU);
   llvm::Error runCtors() const;
   llvm::Error cleanUp();
   llvm::Expected<llvm::JITTargetAddress>

clang/lib/Interpreter/IncrementalParser.cpp

@@ -181,27 +181,9 @@ IncrementalParser::ParseOrWrapTopLevelDecl() {
 
   DiagnosticsEngine &Diags = getCI()->getDiagnostics();
   if (Diags.hasErrorOccurred()) {
-    TranslationUnitDecl *MostRecentTU = C.getTranslationUnitDecl();
-    TranslationUnitDecl *PreviousTU = MostRecentTU->getPreviousDecl();
-    assert(PreviousTU && "Must have a TU from the ASTContext initialization!");
-    TranslationUnitDecl *FirstTU = MostRecentTU->getFirstDecl();
-    assert(FirstTU);
-    FirstTU->RedeclLink.setLatest(PreviousTU);
-    C.TUDecl = PreviousTU;
-    S.TUScope->setEntity(PreviousTU);
-
-    // Clean up the lookup table
-    if (StoredDeclsMap *Map = PreviousTU->getPrimaryContext()->getLookupPtr()) {
-      for (auto I = Map->begin(); I != Map->end(); ++I) {
-        StoredDeclsList &List = I->second;
-        DeclContextLookupResult R = List.getLookupResult();
-        for (NamedDecl *D : R)
-          if (D->getTranslationUnitDecl() == MostRecentTU)
-            List.remove(D);
-        if (List.isNull())
-          Map->erase(I);
-      }
-    }
+    PartialTranslationUnit MostRecentPTU = {C.getTranslationUnitDecl(),
+                                            nullptr};
+    CleanUpPTU(MostRecentPTU);
 
     Diags.Reset(/*soft=*/true);
     Diags.getClient()->clear();
@@ -296,6 +278,24 @@ IncrementalParser::Parse(llvm::StringRef input) {
   return PTU;
 }
 
+void IncrementalParser::CleanUpPTU(PartialTranslationUnit &PTU) {
+  TranslationUnitDecl *MostRecentTU = PTU.TUPart;
+  TranslationUnitDecl *FirstTU = MostRecentTU->getFirstDecl();
+  if (StoredDeclsMap *Map = FirstTU->getPrimaryContext()->getLookupPtr()) {
+    for (auto I = Map->begin(); I != Map->end(); ++I) {
+      StoredDeclsList &List = I->second;
+      DeclContextLookupResult R = List.getLookupResult();
+      for (NamedDecl *D : R) {
+        if (D->getTranslationUnitDecl() == MostRecentTU) {
+          List.remove(D);
+        }
+      }
+      if (List.isNull())
+        Map->erase(I);
+    }
+  }
+}
+
 llvm::StringRef IncrementalParser::GetMangledName(GlobalDecl GD) const {
   CodeGenerator *CG = getCodeGen(Act.get());
   assert(CG);

clang/lib/Interpreter/IncrementalParser.h

@@ -72,6 +72,10 @@ public:
   ///\returns the mangled name of a \c GD.
   llvm::StringRef GetMangledName(GlobalDecl GD) const;
 
+  void CleanUpPTU(PartialTranslationUnit &PTU);
+
+  std::list<PartialTranslationUnit> &getPTUs() { return PTUs; }
+
 private:
   llvm::Expected<PartialTranslationUnit &> ParseOrWrapTopLevelDecl();
 };

clang/lib/Interpreter/Interpreter.cpp

@@ -229,7 +229,7 @@ llvm::Error Interpreter::Execute(PartialTranslationUnit &T) {
       return Err;
   }
   // FIXME: Add a callback to retain the llvm::Module once the JIT is done.
-  if (auto Err = IncrExecutor->addModule(std::move(T.TheModule)))
+  if (auto Err = IncrExecutor->addModule(T))
     return Err;
 
   if (auto Err = IncrExecutor->runCtors())
@@ -267,3 +267,22 @@ Interpreter::getSymbolAddressFromLinkerName(llvm::StringRef Name) const {
 
   return IncrExecutor->getSymbolAddress(Name, IncrementalExecutor::LinkerName);
 }
+
+llvm::Error Interpreter::Undo(unsigned N) {
+
+  std::list<PartialTranslationUnit> &PTUs = IncrParser->getPTUs();
+  if (N > PTUs.size())
+    return llvm::make_error<llvm::StringError>("Operation failed. "
+                                               "Too many undos",
+                                               std::error_code());
+  for (unsigned I = 0; I < N; I++) {
+    if (IncrExecutor) {
+      if (llvm::Error Err = IncrExecutor->removeModule(PTUs.back()))
+        return Err;
+    }
+
+    IncrParser->CleanUpPTU(PTUs.back());
+    PTUs.pop_back();
+  }
+  return llvm::Error::success();
+}

clang/test/Interpreter/code-undo.cpp

@@ -0,0 +1,23 @@
+// RUN: clang-repl "int i = 10;" 'extern "C" int printf(const char*,...);' \
+// RUN:            'auto r1 = printf("i = %d\n", i);' | FileCheck --check-prefix=CHECK-DRIVER %s
+// REQUIRES: host-supports-jit
+// UNSUPPORTED: system-aix
+// CHECK-DRIVER: i = 10
+// RUN: cat %s | clang-repl | FileCheck %s
+extern "C" int printf(const char *, ...);
+int x1 = 0;
+int x2 = 42;
+%undo
+int x2 = 24;
+auto r1 = printf("x1 = %d\n", x1);
+// CHECK: x1 = 0
+auto r2 = printf("x2 = %d\n", x2);
+// CHECK-NEXT: x2 = 24
+
+int foo() { return 1; }
+%undo
+int foo() { return 2; }
+auto r3 = printf("foo() = %d\n", foo());
+// CHECK-NEXT: foo() = 2
+
+%quit

clang/test/Interpreter/execute.cpp

@@ -22,4 +22,4 @@ int r3 = foo();
 struct D { float f = 1.0; D *m = nullptr; D(){} ~D() { printf("D[f=%f, m=0x%llx]\n", f, reinterpret_cast<unsigned long long>(m)); }} d;
 // CHECK: D[f=1.000000, m=0x0]
 
-quit
+%quit

clang/test/Interpreter/plugins.cpp

@@ -6,8 +6,7 @@
 int i = 10;
 extern "C" int printf(const char*,...);
 auto r1 = printf("i = %d\n", i);
-quit
-
+%quit
 
 // CHECK: top-level-decl: "i"
 // CHECK-NEXT: top-level-decl: "r1"

clang/test/Interpreter/sanity.c

@@ -15,4 +15,4 @@ void TestFunc() { ++TestVar; }
 // CHECK-NEXT:     UnaryOperator{{.*}} 'int' lvalue prefix '++'
 // CHECK-NEXT:       DeclRefExpr{{.*}} 'int' lvalue Var [[var_ptr]] 'TestVar' 'int'
 
-quit
+%quit

clang/tools/clang-repl/ClangRepl.cpp

@@ -111,8 +111,14 @@ int main(int argc, const char **argv) {
     llvm::LineEditor LE("clang-repl");
     // FIXME: Add LE.setListCompleter
     while (llvm::Optional<std::string> Line = LE.readLine()) {
-      if (*Line == "quit")
+      if (*Line == R"(%quit)")
         break;
+      if (*Line == R"(%undo)") {
+        if (auto Err = Interp->Undo())
+          llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "error: ");
+        continue;
+      }
+
       if (auto Err = Interp->ParseAndExecute(*Line))
         llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "error: ");
     }

clang/unittests/Interpreter/InterpreterTest.cpp

@@ -128,6 +128,51 @@ TEST(InterpreterTest, DeclsAndStatements) {
   EXPECT_EQ("Parsing failed.", llvm::toString(std::move(Err)));
 }
 
+TEST(InterpreterTest, UndoCommand) {
+  Args ExtraArgs = {"-Xclang", "-diagnostic-log-file", "-Xclang", "-"};
+
+  // Create the diagnostic engine with unowned consumer.
+  std::string DiagnosticOutput;
+  llvm::raw_string_ostream DiagnosticsOS(DiagnosticOutput);
+  auto DiagPrinter = std::make_unique<TextDiagnosticPrinter>(
+      DiagnosticsOS, new DiagnosticOptions());
+
+  auto Interp = createInterpreter(ExtraArgs, DiagPrinter.get());
+
+  // Fail to undo.
+  auto Err1 = Interp->Undo();
+  EXPECT_EQ("Operation failed. Too many undos",
+            llvm::toString(std::move(Err1)));
+  auto Err2 = Interp->Parse("int foo = 42;");
+  EXPECT_TRUE(!!Err2);
+  auto Err3 = Interp->Undo(2);
+  EXPECT_EQ("Operation failed. Too many undos",
+            llvm::toString(std::move(Err3)));
+
+  // Succeed to undo.
+  auto Err4 = Interp->Parse("int x = 42;");
+  EXPECT_TRUE(!!Err4);
+  auto Err5 = Interp->Undo();
+  EXPECT_FALSE(Err5);
+  auto Err6 = Interp->Parse("int x = 24;");
+  EXPECT_TRUE(!!Err6);
+  auto Err7 = Interp->Parse("#define X 42");
+  EXPECT_TRUE(!!Err7);
+  auto Err8 = Interp->Undo();
+  EXPECT_FALSE(Err8);
+  auto Err9 = Interp->Parse("#define X 24");
+  EXPECT_TRUE(!!Err9);
+
+  // Undo input contains errors.
+  auto Err10 = Interp->Parse("int y = ;");
+  EXPECT_FALSE(!!Err10);
+  EXPECT_EQ("Parsing failed.", llvm::toString(Err10.takeError()));
+  auto Err11 = Interp->Parse("int y = 42;");
+  EXPECT_TRUE(!!Err11);
+  auto Err12 = Interp->Undo();
+  EXPECT_FALSE(Err12);
+}
+
 static std::string MangleName(NamedDecl *ND) {
   ASTContext &C = ND->getASTContext();
   std::unique_ptr<MangleContext> MangleC(C.createMangleContext());