Adding doWithCleanup abstraction

[beanz]

Summary

This patch adds a new doWithCleanup abstraction to LLVM that works
with llvm::Error, int, bool, and <insert your own type here>
return types that propagate status.

This abstraction is a more general way of having a conditional cleanup
that can be used with multiple return paths out of a function.

Test cases and a first use case in clang are part of the patch.

[dblaikie]

clang/lib/Frontend/FrontendAction.cpp

553–565

I guess this would be written as:

bool Success = BeginSourceFileImpl(CI, RealInput, HasBegunSourceFile))
if (!Success) {
  if (HasBegunSourceFile)
    CI.getDiagnosticClient().EndSourceFile();
  CI.clearOutputFiles(/*EraseFiles=*/true);
  CI.getLangOpts().setCompilingModule(LangOptions::CMK_None);
  setCurrentInput(FrontendInputFile());
  setCompilerInstance(nullptr);
}
return Success;

Which seems pretty legible - makes me unsure it's worth the complexity of building something like doWithCleanup, not that it's heinously complicated by any means - but all the traits stuff, etc.

& ideally/in general this'd usually be dealt with with some kind of RAII ownership model rather than such a strong stateful transition of an existing object that needs to be undone, rather than destroying an object.

llvm/include/llvm/Support/Cleanup.h

56–59

If this is going to be a template anyway, might as well use raw/generic/templated functor arguments and avoid the dynamic dispatch overhead of std::function? (or at least use llvm::function_ref)

llvm/unittests/Support/CleanupTests.cpp

19

Presumably doWithCleanup doesn't need to be called with an explicit template type argument most of the time/ever? It's probably good to test it in that way too, without passing the explicit template type argument.

clang/include/clang/Frontend/FrontendAction.h

class FrontendAction {
   CompilerInstance *Instance;
   friend class ASTMergeAction;
   friend class WrapperFrontendAction;
 
 private:
   std::unique_ptr<ASTConsumer> CreateWrappedASTConsumer(CompilerInstance &CI,
                                                         StringRef InFile);
 
+  bool BeginSourceFileImpl(CompilerInstance &CI, const FrontendInputFile &Input,
+                           bool &HasBegunSourceFile);
+
 protected:
   /// @name Implementation Action Interface
   /// @{
 
   /// Prepare to execute the action on the given CompilerInstance.
   ///
   /// This is called before executing the action on any inputs, and can modify
   /// the configuration as needed (including adjusting the input list).

clang/lib/Frontend/FrontendAction.cpp

 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/LiteralSupport.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Parse/ParseAST.h"
 #include "clang/Serialization/ASTDeserializationListener.h"
 #include "clang/Serialization/ASTReader.h"
 #include "clang/Serialization/GlobalModuleIndex.h"
-#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/BuryPointer.h"
+#include "llvm/Support/Cleanup.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/raw_ostream.h"
 #include <system_error>
 using namespace clang;
 
getInputBufferForModule(CompilerInstance &CI, Module *M) {
   }
 
   return llvm::MemoryBuffer::getMemBufferCopy(
       HeaderContents, Module::getModuleInputBufferName());
 }
 
 bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
                                      const FrontendInputFile &RealInput) {
+  bool HasBegunSourceFile = false;
+  std::function<bool()> Operation = [&]() {
+    return BeginSourceFileImpl(CI, RealInput, HasBegunSourceFile);
+  };
+  std::function<void()> Cleanup = [&]() {
+    if (HasBegunSourceFile)
+      CI.getDiagnosticClient().EndSourceFile();
+    CI.clearOutputFiles(/*EraseFiles=*/true);
+    CI.getLangOpts().setCompilingModule(LangOptions::CMK_None);
+    setCurrentInput(FrontendInputFile());
+    setCompilerInstance(nullptr);
+  };
+  return llvm::doWithCleanup(Operation, Cleanup);
+}
+
+bool FrontendAction::BeginSourceFileImpl(CompilerInstance &CI,
+                                         const FrontendInputFile &RealInput,
+                                         bool &HasBegunSourceFile) {
   FrontendInputFile Input(RealInput);
   assert(!Instance && "Already processing a source file!");
   assert(!Input.isEmpty() && "Unexpected empty filename!");
   setCurrentInput(Input);
   setCompilerInstance(&CI);
 
-  bool HasBegunSourceFile = false;
   bool ReplayASTFile = Input.getKind().getFormat() == InputKind::Precompiled &&
                        usesPreprocessorOnly();
 
-  // If we fail, reset state since the client will not end up calling the
-  // matching EndSourceFile(). All paths that return true should release this.
-  auto FailureCleanup = llvm::make_scope_exit([&]() {
-    if (HasBegunSourceFile)
-      CI.getDiagnosticClient().EndSourceFile();
-    CI.clearOutputFiles(/*EraseFiles=*/true);
-    CI.getLangOpts().setCompilingModule(LangOptions::CMK_None);
-    setCurrentInput(FrontendInputFile());
-    setCompilerInstance(nullptr);
-  });
-
   if (!BeginInvocation(CI))
     return false;
 
   // If we're replaying the build of an AST file, import it and set up
   // the initial state from its build.
   if (ReplayASTFile) {
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags(&CI.getDiagnostics());
 
if (Input.getKind().getFormat() == InputKind::Precompiled) {
     if (!BeginSourceFileAction(CI))
       return false;
 
     // Create the AST consumer.
     CI.setASTConsumer(CreateWrappedASTConsumer(CI, InputFile));
     if (!CI.hasASTConsumer())
       return false;
 
-    FailureCleanup.release();
     return true;
   }
 
   // Set up the file and source managers, if needed.
   if (!CI.hasFileManager()) {
     if (!CI.createFileManager()) {
       return false;
     }
if (Input.getKind().getLanguage() == Language::LLVM_IR) {
     // Initialize the action.
     if (!BeginSourceFileAction(CI))
       return false;
 
     // Initialize the main file entry.
     if (!CI.InitializeSourceManager(CurrentInput))
       return false;
 
-    FailureCleanup.release();
     return true;
   }
 
   // If the implicit PCH include is actually a directory, rather than
   // a single file, search for a suitable PCH file in that directory.
   if (!CI.getPreprocessorOpts().ImplicitPCHInclude.empty()) {
     FileManager &FileMgr = CI.getFileManager();
     PreprocessorOptions &PPOpts = CI.getPreprocessorOpts();
bool FrontendAction::BeginSourceFileImpl(CompilerInstance &CI,
   if (!CI.getFrontendOpts().OverrideRecordLayoutsFile.empty() &&
       CI.hasASTContext() && !CI.getASTContext().getExternalSource()) {
     IntrusiveRefCntPtr<ExternalASTSource>
       Override(new LayoutOverrideSource(
                      CI.getFrontendOpts().OverrideRecordLayoutsFile));
     CI.getASTContext().setExternalSource(Override);
   }
 
-  FailureCleanup.release();
   return true;
 }
 
 llvm::Error FrontendAction::Execute() {
   CompilerInstance &CI = getCompilerInstance();
 
   if (CI.hasFrontendTimer()) {
     llvm::TimeRegion Timer(CI.getFrontendTimer());

llvm/include/llvm/Support/Cleanup.h

+//===- llvm/Support/Cleanup.h - Cleanup on failure abstraction --*- 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 defines helper abstractions for cleaning up resources on failure.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_CLEANUP_H
+#define LLVM_SUPPORT_CLEANUP_H
+
+#include "llvm/Support/Error.h"
+
+#include <functional>
+#include <type_traits>
+
+namespace llvm {
+
+/// Helper templates for determining if a result is a failure.
+namespace detail {
+
+/// Evaluator for if an llvm::Error contains a failure.
+template <typename ErrT>
+typename std::enable_if<std::is_same<ErrT, llvm::Error>::value, bool>::type
+isFailure(ErrT &Err) {
+  return (bool)Err;
+}
+
+/// Evaluator for c-style int error returns where success is 0.
+template <typename ErrT>
+typename std::enable_if<std::is_same<ErrT, int>::value, bool>::type
+isFailure(ErrT &Err) {
+  return 0 != Err;
+}
+
+/// Evaluator for boolean error returns where success is true.
+template <typename ErrT>
+typename std::enable_if<!(std::is_same<ErrT, int>::value ||
+                          std::is_same<ErrT, llvm::Error>::value),
+                        bool>::type
+isFailure(ErrT &Err) {
+  return !(bool)Err;
+}
+} // namespace detail
+
+/// doWithCleanup abstraction performs an operation and cleanup on failure.
+///
+/// @param Fn std::function to perform which returns the error type.
+/// @param CleanupFn std::function to perform cleanup on failure result.
+/// @param IsFailure (optional) std::function that takes a reference of the
+/// error type, and returns true if the value is an error and false otherwise.
+template <typename ErrT>
+ErrT doWithCleanup(
+    std::function<ErrT()> Fn, std::function<void()> CleanupFn,
+    std::function<bool(ErrT &)> IsFailure = &detail::isFailure<ErrT>) {
+  ErrT Result = Fn();
+  if (IsFailure(Result))
+    CleanupFn();
+  return Result;
+}
+} // namespace llvm
+
+#endif // LLVM_SUPPORT_CLEANUP_H

llvm/unittests/Support/CMakeLists.txt

add_llvm_unittest(SupportTests
   BinaryStreamTest.cpp
   BlockFrequencyTest.cpp
   BranchProbabilityTest.cpp
   CachePruningTest.cpp
   CrashRecoveryTest.cpp
   Casting.cpp
   CheckedArithmeticTest.cpp
   Chrono.cpp
+  CleanupTests.cpp
   CommandLineTest.cpp
   CompressionTest.cpp
   ConvertUTFTest.cpp
   CRCTest.cpp
   DataExtractorTest.cpp
   DebugTest.cpp
   DebugCounterTest.cpp
   DJBTest.cpp

llvm/unittests/Support/CleanupTests.cpp

+//===----- unittests/CleanupTests.cpp - Cleanup.h tests -------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Cleanup.h"
+#include "llvm/Support/Error.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+TEST(Error, DoWithCleanupBool) {
+  bool CleanupRun = false;
+  bool Result =
+      doWithCleanup<bool>([] { return true; }, [&] { CleanupRun = true; });
+  EXPECT_TRUE(Result);
+  EXPECT_FALSE(CleanupRun);
+
+  Result =
+      doWithCleanup<bool>([] { return false; }, [&] { CleanupRun = true; });
+  EXPECT_FALSE(Result);
+  EXPECT_TRUE(CleanupRun);
+}
+
+TEST(Error, DoWithCleanupError) {
+  bool CleanupRun = false;
+  llvm::consumeError(doWithCleanup<llvm::Error>(
+      [&] { return Error::success(); }, [&] { CleanupRun = true; }));
+
+  EXPECT_FALSE(CleanupRun);
+
+  llvm::consumeError(doWithCleanup<llvm::Error>(
+      [&] { return errorCodeToError(inconvertibleErrorCode()); },
+      [&] { CleanupRun = true; }));
+  EXPECT_TRUE(CleanupRun);
+}
+
+TEST(Error, DoWithCleanupInt) {
+  bool CleanupRun = false;
+  int Result = doWithCleanup<int>([] { return 0; }, [&] { CleanupRun = true; });
+  EXPECT_EQ(Result, 0);
+  EXPECT_FALSE(CleanupRun);
+
+  Result = doWithCleanup<int>([] { return 1; }, [&] { CleanupRun = true; });
+  EXPECT_EQ(Result, 1);
+  EXPECT_TRUE(CleanupRun);
+}
+
+enum class PetType { Doggo, Kitteh, Birb, WhoozieWhatsit };
+
+TEST(Error, DoWithCleanupOverrideFailure) {
+  bool CleanupRun = false;
+  PetType Result = doWithCleanup<PetType>(
+      [] { return PetType::Doggo; }, [&] { CleanupRun = true; },
+      [](PetType &E) { return E == PetType::WhoozieWhatsit; });
+  EXPECT_EQ(Result, PetType::Doggo);
+  EXPECT_FALSE(CleanupRun);
+
+  Result = doWithCleanup<PetType>(
+      [] { return PetType::WhoozieWhatsit; }, [&] { CleanupRun = true; },
+      [](PetType &E) { return E == PetType::WhoozieWhatsit; });
+  EXPECT_EQ(Result, PetType::WhoozieWhatsit);
+  EXPECT_TRUE(CleanupRun);
+}
+
+enum class ResultTy {
+  Failure = 0,
+  Success,
+};
+
+TEST(Error, DoWithCleanupEnumNonZeroFailure) {
+  bool CleanupRun = false;
+  ResultTy Result = doWithCleanup<ResultTy>([] { return ResultTy::Success; },
+                                            [&] { CleanupRun = true; });
+  EXPECT_EQ(Result, ResultTy::Success);
+  EXPECT_FALSE(CleanupRun);
+
+  Result = doWithCleanup<ResultTy>([] { return ResultTy::Failure; },
+                                   [&] { CleanupRun = true; });
+  EXPECT_EQ(Result, ResultTy::Failure);
+  EXPECT_TRUE(CleanupRun);
+}