Visitor: Add support for nested visitor clients

nhaehnle · nhaehnle · commit 55c4e43c2c34 · 2023-03-29T13:22:44.000+02:00
The idea here is that we can compose a complex lowering pass out of
smaller "building blocks" that are all implemented in terms of visiting
functions.
diff --git a/example/ExampleMain.cpp b/example/ExampleMain.cpp
@@ -42,12 +42,14 @@ using namespace llvm_dialects;
 enum class Action {
   Build,
   Verify,
+  Visit,
 };
 
 cl::opt<Action> g_action(
     cl::desc("Action to perform:"), cl::init(Action::Build),
     cl::values(clEnumValN(Action::Build, "build", "Example IRBuilder use"),
-               clEnumValN(Action::Verify, "verify", "Verify an input module")));
+               clEnumValN(Action::Verify, "verify", "Verify an input module"),
+               clEnumValN(Action::Visit, "visit", "Example Visitor use")));
 
 // Input sources
 cl::list<std::string> g_inputs(cl::Positional, cl::ZeroOrMore,
@@ -132,6 +134,54 @@ std::unique_ptr<Module> createModuleExampleTypedPtrs(LLVMContext &context) {
   return module;
 }
 
+struct VisitorInnermost {
+  int counter = 0;
+};
+
+struct VisitorNest {
+  raw_ostream *out = nullptr;
+  VisitorInnermost inner;
+};
+
+struct VisitorContainer {
+  int padding;
+  VisitorNest nest;
+};
+
+template <>
+struct llvm_dialects::VisitorPayloadProjection<VisitorNest, raw_ostream> {
+  static raw_ostream &project(VisitorNest &nest) { return *nest.out; }
+};
+
+LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD(VisitorContainer, nest)
+LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD(VisitorNest, inner)
+
+void exampleVisit(Module &module) {
+  auto visitor =
+      VisitorBuilder<VisitorContainer>()
+          .nest<VisitorNest>([](VisitorBuilder<VisitorNest> &b) {
+            b.add<xd::ReadOp>([](VisitorNest &self, xd::ReadOp &op) {
+              *self.out << "visiting ReadOp: " << op << '\n';
+            });
+            b.nest<raw_ostream>([](VisitorBuilder<raw_ostream> &b) {
+              b.add<xd::WriteOp>([](raw_ostream &out, xd::WriteOp &op) {
+                out << "visiting WriteOp: " << op << '\n';
+              });
+            });
+            b.nest<VisitorInnermost>([](VisitorBuilder<VisitorInnermost> &b) {
+              b.add<xd::ITruncOp>([](VisitorInnermost &inner,
+                                     xd::ITruncOp &op) { inner.counter++; });
+            });
+          })
+          .build();
+
+  VisitorContainer container;
+  container.nest.out = &outs();
+  visitor.visit(container, module);
+
+  outs() << "inner.counter = " << container.nest.inner.counter << '\n';
+}
+
 int main(int argc, char **argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv);
 
@@ -143,7 +193,7 @@ int main(int argc, char **argv) {
   if (g_action == Action::Build) {
     auto module = g_typedPointers ? createModuleExampleTypedPtrs(context) : createModuleExample(context);
     module->print(llvm::outs(), nullptr, false);
-  } else if (g_action == Action::Verify) {
+  } else {
     if (g_inputs.size() != 1) {
       errs() << "Need exactly one input module\n";
       return 1;
@@ -164,10 +214,14 @@ int main(int argc, char **argv) {
       return 1;
     }
 
-    if (!verify(*module, outs()))
-      return 1;
-  } else {
-    report_fatal_error("unhandled action");
+    if (g_action == Action::Verify) {
+      if (!verify(*module, outs()))
+        return 1;
+    } else if (g_action == Action::Visit) {
+      exampleVisit(*module);
+    } else {
+      report_fatal_error("unhandled action");
+    }
   }
 
   return 0;
diff --git a/include/llvm-dialects/Dialect/Visitor.h b/include/llvm-dialects/Dialect/Visitor.h
@@ -45,37 +45,107 @@ enum class VisitorStrategy {
   ByFunctionDeclaration,
 };
 
+/// @brief Auxiliary template to support nested visitor clients.
+///
+/// It can be convenient to write visitor clients in a modular fashion and
+/// combine them together to use a single @ref Visitor. Each client module
+/// may use a different payload type.
+///
+/// In this case, each nested client module's payload must be reachable from
+/// the top-level payload, and a specialization of this template must be
+/// provided that "projects" parent payload references to child payload
+/// references.
+///
+/// In cases where the nested payload is a field of the parent payload, the
+/// @ref LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD can be used instead.
+template <typename PayloadT, typename NestedPayloadT>
+struct VisitorPayloadProjection {
+  // Template specializations must implement this static method:
+  //
+  //   static NestedPayloadT &project(PayloadT &);
+};
+
+/// Declare that `PayloadT` can be projected to a nested payload type via
+/// `field`.
+///
+/// This creates a specialization of @ref VisitorPayloadProjection and must
+/// therefore typically be outside of any namespace. The nested type is derived
+/// automatically.
+#define LLVM_DIALECTS_VISITOR_PAYLOAD_PROJECT_FIELD(PayloadT, field)           \
+  template <>                                                                  \
+  struct llvm_dialects::detail::VisitorPayloadOffsetProjection<                \
+      PayloadT,                                                                \
+      std::remove_reference_t<decltype(std::declval<PayloadT>().field)>> {     \
+    static constexpr bool useOffsetof = true;                                  \
+    static constexpr std::size_t offset = offsetof(PayloadT, field);           \
+  };
+
 namespace detail {
 
 class VisitorBase;
 
 using VisitorCallback = void (void *, void *, llvm::Instruction *);
-using VisitorCase = std::tuple<const OpDescription *, void *, VisitorCallback *>;
+using PayloadProjectionCallback = void *(void *);
+
+/// Apply first the byte offset and then the projection function. If projection
+/// is null, stop the projection sequence.
+struct PayloadProjection {
+  std::size_t offset = 0;
+  PayloadProjectionCallback *projection = nullptr;
+};
+
+struct VisitorCase {
+  const OpDescription *description = nullptr;
+  VisitorCallback *callback = nullptr;
+  void *callbackData = nullptr;
+
+  // If non-negative, a byte offset to apply to the payload. If negative,
+  // a shifted index into the projections vector.
+  ssize_t projection = 0;
+};
+
+template <typename PayloadT, typename NestedPayloadT>
+struct VisitorPayloadOffsetProjection {
+  static constexpr bool useOffsetof = false;
+};
 
 class VisitorBuilderBase {
   friend class VisitorBase;
 public:
+  VisitorBuilderBase() = default;
+  explicit VisitorBuilderBase(VisitorBuilderBase *parent) : m_parent(parent) {}
+  ~VisitorBuilderBase();
+
   void setStrategy(VisitorStrategy strategy) { m_strategy = strategy; }
 
-protected:
   void add(const OpDescription &desc, void *extra, VisitorCallback *fn);
 
+public:
+  PayloadProjectionCallback *m_projection = nullptr;
+  size_t m_offsetProjection = 0;
+
 private:
+  VisitorBuilderBase *m_parent = nullptr;
   VisitorStrategy m_strategy = VisitorStrategy::ByFunctionDeclaration;
   llvm::SmallVector<VisitorCase> m_cases;
+  llvm::SmallVector<PayloadProjection> m_projections;
 };
 
 class VisitorBase {
 protected:
-  VisitorBase(VisitorBuilderBase builder);
+  VisitorBase(VisitorBuilderBase &&builder);
 
   void visit(void *payload, llvm::Instruction &inst) const;
   void visit(void *payload, llvm::Function &fn) const;
   void visit(void *payload, llvm::Module &module) const;
 
 private:
+  void call(const VisitorCase &theCase, void *payload,
+            llvm::Instruction &inst) const;
+
   VisitorStrategy m_strategy;
   llvm::SmallVector<VisitorCase> m_cases;
+  llvm::SmallVector<PayloadProjection> m_projections;
 };
 
 } // namespace detail
@@ -91,7 +161,7 @@ class VisitorBase {
 template <typename PayloadT>
 class Visitor : public detail::VisitorBase {
 public:
-  Visitor(detail::VisitorBuilderBase builder)
+  Visitor(detail::VisitorBuilderBase &&builder)
       : VisitorBase(std::move(builder)) {}
 
   void visit(PayloadT &payload, llvm::Instruction &inst) const {
@@ -130,16 +200,20 @@ class Visitor : public detail::VisitorBase {
 /// myVisitor(myPayload, module);
 /// @endcode
 template <typename PayloadT>
-class VisitorBuilder : public detail::VisitorBuilderBase {
+class VisitorBuilder : private detail::VisitorBuilderBase {
+  template <typename OtherT> friend class VisitorBuilder;
+
 public:
   using Payload = PayloadT;
 
+  VisitorBuilder() = default;
+
   VisitorBuilder &setStrategy(VisitorStrategy strategy) {
     VisitorBuilderBase::setStrategy(strategy);
     return *this;
   }
 
-  Visitor<PayloadT> build() { return {std::move(*this)}; }
+  Visitor<PayloadT> build() { return std::move(*this); }
 
   template <typename OpT> VisitorBuilder &add(void (*fn)(PayloadT &, OpT &)) {
     VisitorBuilderBase::add(
@@ -151,6 +225,32 @@ class VisitorBuilder : public detail::VisitorBuilderBase {
         });
     return *this;
   }
+
+  template <typename NestedPayloadT>
+  VisitorBuilder &nest(void (*registration)(VisitorBuilder<NestedPayloadT> &)) {
+    VisitorBuilder<NestedPayloadT> nested{this};
+
+    if constexpr (detail::VisitorPayloadOffsetProjection<
+                      PayloadT, NestedPayloadT>::useOffsetof) {
+      nested.m_offsetProjection =
+          detail::VisitorPayloadOffsetProjection<PayloadT,
+                                                 NestedPayloadT>::offset;
+    } else {
+      nested.m_projection = [](void *payload) -> void * {
+        return static_cast<void *>(
+            &VisitorPayloadProjection<PayloadT, NestedPayloadT>::project(
+                *static_cast<PayloadT *>(payload)));
+      };
+    }
+
+    (*registration)(nested);
+
+    return *this;
+  }
+
+private:
+  explicit VisitorBuilder(VisitorBuilderBase *parent)
+      : VisitorBuilderBase(parent) {}
 };
 
 } // namespace llvm_dialects
diff --git a/lib/Dialect/Visitor.cpp b/lib/Dialect/Visitor.cpp
@@ -29,18 +29,74 @@ using namespace llvm;
 using llvm_dialects::detail::VisitorBuilderBase;
 using llvm_dialects::detail::VisitorBase;
 
+VisitorBuilderBase::~VisitorBuilderBase() {
+  if (m_parent) {
+    // Build up the projection sequence by walking all the way to the root.
+    SmallVector<PayloadProjection> projectionSequence;
+    projectionSequence.emplace_back();
+
+    VisitorBuilderBase *ancestor = this;
+    for (; ancestor->m_parent; ancestor = ancestor->m_parent) {
+      if (ancestor->m_projection) {
+        PayloadProjection next;
+        next.projection = ancestor->m_projection;
+        projectionSequence.push_back(next);
+      } else {
+        projectionSequence.back().offset += ancestor->m_offsetProjection;
+      }
+    }
+
+    ssize_t projection = 0;
+    if (projectionSequence.size() == 1) {
+      projection = projectionSequence[0].offset;
+    } else {
+      projection = -(1 + ancestor->m_projections.size());
+      ancestor->m_projections.append(projectionSequence.rbegin(),
+                                     projectionSequence.rend());
+    }
+
+    for (auto &theCase : m_cases)
+      theCase.projection = projection;
+
+    // Copy all cases to the root.
+    ancestor->m_cases.append(m_cases.begin(), m_cases.end());
+  }
+}
+
 void VisitorBuilderBase::add(const OpDescription &desc, void *extra, VisitorCallback *fn) {
-  m_cases.emplace_back(&desc, extra, fn);
+  VisitorCase theCase;
+  theCase.description = &desc;
+  theCase.callback = fn;
+  theCase.callbackData = extra;
+  theCase.projection = 0;
+  m_cases.emplace_back(theCase);
 }
 
-VisitorBase::VisitorBase(VisitorBuilderBase builder)
-    : m_strategy(builder.m_strategy), m_cases(std::move(builder.m_cases)) {
+VisitorBase::VisitorBase(VisitorBuilderBase &&builder)
+    : m_strategy(builder.m_strategy), m_cases(std::move(builder.m_cases)),
+      m_projections(std::move(builder.m_projections)) {
+  assert(!builder.m_parent);
+}
+
+void VisitorBase::call(const VisitorCase &theCase, void *payload,
+                       Instruction &inst) const {
+  if (theCase.projection >= 0) {
+    payload = (char *)payload + theCase.projection;
+  } else {
+    for (size_t idx = -theCase.projection - 1;; ++idx) {
+      payload = (char *)payload + m_projections[idx].offset;
+      if (!m_projections[idx].projection)
+        break;
+      payload = m_projections[idx].projection(payload);
+    }
+  }
+  theCase.callback(theCase.callbackData, payload, &inst);
 }
 
 void VisitorBase::visit(void *payload, Instruction &inst) const {
-  for (const auto &[desc, extra, callback] : m_cases) {
-    if (desc->matchInstruction(inst))
-      callback(extra, payload, &inst);
+  for (const auto &theCase : m_cases) {
+    if (theCase.description->matchInstruction(inst))
+      call(theCase, payload, inst);
   }
 }
 
@@ -59,15 +115,15 @@ void VisitorBase::visit(void *payload, Function &fn) const {
 
     LLVM_DEBUG(dbgs() << "visit " << decl.getName() << '\n');
 
-    for (const auto &[desc, extra, callback] : m_cases) {
-      if (desc->matchDeclaration(decl)) {
+    for (const auto &theCase : m_cases) {
+      if (theCase.description->matchDeclaration(decl)) {
         for (Use &use : decl.uses()) {
           if (auto *inst = dyn_cast<Instruction>(use.getUser())) {
             if (inst->getFunction() != &fn)
               continue;
-            if (auto *call = dyn_cast<CallInst>(inst)) {
-              if (&use == &call->getCalledOperandUse())
-                callback(extra, payload, call);
+            if (auto *callInst = dyn_cast<CallInst>(inst)) {
+              if (&use == &callInst->getCalledOperandUse())
+                call(theCase, payload, *callInst);
             }
           }
         }
@@ -89,12 +145,12 @@ void VisitorBase::visit(void *payload, Module &module) const {
     if (!decl.isDeclaration())
       continue;
 
-    for (const auto &[desc, extra, callback] : m_cases) {
-      if (desc->matchDeclaration(decl)) {
+    for (const auto &theCase : m_cases) {
+      if (theCase.description->matchDeclaration(decl)) {
         for (Use &use : decl.uses()) {
-          if (auto *call = dyn_cast<CallInst>(use.getUser())) {
-            if (&use == &call->getCalledOperandUse())
-              callback(extra, payload, call);
+          if (auto *callInst = dyn_cast<CallInst>(use.getUser())) {
+            if (&use == &callInst->getCalledOperandUse())
+              call(theCase, payload, *callInst);
           }
         }
       }
diff --git a/test/example/visitor-basic.ll b/test/example/visitor-basic.ll
