[clang-repl] Implement code undo

[junaire]

Summary

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 <[e-mail]>
Co-authored-by: Vassil Vassilev <[e-mail]>
Signed-off-by: Jun Zhang <[e-mail]>

[junaire]

Also remove PTU in PTU list

[junaire]

Improve code a little bit

[junaire]

Remove a NFC change to make review easier

[junaire]

• Support undo N times
• Add a comment

[junaire]

Pass PTU by reference as it can't be NULL, this makes our code safer.

[junaire]

Fix the build

[junaire]

• Add prefix for meta commands
• Add more tests

[junaire]

this depends on https://reviews.llvm.org/D126684

[v.g.vassilev]

In D126682#3546865, @junaire wrote:

this depends on https://reviews.llvm.org/D126684

I would prefer to merge them into one.

[junaire]

Merge D126684 into this patch.

Consider D126684 is almost @Purva-Chaudhari 's work, I'll list her as
co-author when I commit this.

[junaire]

Restore PTU in Interpreter::Undo, this should fix the broken tests.

[junaire]

Allow undo command failed in the edge cases like undoed too many times.

[v.g.vassilev]

clang/lib/Interpreter/IncrementalParser.cpp

183

We should call Restore here.

clang/unittests/Interpreter/InterpreterTest.cpp

295

Can we move this as a regular test called code-undo.cpp?

[junaire]

• split the tests
• add % to all meta commands

[junaire]

clang/lib/Interpreter/IncrementalParser.cpp

183

Though there are many logic duplicates with Restore, I'm not sure if we can refactor it. We have done more here than Restore, and Restore need to pass a PTU, but we only have TranslationDecl here.

[v.g.vassilev]

clang/lib/Interpreter/IncrementalParser.cpp

183

I think we can/should be able to construct the PartialTranslationUnit with only a TranslationUnitDecl having TheModule as nullptr.

[junaire]

Address Vassil's comments

[junaire]

Fix the unintentional format by clang-format

[junaire]

Fix the build

[v.g.vassilev]

I think we are in a good shape to broaden the list of reviewers.

[junaire]

• Add unittest
• Fix a crash when previous input failed to parse

[junaire]

Rebase

[v.g.vassilev]

Thanks for working on this. Here are some more comments.

clang/include/clang/Interpreter/Interpreter.h

72

I am not sure if that's the best function. We

72

Please ignore this comment it was superseded by another one and I forgot to delete it.

clang/lib/Interpreter/IncrementalParser.cpp

281

I think we can merge this function with Undo where we can conditionally check if we have a llvm::Module to recover.

clang/lib/Interpreter/Interpreter.cpp

273

I am not sure how feasible is this usecase but in -fsyntax-only mode we can still need undo. That means somebody created an Interpreter instance for frontend-only operations, such as lookups or template instantiations. In that case I think it also would make sense to support the undo operation.

278

In some cases we can have PTU's that do not generate code. For example, define A 1 will not produce a corresponding llvm::Module. I think here we should count the size of the PartialTranslationUnitDecls.

clang/test/Interpreter/code-undo.cpp

7

Can you add here something like int j = 0 and print its value after the undo. This way we will make sure we do not undo too much.

[junaire]

Address @v.g.vassilev 's comments.

[junaire]

clang/lib/Interpreter/IncrementalParser.cpp

281

I'm not sure I understand you correctly, do you mean to merge Interpreter::Restore with Interpreter::Undo? I think IncrementalParser::Restore is used above?

[junaire]

• Rename IncrementalParser::Restore to IncrementalParser::CleanupPTU
• Drop unused code

[junaire]

• Update a comment
• Add a FIXME back

[junaire]

Update commit message

[junaire]

Rename CleanupPTU => CleanUpPTU

[junaire]

Try to fix the failing permerge test.

[v.g.vassilev]

LGTM!

[junaire]

land in https://reviews.llvm.org/rGdea5a9cc929048be261a4c030407e4d7e1e70fec

clang/include/clang/Interpreter/Interpreter.h

llvm::Error ParseAndExecute(llvm::StringRef Code) {
     auto PTU = Parse(Code);
     if (!PTU)
       return PTU.takeError();
     if (PTU->TheModule)
       return Execute(*PTU);
     return llvm::Error::success();
   }
 
+  void Restore(PartialTranslationUnit &PTU);
+
+  /// Undo previous parse results for N times. It'll stop and report an error
+  /// if an error occurs.
+  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.
   llvm::Expected<llvm::JITTargetAddress> getSymbolAddress(GlobalDecl GD) const;
 
   /// \returns the \c JITTargetAddress of a given name as written in the IR.
   llvm::Expected<llvm::JITTargetAddress>
   getSymbolAddress(llvm::StringRef IRName) const;

clang/lib/Interpreter/IncrementalExecutor.h

 //===--- IncrementalExecutor.h - Incremental Execution ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the class which performs incremental code execution.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 #include <memory>
 
 namespace llvm {
 class Error;
 class Module;
 namespace orc {
 class LLJIT;
 class ThreadSafeContext;
 } // namespace orc
 } // 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 };
 
   IncrementalExecutor(llvm::orc::ThreadSafeContext &TSC, llvm::Error &Err,
                       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::Expected<llvm::JITTargetAddress>
   getSymbolAddress(llvm::StringRef Name, SymbolNameKind NameKind) const;
   llvm::orc::LLJIT *getExecutionEngine() const { return Jit.get(); }
+
+  std::size_t getAvailableModuleSize() const { return ResourceTrackers.size(); }
 };
 
 } // end namespace clang
 
 #endif // LLVM_CLANG_LIB_INTERPRETER_INCREMENTALEXECUTOR_H

clang/lib/Interpreter/IncrementalExecutor.cpp

 //===--- IncrementalExecutor.cpp - Incremental Execution --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the class which performs incremental code execution.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/Module.h"
IncrementalExecutor::IncrementalExecutor(llvm::orc::ThreadSafeContext &TSC,
   else {
     Err = PSGOrErr.takeError();
     return;
   }
 }
 
 IncrementalExecutor::~IncrementalExecutor() {}
 
-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 {
   return Jit->initialize(Jit->getMainJITDylib());
 }
 
 llvm::Expected<llvm::JITTargetAddress>
 IncrementalExecutor::getSymbolAddress(llvm::StringRef Name,

clang/lib/Interpreter/IncrementalParser.h

public:
   ///\returns a \c PartialTranslationUnit which holds information about the
   /// \c TranslationUnitDecl and \c llvm::Module corresponding to the input.
   llvm::Expected<PartialTranslationUnit &> Parse(llvm::StringRef Input);
 
   /// Uses the CodeGenModule mangled name cache and avoids recomputing.
   ///\returns the mangled name of a \c GD.
   llvm::StringRef GetMangledName(GlobalDecl GD) const;
 
+  void Restore(PartialTranslationUnit &PTU);
+
+  std::list<PartialTranslationUnit> &getPTUs() { return PTUs; }
+
 private:
   llvm::Expected<PartialTranslationUnit &> ParseOrWrapTopLevelDecl();
 };
 } // end namespace clang
 
 #endif // LLVM_CLANG_LIB_INTERPRETER_INCREMENTALPARSER_H

clang/lib/Interpreter/IncrementalParser.cpp

for (bool AtEOF = P->ParseFirstTopLevelDecl(ADecl, ImportState); !AtEOF;
     // skipping something.
     if (ADecl && !Consumer->HandleTopLevelDecl(ADecl.get()))
       return llvm::make_error<llvm::StringError>("Parsing failed. "
                                                  "The consumer rejected a decl",
                                                  std::error_code());
   }
 
   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};
+    Restore(MostRecentPTU);
 
     // FIXME: Do not reset the pragma handlers.
     Diags.Reset();
     return llvm::make_error<llvm::StringError>("Parsing failed.",
                                                std::error_code());
   }
 
   // Process any TopLevelDecls generated by #pragma weak.
CG->StartModule("incr_module_" + std::to_string(PTUs.size()),
                     M->getContext());
 
     PTU->TheModule = std::move(M);
   }
 
   return PTU;
 }
 
+void IncrementalParser::Restore(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);
   return CG->GetMangledName(GD);
 }
 
 } // end namespace clang

clang/lib/Interpreter/Interpreter.cpp

if (!IncrExecutor) {
     const llvm::Triple &Triple =
         getCompilerInstance()->getASTContext().getTargetInfo().getTriple();
     llvm::Error Err = llvm::Error::success();
     IncrExecutor = std::make_unique<IncrementalExecutor>(*TSCtx, Err, Triple);
 
     if (Err)
       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())
     return Err;
 
   return llvm::Error::success();
 }
 
+void Interpreter::Restore(PartialTranslationUnit &PTU) {
+  IncrParser->Restore(PTU);
+}
+
 llvm::Expected<llvm::JITTargetAddress>
 Interpreter::getSymbolAddress(GlobalDecl GD) const {
   if (!IncrExecutor)
     return llvm::make_error<llvm::StringError>("Operation failed. "
                                                "No execution engine",
                                                std::error_code());
   llvm::StringRef MangledName = IncrParser->GetMangledName(GD);
   return getSymbolAddress(MangledName);

 Interpreter::getSymbolAddressFromLinkerName(llvm::StringRef Name) const {
   if (!IncrExecutor)
     return llvm::make_error<llvm::StringError>("Operation failed. "
                                                "No execution engine",
                                                std::error_code());
 
   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;
+    }
+
+    Restore(PTUs.back());
+    PTUs.pop_back();
+  }
+  return llvm::Error::success();
+}

clang/test/Interpreter/code-undo.cpp

-// RUN: clang-repl "int x = 10;" "int y=7; err;" "int y = 10;"
 // 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 i = 42;
-auto r1 = printf("i = %d\n", i);
-// CHECK: i = 42
+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
 
-struct S { float f = 1.0; S *m = nullptr;} s;
+int foo() { return 1; }
+%undo
+int foo() { return 2; }
+auto r3 = printf("foo() = %d\n", foo());
+// CHECK-NEXT: foo() = 2
 
-auto r2 = printf("S[f=%f, m=0x%llx]\n", s.f, reinterpret_cast<unsigned long long>(s.m));
-// CHECK-NEXT: S[f=1.000000, m=0x0]
-
-inline int foo() { return 42; }
-int r3 = foo();
-
-quit
+%quit

clang/test/Interpreter/execute.cpp

 struct S { float f = 1.0; S *m = nullptr;} s;
 
 auto r2 = printf("S[f=%f, m=0x%llx]\n", s.f, reinterpret_cast<unsigned long long>(s.m));
 // CHECK-NEXT: S[f=1.000000, m=0x0]
 
 inline int foo() { return 42; }
 int r3 = foo();
 
-quit
+%quit

clang/test/Interpreter/plugins.cpp

 // RUN: cat %s | clang-repl -Xcc -Xclang -Xcc -load -Xcc -Xclang \
 // RUN:            -Xcc %llvmshlibdir/PrintFunctionNames%pluginext -Xcc -Xclang\
 // RUN:            -Xcc -add-plugin -Xcc -Xclang -Xcc print-fns 2>&1 | FileCheck %s
 // REQUIRES: host-supports-jit, plugins, examples
 
 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"
 // CHECK-NEXT: i = 10

clang/test/Interpreter/sanity.c

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

clang/tools/clang-repl/ClangRepl.cpp

for (const std::string &input : OptInputs) {
     if (auto Err = Interp->ParseAndExecute(input))
       llvm::logAllUnhandledErrors(std::move(Err), llvm::errs(), "error: ");
   }
 
   if (OptInputs.empty()) {
     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: ");
     }
   }
 
   // Our error handler depends on the Diagnostics object, which we're
   // potentially about to delete. Uninstall the handler now so that any
   // later errors use the default handling behavior instead.
   llvm::remove_fatal_error_handler();
 
   llvm::llvm_shutdown();
 
   return 0;
 }

clang/unittests/Interpreter/InterpreterTest.cpp

TEST(InterpreterTest, DeclsAndStatements) {
   EXPECT_FALSE(!!R2);
   using ::testing::HasSubstr;
   EXPECT_THAT(DiagnosticsOS.str(),
               HasSubstr("error: unknown type name 'var1'"));
   auto Err = R2.takeError();
   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);
+  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());
   std::string mangledName;
   llvm::raw_string_ostream RawStr(mangledName);
   MangleC->mangleName(ND, RawStr);
   return RawStr.str();
 }
#endif
   NamedDecl *TmpltSpec = llvm::cast<DeclRefExpr>(UO->getSubExpr())->getDecl();
 
   std::string MangledName = MangleName(TmpltSpec);
   typedef int (*TemplateSpecFn)(void *);
   auto fn = (TemplateSpecFn)cantFail(Interp->getSymbolAddress(MangledName));
   EXPECT_EQ(42, fn(NewA));
 }
 
 } // end anonymous namespace