[ORC] Add ObjectTransformLayer

[JosephTremoulet]

Summary

This is a utility for clients that want to insert a layer that modifies
each ObjectFile and then passes it along to the next layer.

[lhames]

This looks great. Thanks very much Joseph!

include/llvm/ExecutionEngine/Orc/ObjectTransformLayer.h

+//===- ObjectTransformLayer.h - Run all objects through functor -*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Run all objects passed in through a user supplied functor.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_OBJECTTRANSFORMLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_OBJECTTRANSFORMLAYER_H
+
+#include "JITSymbol.h"
+
+namespace llvm {
+namespace orc {
+
+/// @brief Object mutating layer.
+///
+///   This layer accepts sets of ObjectFiles (via addObjectSet). It
+/// immediately applies the user supplied functor to each object, then adds
+/// the set of transformed objects to the layer below.
+template <typename BaseLayerT, typename TransformFtor>
+class ObjectTransformLayer {
+public:
+  /// @brief Handle to a set of added objects.
+  typedef typename BaseLayerT::ObjSetHandleT ObjSetHandleT;
+
+  /// @brief Construct an ObjectTransformLayer with the given BaseLayer
+  ObjectTransformLayer(BaseLayerT &BaseLayer,
+                       TransformFtor Transform = TransformFtor())
+      : BaseLayer(BaseLayer), Transform(std::move(Transform)) {}
+
+  /// @brief Apply the transform functor to each object in the object set, then
+  ///        add the resulting set of objects to the base layer, along with the
+  ///        memory manager and symbol resolver.
+  ///
+  /// @return A handle for the added objects.
+  template <typename ObjSetT, typename MemoryManagerPtrT,
+            typename SymbolResolverPtrT>
+  ObjSetHandleT addObjectSet(ObjSetT &Objects, MemoryManagerPtrT MemMgr,
+                             SymbolResolverPtrT Resolver) {
+
+    for (auto I = Objects.begin(), E = Objects.end(); I != E; ++I)
+      *I = Transform(std::move(*I));
+
+    return BaseLayer.addObjectSet(Objects, std::move(MemMgr),
+                                  std::move(Resolver));
+  }
+
+  /// @brief Remove the object set associated with the handle H.
+  void removeObjectSet(ObjSetHandleT H) { BaseLayer.removeObjectSet(H); }
+
+  /// @brief Search for the given named symbol.
+  /// @param Name The name of the symbol to search for.
+  /// @param ExportedSymbolsOnly If true, search only for exported symbols.
+  /// @return A handle for the given named symbol, if it exists.
+  JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
+    return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
+  }
+
+  /// @brief Get the address of the given symbol in the context of the set of
+  ///        objects represented by the handle H. This call is forwarded to the
+  ///        base layer's implementation.
+  /// @param H The handle for the object set to search in.
+  /// @param Name The name of the symbol to search for.
+  /// @param ExportedSymbolsOnly If true, search only for exported symbols.
+  /// @return A handle for the given named symbol, if it is found in the
+  ///         given object set.
+  JITSymbol findSymbolIn(ObjSetHandleT H, const std::string &Name,
+                         bool ExportedSymbolsOnly) {
+    return BaseLayer.findSymbolIn(H, Name, ExportedSymbolsOnly);
+  }
+
+  /// @brief Immediately emit and finalize the object set represented by the
+  ///        given handle.
+  /// @param H Handle for object set to emit/finalize.
+  void emitAndFinalize(ObjSetHandleT H) { BaseLayer.emitAndFinalize(H); }
+
+  /// @brief Map section addresses for the objects associated with the handle H.
+  void mapSectionAddress(ObjSetHandleT H, const void *LocalAddress,
+                         TargetAddress TargetAddr) {
+    BaseLayer.mapSectionAddress(H, LocalAddress, TargetAddr);
+  }
+
+  // Ownership hack.
+  // FIXME: Remove this as soon as RuntimeDyldELF can apply relocations without
+  //        referencing the original object.
+  template <typename OwningMBSet>
+  void takeOwnershipOfBuffers(ObjSetHandleT H, OwningMBSet MBs) {
+    BaseLayer.takeOwnershipOfBuffers(H, std::move(MBs));
+  }
+
+  /// @brief Access the transform functor directly.
+  TransformFtor &getTransform() { return Transform; }
+
+  /// @brief Access the mumate functor directly.
+  const TransformFtor &getTransform() const { return Transform; }
+
+private:
+  BaseLayerT &BaseLayer;
+  TransformFtor Transform;
+};
+
+} // End namespace orc.
+} // End namespace llvm.
+
+#endif // LLVM_EXECUTIONENGINE_ORC_OBJECTTRANSFORMLAYER_H

unittests/ExecutionEngine/Orc/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Core
   OrcJIT
   Support
   )
 
 add_llvm_unittest(OrcJITTests
   IndirectionUtilsTest.cpp
   LazyEmittingLayerTest.cpp
+  ObjectTransformLayerTest.cpp
   OrcTestCommon.cpp
   )

unittests/ExecutionEngine/Orc/ObjectTransformLayerTest.cpp

+//===- ObjectTransformLayerTest.cpp - Unit tests for ObjectTransformLayer -===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
+#include "llvm/ADT/SmallVector.h"
+#include "gtest/gtest.h"
+
+using namespace llvm::orc;
+
+namespace {
+
+// Stand-in for RuntimeDyld::MemoryManager
+typedef int MockMemoryManager;
+
+// Stand-in for RuntimeDyld::SymbolResolver
+typedef int MockSymbolResolver;
+
+// stand-in for object::ObjectFile
+typedef int MockObjectFile;
+
+// stand-in for llvm::MemoryBuffer set
+typedef int MockMemoryBufferSet;
+
+// Mock transform that operates on unique pointers to object files, and
+// allocats new object files rather than mutating the given ones.
+struct AllocatingTransform {
+  std::unique_ptr<MockObjectFile>
+  operator()(std::unique_ptr<MockObjectFile> Obj) const {
+    return std::make_unique<MockObjectFile>(*Obj + 1);
+  }
+};
+
+// Mock base layer for verifying behavior of transform layer.
+// Each method "T foo(args)" is accompanied by two auxiliary methods:
+//  - "void expectFoo(args)", to be called before calling foo on the transform
+//      layer; saves values of args, which mock layer foo then verifies against.
+// - "void verifyFoo(T)", to be called after foo, which verifies that the
+//      transform layer called the base layer and forwarded any return value.
+class MockBaseLayer {
+public:
+  typedef int ObjSetHandleT;
+
+  MockBaseLayer() : MockSymbol(nullptr) { resetExpectations(); }
+
+  template <typename ObjSetT, typename MemoryManagerPtrT,
+            typename SymbolResolverPtrT>
+  ObjSetHandleT addObjectSet(ObjSetT &Objects, MemoryManagerPtrT MemMgr,
+                             SymbolResolverPtrT Resolver) {
+    EXPECT_EQ(MockManager, *MemMgr) << "MM should pass through";
+    EXPECT_EQ(MockResolver, *Resolver) << "Resolver should pass through";
+    int I = 0;
+    for (auto &ObjPtr : Objects) {
+      EXPECT_EQ(MockObjects[I++] + 1, *ObjPtr) << "Transform should be applied";
+    }
+    EXPECT_EQ(MockObjects.size(), I) << "Number of objects should match";
+    LastCalled = "addObjectSet";
+    MockObjSetHandle = 111;
+    return MockObjSetHandle;
+  }
+  template <typename ObjSetT>
+  void expectAddObjectSet(ObjSetT &Objects, MockMemoryManager *MemMgr,
+                          MockSymbolResolver *Resolver) {
+    MockManager = *MemMgr;
+    MockResolver = *Resolver;
+    for (auto &ObjPtr : Objects) {
+      MockObjects.push_back(*ObjPtr);
+    }
+  }
+  void verifyAddObjectSet(ObjSetHandleT Returned) {
+    EXPECT_EQ("addObjectSet", LastCalled);
+    EXPECT_EQ(MockObjSetHandle, Returned) << "Return should pass through";
+    resetExpectations();
+  }
+
+  void removeObjectSet(ObjSetHandleT H) {
+    EXPECT_EQ(MockObjSetHandle, H);
+    LastCalled = "removeObjectSet";
+  }
+  void expectRemoveObjectSet(ObjSetHandleT H) { MockObjSetHandle = H; }
+  void verifyRemoveObjectSet() {
+    EXPECT_EQ("removeObjectSet", LastCalled);
+    resetExpectations();
+  }
+
+  JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
+    EXPECT_EQ(MockName, Name) << "Name should pass through";
+    EXPECT_EQ(MockBool, ExportedSymbolsOnly) << "Flag should pass through";
+    LastCalled = "findSymbol";
+    MockSymbol = JITSymbol(122, llvm::JITSymbolFlags::None);
+    return MockSymbol;
+  }
+  void expectFindSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
+    MockName = Name;
+    MockBool = ExportedSymbolsOnly;
+  }
+  void verifyFindSymbol(llvm::orc::JITSymbol Returned) {
+    EXPECT_EQ("findSymbol", LastCalled);
+    EXPECT_EQ(MockSymbol.getAddress(), Returned.getAddress())
+        << "Return should pass through";
+    resetExpectations();
+  }
+
+  JITSymbol findSymbolIn(ObjSetHandleT H, const std::string &Name,
+                         bool ExportedSymbolsOnly) {
+    EXPECT_EQ(MockObjSetHandle, H) << "Handle should pass through";
+    EXPECT_EQ(MockName, Name) << "Name should pass through";
+    EXPECT_EQ(MockBool, ExportedSymbolsOnly) << "Flag should pass through";
+    LastCalled = "findSymbolIn";
+    MockSymbol = JITSymbol(122, llvm::JITSymbolFlags::None);
+    return MockSymbol;
+  }
+  void expectFindSymbolIn(ObjSetHandleT H, const std::string &Name,
+                          bool ExportedSymbolsOnly) {
+    MockObjSetHandle = H;
+    MockName = Name;
+    MockBool = ExportedSymbolsOnly;
+  }
+  void verifyFindSymbolIn(llvm::orc::JITSymbol Returned) {
+    EXPECT_EQ("findSymbolIn", LastCalled);
+    EXPECT_EQ(MockSymbol.getAddress(), Returned.getAddress())
+        << "Return should pass through";
+    resetExpectations();
+  }
+
+  void emitAndFinalize(ObjSetHandleT H) {
+    EXPECT_EQ(MockObjSetHandle, H) << "Handle should pass through";
+    LastCalled = "emitAndFinalize";
+  }
+  void expectEmitAndFinalize(ObjSetHandleT H) { MockObjSetHandle = H; }
+  void verifyEmitAndFinalize() {
+    EXPECT_EQ("emitAndFinalize", LastCalled);
+    resetExpectations();
+  }
+
+  void mapSectionAddress(ObjSetHandleT H, const void *LocalAddress,
+                         TargetAddress TargetAddr) {
+    EXPECT_EQ(MockObjSetHandle, H);
+    EXPECT_EQ(MockLocalAddress, LocalAddress);
+    EXPECT_EQ(MockTargetAddress, TargetAddr);
+    LastCalled = "mapSectionAddress";
+  }
+  void expectMapSectionAddress(ObjSetHandleT H, const void *LocalAddress,
+                               TargetAddress TargetAddr) {
+    MockObjSetHandle = H;
+    MockLocalAddress = LocalAddress;
+    MockTargetAddress = TargetAddr;
+  }
+  void verifyMapSectionAddress() {
+    EXPECT_EQ("mapSectionAddress", LastCalled);
+    resetExpectations();
+  }
+
+  template <typename OwningMBSet>
+  void takeOwnershipOfBuffers(ObjSetHandleT H, OwningMBSet MBs) {
+    EXPECT_EQ(MockObjSetHandle, H);
+    EXPECT_EQ(MockBufferSet, *MBs);
+    LastCalled = "takeOwnershipOfBuffers";
+  }
+  void expectTakeOwnershipOfBuffers(ObjSetHandleT H, MockMemoryBufferSet *MBs) {
+    MockObjSetHandle = H;
+    MockBufferSet = *MBs;
+  }
+  void verifyTakeOwnershipOfBuffers() {
+    EXPECT_EQ("takeOwnershipOfBuffers", LastCalled);
+    resetExpectations();
+  }
+
+private:
+  // Backing fields for remembering parameter/return values
+  std::string LastCalled;
+  MockMemoryManager MockManager;
+  MockSymbolResolver MockResolver;
+  std::vector<MockObjectFile> MockObjects;
+  ObjSetHandleT MockObjSetHandle;
+  std::string MockName;
+  bool MockBool;
+  JITSymbol MockSymbol;
+  const void *MockLocalAddress;
+  TargetAddress MockTargetAddress;
+  MockMemoryBufferSet MockBufferSet;
+
+  // Clear remembered parameters between calls
+  void resetExpectations() {
+    LastCalled = "nothing";
+    MockManager = 0;
+    MockResolver = 0;
+    MockObjects.clear();
+    MockObjSetHandle = 0;
+    MockName = "bogus";
+    MockSymbol = JITSymbol(nullptr);
+    MockLocalAddress = nullptr;
+    MockTargetAddress = 0;
+    MockBufferSet = 0;
+  }
+};
+
+// Test each operation on ObjectTransformLayer.
+TEST(ObjectTransformLayerTest, Main) {
+  MockBaseLayer M;
+
+  // Create one object transform layer using a transform (as a functor)
+  // that allocates new objects, and deals in unique pointers.
+  ObjectTransformLayer<MockBaseLayer, AllocatingTransform> T1(M);
+
+  // Create a second object transform layer using a transform (as a lambda)
+  // that mutates objects in place, and deals in naked pointers
+  ObjectTransformLayer<MockBaseLayer,
+                       std::function<MockObjectFile *(MockObjectFile *)>>
+  T2(M, [](MockObjectFile *Obj) {
+    ++(*Obj);
+    return Obj;
+  });
+
+  // Instantiate some mock objects to use below
+  MockObjectFile MockObject1 = 211;
+  MockObjectFile MockObject2 = 222;
+  MockMemoryManager MockManager = 233;
+  MockSymbolResolver MockResolver = 244;
+
+  // Test addObjectSet with T1 (allocating, unique pointers)
+  std::vector<std::unique_ptr<MockObjectFile>> Objs1;
+  Objs1.push_back(std::make_unique<MockObjectFile>(MockObject1));
+  Objs1.push_back(std::make_unique<MockObjectFile>(MockObject2));
+  auto MM = std::make_unique<MockMemoryManager>(MockManager);
+  auto SR = std::make_unique<MockSymbolResolver>(MockResolver);
+  M.expectAddObjectSet(Objs1, MM.get(), SR.get());
+  auto H = T1.addObjectSet(Objs1, std::move(MM), std::move(SR));
+  M.verifyAddObjectSet(H);
+
+  // Test addObjectSet with T2 (mutating, naked pointers)
+  llvm::SmallVector<MockObjectFile *, 2> Objs2;
+  Objs2.push_back(&MockObject1);
+  Objs2.push_back(&MockObject2);
+  M.expectAddObjectSet(Objs2, &MockManager, &MockResolver);
+  H = T2.addObjectSet(Objs2, &MockManager, &MockResolver);
+  M.verifyAddObjectSet(H);
+  EXPECT_EQ(212, MockObject1) << "Expected mutation";
+  EXPECT_EQ(223, MockObject2) << "Expected mutation";
+
+  // Test removeObjectSet
+  M.expectRemoveObjectSet(H);
+  T1.removeObjectSet(H);
+  M.verifyRemoveObjectSet();
+
+  // Test findSymbol
+  std::string Name = "foo";
+  bool ExportedOnly = true;
+  M.expectFindSymbol(Name, ExportedOnly);
+  JITSymbol Symbol = T2.findSymbol(Name, ExportedOnly);
+  M.verifyFindSymbol(Symbol);
+
+  // Test findSymbolIn
+  Name = "bar";
+  ExportedOnly = false;
+  M.expectFindSymbolIn(H, Name, ExportedOnly);
+  Symbol = T1.findSymbolIn(H, Name, ExportedOnly);
+  M.verifyFindSymbolIn(Symbol);
+
+  // Test emitAndFinalize
+  M.expectEmitAndFinalize(H);
+  T2.emitAndFinalize(H);
+  M.verifyEmitAndFinalize();
+
+  // Test mapSectionAddress
+  char Buffer[24];
+  TargetAddress MockAddress = 255;
+  M.expectMapSectionAddress(H, Buffer, MockAddress);
+  T1.mapSectionAddress(H, Buffer, MockAddress);
+  M.verifyMapSectionAddress();
+
+  // Test takeOwnershipOfBuffers, using unique pointer to buffer set
+  auto MockBufferSetPtr = std::make_unique<MockMemoryBufferSet>(366);
+  M.expectTakeOwnershipOfBuffers(H, MockBufferSetPtr.get());
+  T2.takeOwnershipOfBuffers(H, std::move(MockBufferSetPtr));
+  M.verifyTakeOwnershipOfBuffers();
+
+  // Test takeOwnershipOfBuffers, using naked pointer to buffer set
+  MockMemoryBufferSet MockBufferSet = 266;
+  M.expectTakeOwnershipOfBuffers(H, &MockBufferSet);
+  T1.takeOwnershipOfBuffers(H, &MockBufferSet);
+  M.verifyTakeOwnershipOfBuffers();
+
+  // Verify transform getter (non-const)
+  MockObjectFile Mutatee = 277;
+  MockObjectFile *Out = T2.getTransform()(&Mutatee);
+  EXPECT_EQ(&Mutatee, Out) << "Expected in-place transform";
+  EXPECT_EQ(278, Mutatee) << "Expected incrementing transform";
+
+  // Verify transform getter (const)
+  auto OwnedObj = std::make_unique<MockObjectFile>(288);
+  const auto &T1C = T1;
+  OwnedObj = T1C.getTransform()(std::move(OwnedObj));
+  EXPECT_EQ(289, *OwnedObj) << "Expected incrementing transform";
+}
+}