Visitor: encapsulate projection encoding in a wrapper struct

Also change the encoding to use the MSB but not negation. This shouldn't
make a big difference, but if anything is probably slightly better.

Author: nhaehnle

include/llvm-dialects/Dialect/Visitor.h

@@ -142,9 +142,36 @@ using VisitorCallback = void(const VisitorCallbackData &, void *,
 using PayloadProjectionCallback = void *(void *);
 
 struct VisitorHandler {
-  // 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;
+  // Either a byte offset to apply to the payload or an index into the payload
+  // projections vector for more complex projections.
+  struct Projection {
+    static constexpr std::size_t Msb = (std::size_t)1
+                                       << (8 * sizeof(std::size_t) - 1);
+
+    std::size_t data = 0;
+
+    void setOffset(std::size_t offset) {
+      assert((offset & Msb) == 0);
+      data = offset;
+    }
+
+    void setIndex(std::size_t index) {
+      assert((index & Msb) == 0);
+      data = index | Msb;
+    }
+
+    bool isOffset() const { return (data & Msb) == 0; }
+    std::size_t getOffset() const {
+      assert((data & Msb) == 0);
+      return data;
+    }
+    std::size_t getIndex() const {
+      assert((data & Msb) != 0);
+      return data & ~Msb;
+    }
+  };
+
+  Projection projection;
 
   VisitorCallback *callback = nullptr;
   VisitorCallbackData data;
@@ -173,7 +200,7 @@ class VisitorTemplate {
 public:
   void setStrategy(VisitorStrategy strategy);
   void add(VisitorKey key, VisitorCallback *fn, VisitorCallbackData data,
-           ssize_t projection);
+           VisitorHandler::Projection projection);
 
 private:
   VisitorStrategy m_strategy = VisitorStrategy::Default;
@@ -225,7 +252,7 @@ class VisitorBuilderBase {
 
   VisitorTemplate m_ownedTemplate;
   VisitorTemplate *m_template;
-  ssize_t m_projection;
+  VisitorHandler::Projection m_projection;
 };
 
 /// @brief Base class for Visitors

lib/Dialect/Visitor.cpp

@@ -45,7 +45,8 @@ void VisitorTemplate::setStrategy(VisitorStrategy strategy) {
 }
 
 void VisitorTemplate::add(VisitorKey key, VisitorCallback *fn,
-                          VisitorCallbackData data, ssize_t projection) {
+                          VisitorCallbackData data,
+                          VisitorHandler::Projection projection) {
   VisitorHandler handler;
   handler.callback = fn;
   handler.data = data;
@@ -84,20 +85,19 @@ void VisitorTemplate::add(VisitorKey key, VisitorCallback *fn,
   }
 }
 
-VisitorBuilderBase::VisitorBuilderBase()
-    : m_template(&m_ownedTemplate), m_projection(0) {}
+VisitorBuilderBase::VisitorBuilderBase() : m_template(&m_ownedTemplate) {}
 
 VisitorBuilderBase::VisitorBuilderBase(VisitorBuilderBase *parent,
                                        PayloadProjectionCallback *projection)
     : m_template(parent->m_template) {
   // Extract and update the parent's projection.
   SmallVector<PayloadProjection> sequence;
 
-  if (parent->m_projection >= 0) {
+  if (parent->m_projection.isOffset()) {
     sequence.emplace_back();
-    sequence.back().offset = parent->m_projection;
+    sequence.back().offset = parent->m_projection.getOffset();
   } else {
-    for (size_t idx = -parent->m_projection - 1;; ++idx) {
+    for (size_t idx = parent->m_projection.getIndex();; ++idx) {
       sequence.push_back(m_template->m_projections[idx]);
       if (!sequence.back().projection)
         break;
@@ -108,7 +108,7 @@ VisitorBuilderBase::VisitorBuilderBase(VisitorBuilderBase *parent,
   sequence.emplace_back();
 
   // Setup our projection.
-  m_projection = -(m_template->m_projections.size() + 1);
+  m_projection.setIndex(m_template->m_projections.size());
   m_template->m_projections.insert(m_template->m_projections.end(),
                                    sequence.begin(), sequence.end());
 }
@@ -118,20 +118,20 @@ VisitorBuilderBase::VisitorBuilderBase(VisitorBuilderBase *parent,
     : m_template(parent->m_template) {
   if (offset == 0) {
     m_projection = parent->m_projection;
-  } else if (parent->m_projection >= 0) {
-    m_projection = parent->m_projection + offset;
+  } else if (parent->m_projection.isOffset()) {
+    m_projection.setOffset(parent->m_projection.getOffset() + offset);
   } else {
     // Extract and update the parent's projection.
     SmallVector<PayloadProjection> sequence;
-    for (size_t idx = -parent->m_projection - 1;; ++idx) {
+    for (size_t idx = parent->m_projection.getIndex();; ++idx) {
       sequence.push_back(m_template->m_projections[idx]);
       if (!sequence.back().projection)
         break;
     }
     sequence.back().offset += offset;
 
     // Setup our projection.
-    m_projection = -(m_template->m_projections.size() + 1);
+    m_projection.setIndex(m_template->m_projections.size());
     m_template->m_projections.insert(m_template->m_projections.end(),
                                      sequence.begin(), sequence.end());
   }
@@ -212,10 +212,10 @@ void VisitorBase::call(HandlerRange handlers, void *payload,
 
 void VisitorBase::call(const VisitorHandler &handler, void *payload,
                        Instruction &inst) const {
-  if (handler.projection >= 0) {
-    payload = (char *)payload + handler.projection;
+  if (handler.projection.isOffset()) {
+    payload = (char *)payload + handler.projection.getOffset();
   } else {
-    for (size_t idx = -handler.projection - 1;; ++idx) {
+    for (size_t idx = handler.projection.getIndex();; ++idx) {
       payload = (char *)payload + m_projections[idx].offset;
       if (!m_projections[idx].projection)
         break;