[ORC] Add writePointers to ExecutorProcessControl's MemoryAccess

[sunho]

Summary

Add pointer write functionality to MemoryAccess that is needed for implementing redirection manager. It also refactors the code a bit by introducing InProcessMemoryAccess class.

[sunho]

Update

[lhames]

LGTM. Thanks @sunho!

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericMemoryAccess.h

 public:
   /// Function addresses for memory access.
   struct FuncAddrs {
     ExecutorAddr WriteUInt8s;
     ExecutorAddr WriteUInt16s;
     ExecutorAddr WriteUInt32s;
     ExecutorAddr WriteUInt64s;
     ExecutorAddr WriteBuffers;
+    ExecutorAddr WritePointers;
   };
 
   /// Create an EPCGenericMemoryAccess instance from a given set of
   /// function addrs.
   EPCGenericMemoryAccess(ExecutorProcessControl &EPC, FuncAddrs FAs)
       : EPC(EPC), FAs(FAs) {}
 
   void writeUInt8sAsync(ArrayRef<tpctypes::UInt8Write> Ws,
public:
 
   void writeBuffersAsync(ArrayRef<tpctypes::BufferWrite> Ws,
                          WriteResultFn OnWriteComplete) override {
     using namespace shared;
     EPC.callSPSWrapperAsync<void(SPSSequence<SPSMemoryAccessBufferWrite>)>(
         FAs.WriteBuffers, std::move(OnWriteComplete), Ws);
   }
 
+  void writePointersAsync(ArrayRef<tpctypes::PointerWrite> Ws,
+                          WriteResultFn OnWriteComplete) override {
+    using namespace shared;
+    EPC.callSPSWrapperAsync<void(SPSSequence<SPSMemoryAccessPointerWrite>)>(
+        FAs.WritePointers, std::move(OnWriteComplete), Ws);
+  }
+
 private:
   ExecutorProcessControl &EPC;
   FuncAddrs FAs;
 };
 
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_EPCGENERICMEMORYACCESS_H

llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h

virtual void writeUInt32sAsync(ArrayRef<tpctypes::UInt32Write> Ws,
                                    WriteResultFn OnWriteComplete) = 0;
 
     virtual void writeUInt64sAsync(ArrayRef<tpctypes::UInt64Write> Ws,
                                    WriteResultFn OnWriteComplete) = 0;
 
     virtual void writeBuffersAsync(ArrayRef<tpctypes::BufferWrite> Ws,
                                    WriteResultFn OnWriteComplete) = 0;
 
+    virtual void writePointersAsync(ArrayRef<tpctypes::PointerWrite> Ws,
+                                    WriteResultFn OnWriteComplete) = 0;
+
     Error writeUInt8s(ArrayRef<tpctypes::UInt8Write> Ws) {
       std::promise<MSVCPError> ResultP;
       auto ResultF = ResultP.get_future();
       writeUInt8sAsync(Ws,
                        [&](Error Err) { ResultP.set_value(std::move(Err)); });
       return ResultF.get();
     }
 
public:
 
     Error writeBuffers(ArrayRef<tpctypes::BufferWrite> Ws) {
       std::promise<MSVCPError> ResultP;
       auto ResultF = ResultP.get_future();
       writeBuffersAsync(Ws,
                         [&](Error Err) { ResultP.set_value(std::move(Err)); });
       return ResultF.get();
     }
+
+    Error writePointers(ArrayRef<tpctypes::PointerWrite> Ws) {
+      std::promise<MSVCPError> ResultP;
+      auto ResultF = ResultP.get_future();
+      writePointersAsync(Ws,
+                         [&](Error Err) { ResultP.set_value(std::move(Err)); });
+      return ResultF.get();
+    }
   };
 
   /// A pair of a dylib and a set of symbols to be looked up.
   struct LookupRequest {
     LookupRequest(tpctypes::DylibHandle Handle, const SymbolLookupSet &Symbols)
         : Handle(Handle), Symbols(Symbols) {}
     tpctypes::DylibHandle Handle;
     const SymbolLookupSet &Symbols;
protected:
   unsigned PageSize = 0;
   JITDispatchInfo JDI;
   MemoryAccess *MemAccess = nullptr;
   jitlink::JITLinkMemoryManager *MemMgr = nullptr;
   StringMap<std::vector<char>> BootstrapMap;
   StringMap<ExecutorAddr> BootstrapSymbols;
 };
 
+class InProcessMemoryAccess : public ExecutorProcessControl::MemoryAccess {
+public:
+  InProcessMemoryAccess(bool IsArch64Bit) : IsArch64Bit(IsArch64Bit) {}
+  void writeUInt8sAsync(ArrayRef<tpctypes::UInt8Write> Ws,
+                        WriteResultFn OnWriteComplete) override;
+
+  void writeUInt16sAsync(ArrayRef<tpctypes::UInt16Write> Ws,
+                         WriteResultFn OnWriteComplete) override;
+
+  void writeUInt32sAsync(ArrayRef<tpctypes::UInt32Write> Ws,
+                         WriteResultFn OnWriteComplete) override;
+
+  void writeUInt64sAsync(ArrayRef<tpctypes::UInt64Write> Ws,
+                         WriteResultFn OnWriteComplete) override;
+
+  void writeBuffersAsync(ArrayRef<tpctypes::BufferWrite> Ws,
+                         WriteResultFn OnWriteComplete) override;
+
+  void writePointersAsync(ArrayRef<tpctypes::PointerWrite> Ws,
+                          WriteResultFn OnWriteComplete) override;
+
+private:
+  bool IsArch64Bit;
+};
+
 /// A ExecutorProcessControl instance that asserts if any of its methods are
 /// used. Suitable for use is unit tests, and by ORC clients who haven't moved
 /// to ExecutorProcessControl-based APIs yet.
-class UnsupportedExecutorProcessControl : public ExecutorProcessControl {
+class UnsupportedExecutorProcessControl : public ExecutorProcessControl,
+                                          private InProcessMemoryAccess {
 public:
   UnsupportedExecutorProcessControl(
       std::shared_ptr<SymbolStringPool> SSP = nullptr,
-      std::unique_ptr<TaskDispatcher> D = nullptr,
-      const std::string &TT = "", unsigned PageSize = 0)
-      : ExecutorProcessControl(SSP ? std::move(SSP)
-                               : std::make_shared<SymbolStringPool>(),
-                               D ? std::move(D)
-                               : std::make_unique<InPlaceTaskDispatcher>()) {
+      std::unique_ptr<TaskDispatcher> D = nullptr, const std::string &TT = "",
+      unsigned PageSize = 0)
+      : ExecutorProcessControl(
+            SSP ? std::move(SSP) : std::make_shared<SymbolStringPool>(),
+            D ? std::move(D) : std::make_unique<InPlaceTaskDispatcher>()),
+        InProcessMemoryAccess(Triple(TT).isArch64Bit()) {
     this->TargetTriple = Triple(TT);
     this->PageSize = PageSize;
+    this->MemAccess = this;
   }
 
   Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override {
     llvm_unreachable("Unsupported");
   }
 
   Expected<std::vector<tpctypes::LookupResult>>
   lookupSymbols(ArrayRef<LookupRequest> Request) override {
void callWrapperAsync(ExecutorAddr WrapperFnAddr,
                         ArrayRef<char> ArgBuffer) override {
     llvm_unreachable("Unsupported");
   }
 
   Error disconnect() override { return Error::success(); }
 };
 
 /// A ExecutorProcessControl implementation targeting the current process.
-class SelfExecutorProcessControl
-    : public ExecutorProcessControl,
-      private ExecutorProcessControl::MemoryAccess {
+class SelfExecutorProcessControl : public ExecutorProcessControl,
+                                   private InProcessMemoryAccess {
 public:
   SelfExecutorProcessControl(
       std::shared_ptr<SymbolStringPool> SSP, std::unique_ptr<TaskDispatcher> D,
       Triple TargetTriple, unsigned PageSize,
       std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr);
 
   /// Create a SelfExecutorProcessControl with the given symbol string pool and
   /// memory manager.
public:
 
   void callWrapperAsync(ExecutorAddr WrapperFnAddr,
                         IncomingWFRHandler OnComplete,
                         ArrayRef<char> ArgBuffer) override;
 
   Error disconnect() override;
 
 private:
-  void writeUInt8sAsync(ArrayRef<tpctypes::UInt8Write> Ws,
-                        WriteResultFn OnWriteComplete) override;
-
-  void writeUInt16sAsync(ArrayRef<tpctypes::UInt16Write> Ws,
-                         WriteResultFn OnWriteComplete) override;
-
-  void writeUInt32sAsync(ArrayRef<tpctypes::UInt32Write> Ws,
-                         WriteResultFn OnWriteComplete) override;
-
-  void writeUInt64sAsync(ArrayRef<tpctypes::UInt64Write> Ws,
-                         WriteResultFn OnWriteComplete) override;
-
-  void writeBuffersAsync(ArrayRef<tpctypes::BufferWrite> Ws,
-                         WriteResultFn OnWriteComplete) override;
-
   static shared::CWrapperFunctionResult
   jitDispatchViaWrapperFunctionManager(void *Ctx, const void *FnTag,
                                        const char *Data, size_t Size);
 
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   char GlobalManglingPrefix = 0;
 };
 
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_EXECUTORPROCESSCONTROL_H

llvm/include/llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h

struct BufferWrite {
   BufferWrite() = default;
   BufferWrite(ExecutorAddr Addr, StringRef Buffer)
       : Addr(Addr), Buffer(Buffer) {}
 
   ExecutorAddr Addr;
   StringRef Buffer;
 };
 
+/// Describes a write to a pointer.
+/// For use with TargetProcessControl::MemoryAccess objects.
+struct PointerWrite {
+  PointerWrite() = default;
+  PointerWrite(ExecutorAddr Addr, ExecutorAddr Value)
+      : Addr(Addr), Value(Value) {}
+
+  ExecutorAddr Addr;
+  ExecutorAddr Value;
+};
+
 /// A handle used to represent a loaded dylib in the target process.
 using DylibHandle = ExecutorAddr;
 
 using LookupResult = std::vector<ExecutorAddr>;
 
 } // end namespace tpctypes
 
 namespace shared {

 using SPSMemoryAccessUIntWrite = SPSTuple<SPSExecutorAddr, T>;
 
 using SPSMemoryAccessUInt8Write = SPSMemoryAccessUIntWrite<uint8_t>;
 using SPSMemoryAccessUInt16Write = SPSMemoryAccessUIntWrite<uint16_t>;
 using SPSMemoryAccessUInt32Write = SPSMemoryAccessUIntWrite<uint32_t>;
 using SPSMemoryAccessUInt64Write = SPSMemoryAccessUIntWrite<uint64_t>;
 
 using SPSMemoryAccessBufferWrite = SPSTuple<SPSExecutorAddr, SPSSequence<char>>;
+using SPSMemoryAccessPointerWrite = SPSTuple<SPSExecutorAddr, SPSExecutorAddr>;
 
 template <>
 class SPSSerializationTraits<SPSRemoteAllocGroup, tpctypes::RemoteAllocGroup> {
   enum WireBits {
     ReadBit = 1 << 0,
     WriteBit = 1 << 1,
     ExecBit = 1 << 2,
     FinalizeBit = 1 << 3
public:
 
   static bool deserialize(SPSInputBuffer &IB, tpctypes::BufferWrite &W) {
     return SPSTuple<SPSExecutorAddr,
                     SPSSequence<char>>::AsArgList ::deserialize(IB, W.Addr,
                                                                 W.Buffer);
   }
 };
 
+template <>
+class SPSSerializationTraits<SPSMemoryAccessPointerWrite,
+                             tpctypes::PointerWrite> {
+public:
+  static size_t size(const tpctypes::PointerWrite &W) {
+    return SPSTuple<SPSExecutorAddr, SPSExecutorAddr>::AsArgList::size(W.Addr,
+                                                                       W.Value);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const tpctypes::PointerWrite &W) {
+    return SPSTuple<SPSExecutorAddr, SPSExecutorAddr>::AsArgList::serialize(
+        OB, W.Addr, W.Value);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, tpctypes::PointerWrite &W) {
+    return SPSTuple<SPSExecutorAddr, SPSExecutorAddr>::AsArgList::deserialize(
+        IB, W.Addr, W.Value);
+  }
+};
+
 } // end namespace shared
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_SHARED_TARGETPROCESSCONTROLTYPES_H

llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp

 ExecutorProcessControl::MemoryAccess::~MemoryAccess() = default;
 
 ExecutorProcessControl::~ExecutorProcessControl() = default;
 
 SelfExecutorProcessControl::SelfExecutorProcessControl(
     std::shared_ptr<SymbolStringPool> SSP, std::unique_ptr<TaskDispatcher> D,
     Triple TargetTriple, unsigned PageSize,
     std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr)
-    : ExecutorProcessControl(std::move(SSP), std::move(D)) {
+    : ExecutorProcessControl(std::move(SSP), std::move(D)),
+      InProcessMemoryAccess(TargetTriple.isArch64Bit()) {
 
   OwnedMemMgr = std::move(MemMgr);
   if (!OwnedMemMgr)
     OwnedMemMgr = std::make_unique<jitlink::InProcessMemoryManager>(
         sys::Process::getPageSizeEstimate());
 
   this->TargetTriple = std::move(TargetTriple);
   this->PageSize = PageSize;
void SelfExecutorProcessControl::callWrapperAsync(ExecutorAddr WrapperFnAddr,
   SendResult(WrapperFn(ArgBuffer.data(), ArgBuffer.size()));
 }
 
 Error SelfExecutorProcessControl::disconnect() {
   D->shutdown();
   return Error::success();
 }
 
-void SelfExecutorProcessControl::writeUInt8sAsync(
-    ArrayRef<tpctypes::UInt8Write> Ws, WriteResultFn OnWriteComplete) {
+void InProcessMemoryAccess::writeUInt8sAsync(ArrayRef<tpctypes::UInt8Write> Ws,
+                                             WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)
     *W.Addr.toPtr<uint8_t *>() = W.Value;
   OnWriteComplete(Error::success());
 }
 
-void SelfExecutorProcessControl::writeUInt16sAsync(
+void InProcessMemoryAccess::writeUInt16sAsync(
     ArrayRef<tpctypes::UInt16Write> Ws, WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)
     *W.Addr.toPtr<uint16_t *>() = W.Value;
   OnWriteComplete(Error::success());
 }
 
-void SelfExecutorProcessControl::writeUInt32sAsync(
+void InProcessMemoryAccess::writeUInt32sAsync(
     ArrayRef<tpctypes::UInt32Write> Ws, WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)
     *W.Addr.toPtr<uint32_t *>() = W.Value;
   OnWriteComplete(Error::success());
 }
 
-void SelfExecutorProcessControl::writeUInt64sAsync(
+void InProcessMemoryAccess::writeUInt64sAsync(
     ArrayRef<tpctypes::UInt64Write> Ws, WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)
     *W.Addr.toPtr<uint64_t *>() = W.Value;
   OnWriteComplete(Error::success());
 }
 
-void SelfExecutorProcessControl::writeBuffersAsync(
+void InProcessMemoryAccess::writeBuffersAsync(
     ArrayRef<tpctypes::BufferWrite> Ws, WriteResultFn OnWriteComplete) {
   for (auto &W : Ws)
     memcpy(W.Addr.toPtr<char *>(), W.Buffer.data(), W.Buffer.size());
   OnWriteComplete(Error::success());
 }
 
+void InProcessMemoryAccess::writePointersAsync(
+    ArrayRef<tpctypes::PointerWrite> Ws, WriteResultFn OnWriteComplete) {
+  if (IsArch64Bit) {
+    for (auto &W : Ws)
+      *W.Addr.toPtr<uint64_t *>() = W.Value.getValue();
+  } else {
+    for (auto &W : Ws)
+      *W.Addr.toPtr<uint32_t *>() = static_cast<uint32_t>(W.Value.getValue());
+  }
+
+  OnWriteComplete(Error::success());
+}
+
 shared::CWrapperFunctionResult
 SelfExecutorProcessControl::jitDispatchViaWrapperFunctionManager(
     void *Ctx, const void *FnTag, const char *Data, size_t Size) {
 
   LLVM_DEBUG({
     dbgs() << "jit-dispatch call with tag " << FnTag << " and " << Size
            << " byte payload.\n";
   });

llvm/unittests/ExecutionEngine/Orc/EPCGenericMemoryAccessTest.cpp

return WrapperFunction<void(SPSSequence<SPSMemoryAccessBufferWrite>)>::handle(
              [](std::vector<tpctypes::BufferWrite> Ws) {
                for (auto &W : Ws)
                  memcpy(W.Addr.template toPtr<char *>(), W.Buffer.data(),
                         W.Buffer.size());
              })
       .release();
 }
 
+llvm::orc::shared::CWrapperFunctionResult testWritePointers(const char *ArgData,
+                                                            size_t ArgSize) {
+  return WrapperFunction<void(SPSSequence<SPSMemoryAccessPointerWrite>)>::
+      handle(ArgData, ArgSize,
+             [](std::vector<tpctypes::PointerWrite> Ws) {
+               for (auto &W : Ws)
+                 *W.Addr.template toPtr<uint64_t *>() = W.Value.getValue();
+             })
+          .release();
+}
+
 TEST(EPCGenericMemoryAccessTest, MemWrites) {
   auto SelfEPC = cantFail(SelfExecutorProcessControl::Create());
 
   EPCGenericMemoryAccess::FuncAddrs FAs;
   FAs.WriteUInt8s = ExecutorAddr::fromPtr(
       &testWriteUInts<tpctypes::UInt8Write, SPSMemoryAccessUInt8Write>);
   FAs.WriteUInt16s = ExecutorAddr::fromPtr(
       &testWriteUInts<tpctypes::UInt16Write, SPSMemoryAccessUInt16Write>);
   FAs.WriteUInt32s = ExecutorAddr::fromPtr(
       &testWriteUInts<tpctypes::UInt32Write, SPSMemoryAccessUInt32Write>);
   FAs.WriteUInt64s = ExecutorAddr::fromPtr(
       &testWriteUInts<tpctypes::UInt64Write, SPSMemoryAccessUInt64Write>);
   FAs.WriteBuffers = ExecutorAddr::fromPtr(&testWriteBuffers);
+  FAs.WritePointers = ExecutorAddr::fromPtr(&testWritePointers);
 
   auto MemAccess = std::make_unique<EPCGenericMemoryAccess>(*SelfEPC, FAs);
 
   uint8_t Test_UInt8_1 = 0;
   uint8_t Test_UInt8_2 = 0;
   uint16_t Test_UInt16 = 0;
   uint32_t Test_UInt32 = 0;
   uint64_t Test_UInt64 = 0;
+  uint64_t Test_Pointer = 0;
   char Test_Buffer[21];
 
   auto Err1 =
       MemAccess->writeUInt8s({{ExecutorAddr::fromPtr(&Test_UInt8_1), 1},
                               {ExecutorAddr::fromPtr(&Test_UInt8_2), 0xFE}});
 
   EXPECT_THAT_ERROR(std::move(Err1), Succeeded());
   EXPECT_EQ(Test_UInt8_1, 1U);
TEST(EPCGenericMemoryAccessTest, MemWrites) {
   EXPECT_EQ(Test_UInt64, 1U);
 
   StringRef TestMsg("test-message");
   auto Err5 =
       MemAccess->writeBuffers({{ExecutorAddr::fromPtr(&Test_Buffer), TestMsg}});
   EXPECT_THAT_ERROR(std::move(Err5), Succeeded());
   EXPECT_EQ(StringRef(Test_Buffer, TestMsg.size()), TestMsg);
 
+  auto Err6 = MemAccess->writePointers(
+      {{ExecutorAddr::fromPtr(&Test_Pointer), ExecutorAddr(1U)}});
+  EXPECT_THAT_ERROR(std::move(Err6), Succeeded());
+  EXPECT_EQ(Test_Pointer, 1U);
+
   cantFail(SelfEPC->disconnect());
 }
 
 } // namespace