Visitor: Add convenient way of using member functions

Using member functions of the payload type seems likely to be a very
common pattern, especially for converting code that uses the
llvm::InstVisitor.

This complicates the VisitorHandler setup a bit. At some point, one
has to wonder if it wouldn't be better to just use std::function. The
current system still has some performance advantages:

* It is able to handle generic projections while still handling the
  common case of projections to fields cheaply, with a single, perfectly
  predictable branch.

* VisitorHandler is 32 bytes; so is std::function, so wrapping
  std::function would be more expensive in terms of cache pressure.

And we specifically do not want the full power of std::function. Users
of this facility should be strongly encouraged to use the pattern where
a visitor is constructed once as a static variable, so that the cost of
building hash maps is amortized.

Author: nhaehnle

example/ExampleMain.cpp

@@ -146,6 +146,10 @@ struct VisitorInnermost {
 struct VisitorNest {
   raw_ostream *out = nullptr;
   VisitorInnermost inner;
+
+  void visitBinaryOperator(BinaryOperator &inst) {
+    *out << "visiting BinaryOperator: " << inst << '\n';
+  }
 };
 
 struct VisitorContainer {
@@ -179,9 +183,7 @@ template <bool rpot> const Visitor<VisitorContainer> &getExampleVisitor() {
                 [](VisitorNest &self, UnaryInstruction &inst) {
                   *self.out << "visiting UnaryInstruction: " << inst << '\n';
                 });
-            b.add<BinaryOperator>([](VisitorNest &self, BinaryOperator &inst) {
-              *self.out << "visiting BinaryOperator: " << inst << '\n';
-            });
+            b.add(&VisitorNest::visitBinaryOperator);
             b.nest<raw_ostream>([](VisitorBuilder<raw_ostream> &b) {
               b.add<xd::WriteOp>([](raw_ostream &out, xd::WriteOp &op) {
                 out << "visiting WriteOp: " << op << '\n';

include/llvm-dialects/Dialect/Visitor.h

@@ -123,16 +123,31 @@ class VisitorKey {
   unsigned m_intrinsicId = 0;
 };
 
-using VisitorCallback = void (void *, void *, llvm::Instruction *);
+struct VisitorCallbackData {
+private:
+  static constexpr size_t SizeofMemberFunctionPointer =
+      sizeof(void(VisitorCallbackData::*)());
+  static constexpr size_t SizeofFunctionPointer = sizeof(void (*)());
+
+public:
+  // Size is 16 bytes in the 64-bit Itanium ABI.
+  static constexpr size_t Size =
+      std::max(SizeofFunctionPointer, SizeofMemberFunctionPointer);
+
+  char data[Size];
+};
+
+using VisitorCallback = void(const VisitorCallbackData &, void *,
+                             llvm::Instruction *);
 using PayloadProjectionCallback = void *(void *);
 
 struct VisitorHandler {
-  VisitorCallback *callback = nullptr;
-  void *callbackData = nullptr;
-
   // If non-negative, a byte offset to apply to the payload. If negative,
   // a shifted index into the payload projections vector.
   ssize_t projection = 0;
+
+  VisitorCallback *callback = nullptr;
+  VisitorCallbackData data;
 };
 
 /// Apply first the byte offset and then the projection function. If projection
@@ -157,7 +172,7 @@ class VisitorTemplate {
 
 public:
   void setStrategy(VisitorStrategy strategy);
-  void add(VisitorKey key, void *extra, VisitorCallback *fn,
+  void add(VisitorKey key, VisitorCallback *fn, VisitorCallbackData data,
            ssize_t projection);
 
 private:
@@ -200,7 +215,7 @@ class VisitorBuilderBase {
 
   void setStrategy(VisitorStrategy strategy);
 
-  void add(VisitorKey key, void *extra, VisitorCallback *fn);
+  void add(VisitorKey key, VisitorCallback *fn, VisitorCallbackData data);
 
   VisitorBase build();
 
@@ -316,12 +331,23 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
     return *this;
   }
 
+  template <typename OpT> VisitorBuilder &add(void (PayloadT::*fn)(OpT &)) {
+    addMemberFnCase<OpT>(detail::VisitorKey::op<OpT>(), fn);
+    return *this;
+  }
+
   VisitorBuilder &addIntrinsic(unsigned id,
                                void (*fn)(PayloadT &, llvm::IntrinsicInst &)) {
     addCase<llvm::IntrinsicInst>(detail::VisitorKey::intrinsic(id), fn);
     return *this;
   }
 
+  VisitorBuilder &addIntrinsic(unsigned id,
+                               void (PayloadT::*fn)(llvm::IntrinsicInst &)) {
+    addMemberFnCase<llvm::IntrinsicInst>(detail::VisitorKey::intrinsic(id), fn);
+    return *this;
+  }
+
   template <typename NestedPayloadT>
   VisitorBuilder &nest(void (*registration)(VisitorBuilder<NestedPayloadT> &)) {
     if constexpr (detail::VisitorPayloadOffsetProjection<
@@ -354,14 +380,36 @@ class VisitorBuilder : private detail::VisitorBuilderBase {
 
   template <typename OpT>
   void addCase(detail::VisitorKey key, void (*fn)(PayloadT &, OpT &)) {
-    VisitorBuilderBase::add(key, (void *)fn, &VisitorBuilder::forwarder<OpT>);
+    detail::VisitorCallbackData data{};
+    static_assert(sizeof(fn) <= sizeof(data.data));
+    memcpy(&data.data, &fn, sizeof(fn));
+    VisitorBuilderBase::add(key, &VisitorBuilder::forwarder<OpT>, data);
   }
 
   template <typename OpT>
-  static void forwarder(void *extra, void *payload, llvm::Instruction *op) {
-    auto fn = (void (*)(PayloadT &, OpT &))extra;
+  void addMemberFnCase(detail::VisitorKey key, void (PayloadT::*fn)(OpT &)) {
+    detail::VisitorCallbackData data{};
+    static_assert(sizeof(fn) <= sizeof(data.data));
+    memcpy(&data.data, &fn, sizeof(fn));
+    VisitorBuilderBase::add(key, &VisitorBuilder::memberFnForwarder<OpT>, data);
+  }
+
+  template <typename OpT>
+  static void forwarder(const detail::VisitorCallbackData &data, void *payload,
+                        llvm::Instruction *op) {
+    void (*fn)(PayloadT &, OpT &);
+    memcpy(&fn, &data.data, sizeof(fn));
     fn(*static_cast<PayloadT *>(payload), *llvm::cast<OpT>(op));
   }
+
+  template <typename OpT>
+  static void memberFnForwarder(const detail::VisitorCallbackData &data,
+                                void *payload, llvm::Instruction *op) {
+    void (PayloadT::*fn)(OpT &);
+    memcpy(&fn, &data.data, sizeof(fn));
+    PayloadT *self = static_cast<PayloadT *>(payload);
+    (self->*fn)(*llvm::cast<OpT>(op));
+  }
 };
 
 } // namespace llvm_dialects

lib/Dialect/Visitor.cpp

@@ -30,6 +30,7 @@ using namespace llvm;
 
 using llvm_dialects::detail::VisitorBase;
 using llvm_dialects::detail::VisitorBuilderBase;
+using llvm_dialects::detail::VisitorCallbackData;
 using llvm_dialects::detail::VisitorHandler;
 using llvm_dialects::detail::VisitorKey;
 using llvm_dialects::detail::VisitorTemplate;
@@ -43,11 +44,11 @@ void VisitorTemplate::setStrategy(VisitorStrategy strategy) {
   m_strategy = strategy;
 }
 
-void VisitorTemplate::add(VisitorKey key, void *extra, VisitorCallback *fn,
-                          ssize_t projection) {
+void VisitorTemplate::add(VisitorKey key, VisitorCallback *fn,
+                          VisitorCallbackData data, ssize_t projection) {
   VisitorHandler handler;
   handler.callback = fn;
-  handler.callbackData = extra;
+  handler.data = data;
   handler.projection = projection;
   m_handlers.emplace_back(handler);
 
@@ -140,8 +141,9 @@ void VisitorBuilderBase::setStrategy(VisitorStrategy strategy) {
   m_template->setStrategy(strategy);
 }
 
-void VisitorBuilderBase::add(VisitorKey key, void *extra, VisitorCallback *fn) {
-  m_template->add(key, extra, fn, m_projection);
+void VisitorBuilderBase::add(VisitorKey key, VisitorCallback *fn,
+                             VisitorCallbackData data) {
+  m_template->add(key, fn, data, m_projection);
 }
 
 VisitorBase VisitorBuilderBase::build() {
@@ -220,7 +222,7 @@ void VisitorBase::call(const VisitorHandler &handler, void *payload,
       payload = m_projections[idx].projection(payload);
     }
   }
-  handler.callback(handler.callbackData, payload, &inst);
+  handler.callback(handler.data, payload, &inst);
 }
 
 void VisitorBase::visit(void *payload, Instruction &inst) const {