Line data    Source code

       1             : #ifdef __PLUMED_HAS_ASMJIT
       2             : // This file is part of AsmJit project <https://asmjit.com>
       3             : //
       4             : // See asmjit.h or LICENSE.md for license and copyright information
       5             : // SPDX-License-Identifier: Zlib
       6             : 
       7             : #ifndef ASMJIT_CORE_ZONETREE_H_INCLUDED
       8             : #define ASMJIT_CORE_ZONETREE_H_INCLUDED
       9             : 
      10             : #include "support.h"
      11             : 
      12             : ASMJIT_BEGIN_NAMESPACE
      13             : 
      14             : //! \addtogroup asmjit_zone
      15             : //! \{
      16             : 
      17             : //! RB-Tree node.
      18             : //!
      19             : //! The color is stored in a least significant bit of the `left` node.
      20             : //!
      21             : //! WARNING: Always use accessors to access left and right children.
      22             : class ZoneTreeNode {
      23             : public:
      24             :   ASMJIT_NONCOPYABLE(ZoneTreeNode)
      25             : 
      26             :   //! \name Constants
      27             :   //! \{
      28             : 
      29             :   enum : uintptr_t {
      30             :     kRedMask = 0x1,
      31             :     kPtrMask = ~kRedMask
      32             :   };
      33             : 
      34             :   //! \}
      35             : 
      36             :   //! \name Members
      37             :   //! \{
      38             : 
      39             :   uintptr_t _rbNodeData[2];
      40             : 
      41             :   //! \}
      42             : 
      43             :   //! \name Construction & Destruction
      44             :   //! \{
      45             : 
      46             :   inline ZoneTreeNode() noexcept
      47       27392 :     : _rbNodeData { 0, 0 } {}
      48             : 
      49             :   //! \}
      50             : 
      51             :   //! \name Accessors
      52             :   //! \{
      53             : 
      54           0 :   inline bool isRed() const noexcept { return static_cast<bool>(_rbNodeData[0] & kRedMask); }
      55             : 
      56           0 :   inline bool hasChild(size_t i) const noexcept { return _rbNodeData[i] > kRedMask; }
      57             :   inline bool hasLeft() const noexcept { return _rbNodeData[0] > kRedMask; }
      58             :   inline bool hasRight() const noexcept { return _rbNodeData[1] != 0; }
      59             : 
      60             :   template<typename T = ZoneTreeNode>
      61             :   inline T* child(size_t i) const noexcept { return static_cast<T*>(_getChild(i)); }
      62             :   template<typename T = ZoneTreeNode>
      63             :   inline T* left() const noexcept { return static_cast<T*>(_getLeft()); }
      64             :   template<typename T = ZoneTreeNode>
      65             :   inline T* right() const noexcept { return static_cast<T*>(_getRight()); }
      66             : 
      67             :   //! \}
      68             : 
      69             :   //! \cond INTERNAL
      70             :   //! \name Internal
      71             :   //! \{
      72             : 
      73           0 :   inline ZoneTreeNode* _getChild(size_t i) const noexcept { return (ZoneTreeNode*)(_rbNodeData[i] & kPtrMask); }
      74           0 :   inline ZoneTreeNode* _getLeft() const noexcept { return (ZoneTreeNode*)(_rbNodeData[0] & kPtrMask); }
      75           0 :   inline ZoneTreeNode* _getRight() const noexcept { return (ZoneTreeNode*)(_rbNodeData[1]); }
      76             : 
      77           0 :   inline void _setChild(size_t i, ZoneTreeNode* node) noexcept { _rbNodeData[i] = (_rbNodeData[i] & kRedMask) | (uintptr_t)node; }
      78             :   inline void _setLeft(ZoneTreeNode* node) noexcept { _rbNodeData[0] = (_rbNodeData[0] & kRedMask) | (uintptr_t)node; }
      79           0 :   inline void _setRight(ZoneTreeNode* node) noexcept { _rbNodeData[1] = (uintptr_t)node; }
      80             : 
      81           0 :   inline void _makeRed() noexcept { _rbNodeData[0] |= kRedMask; }
      82           0 :   inline void _makeBlack() noexcept { _rbNodeData[0] &= kPtrMask; }
      83             : 
      84             :   //! Tests whether the node is RED (RED node must be non-null and must have RED flag set).
      85           0 :   static inline bool _isValidRed(ZoneTreeNode* node) noexcept { return node && node->isRed(); }
      86             : 
      87             :   //! \}
      88             :   //! \endcond
      89             : };
      90             : 
      91             : //! RB-Tree node casted to `NodeT`.
      92             : template<typename NodeT>
      93             : class ZoneTreeNodeT : public ZoneTreeNode {
      94             : public:
      95             :   ASMJIT_NONCOPYABLE(ZoneTreeNodeT)
      96             : 
      97             :   //! \name Construction & Destruction
      98             :   //! \{
      99             : 
     100             :   inline ZoneTreeNodeT() noexcept
     101             :     : ZoneTreeNode() {}
     102             : 
     103             :   //! \}
     104             : 
     105             :   //! \name Accessors
     106             :   //! \{
     107             : 
     108             :   inline NodeT* child(size_t i) const noexcept { return static_cast<NodeT*>(_getChild(i)); }
     109             :   inline NodeT* left() const noexcept { return static_cast<NodeT*>(_getLeft()); }
     110             :   inline NodeT* right() const noexcept { return static_cast<NodeT*>(_getRight()); }
     111             : 
     112             :   //! \}
     113             : };
     114             : 
     115             : //! RB-Tree.
     116             : template<typename NodeT>
     117             : class ZoneTree {
     118             : public:
     119             :   ASMJIT_NONCOPYABLE(ZoneTree)
     120             : 
     121             :   typedef NodeT Node;
     122             :   NodeT* _root;
     123             : 
     124             :   //! \name Construction & Destruction
     125             :   //! \{
     126             : 
     127       54914 :   inline ZoneTree() noexcept
     128       54914 :     : _root(nullptr) {}
     129             : 
     130             :   inline ZoneTree(ZoneTree&& other) noexcept
     131             :     : _root(other._root) {}
     132             : 
     133       54914 :   inline void reset() noexcept { _root = nullptr; }
     134             : 
     135             :   //! \}
     136             : 
     137             :   //! \name Accessors
     138             :   //! \{
     139             : 
     140             :   inline bool empty() const noexcept { return _root == nullptr; }
     141           0 :   inline NodeT* root() const noexcept { return static_cast<NodeT*>(_root); }
     142             : 
     143             :   //! \}
     144             : 
     145             :   //! \name Utilities
     146             :   //! \{
     147             : 
     148             :   inline void swap(ZoneTree& other) noexcept {
     149             :     std::swap(_root, other._root);
     150             :   }
     151             : 
     152             :   template<typename CompareT = Support::Compare<Support::SortOrder::kAscending>>
     153       27392 :   void insert(NodeT* ASMJIT_NONNULL(node), const CompareT& cmp = CompareT()) noexcept {
     154             :     // Node to insert must not contain garbage.
     155             :     ASMJIT_ASSERT(!node->hasLeft());
     156             :     ASMJIT_ASSERT(!node->hasRight());
     157             :     ASMJIT_ASSERT(!node->isRed());
     158             : 
     159       27392 :     if (!_root) {
     160       27392 :       _root = node;
     161       27392 :       return;
     162             :     }
     163             : 
     164             :     ZoneTreeNode head;           // False root node,
     165             :     head._setRight(_root);       // having root on the right.
     166             : 
     167             :     ZoneTreeNode* g = nullptr;   // Grandparent.
     168             :     ZoneTreeNode* p = nullptr;   // Parent.
     169             :     ZoneTreeNode* t = &head;     // Iterator.
     170             :     ZoneTreeNode* q = _root;     // Query.
     171             : 
     172             :     size_t dir = 0;              // Direction for accessing child nodes.
     173             :     size_t last = 0;             // Not needed to initialize, but makes some tools happy.
     174             : 
     175             :     node->_makeRed();            // New nodes are always red and violations fixed appropriately.
     176             : 
     177             :     // Search down the tree.
     178             :     for (;;) {
     179           0 :       if (!q) {
     180             :         // Insert new node at the bottom.
     181             :         q = node;
     182             :         p->_setChild(dir, node);
     183             :       }
     184             :       else if (_isValidRed(q->_getLeft()) && _isValidRed(q->_getRight())) {
     185             :         // Color flip.
     186             :         q->_makeRed();
     187             :         q->_getLeft()->_makeBlack();
     188             :         q->_getRight()->_makeBlack();
     189             :       }
     190             : 
     191             :       // Fix red violation.
     192             :       if (_isValidRed(q) && _isValidRed(p)) {
     193           0 :         ASMJIT_ASSUME(g != nullptr);
     194             :         ASMJIT_ASSUME(p != nullptr);
     195           0 :         t->_setChild(t->_getRight() == g,
     196           0 :                      q == p->_getChild(last) ? _singleRotate(g, !last) : _doubleRotate(g, !last));
     197             :       }
     198             : 
     199             :       // Stop if found.
     200           0 :       if (q == node)
     201             :         break;
     202             : 
     203             :       last = dir;
     204           0 :       dir = cmp(*static_cast<NodeT*>(q), *static_cast<NodeT*>(node)) < 0;
     205             : 
     206             :       // Update helpers.
     207           0 :       if (g) t = g;
     208             : 
     209             :       g = p;
     210             :       p = q;
     211             :       q = q->_getChild(dir);
     212             :     }
     213             : 
     214             :     // Update root and make it black.
     215           0 :     _root = static_cast<NodeT*>(head._getRight());
     216             :     _root->_makeBlack();
     217             :   }
     218             : 
     219             :   //! Remove node from RBTree.
     220             :   template<typename CompareT = Support::Compare<Support::SortOrder::kAscending>>
     221           0 :   void remove(ZoneTreeNode* ASMJIT_NONNULL(node), const CompareT& cmp = CompareT()) noexcept {
     222             :     ZoneTreeNode head;           // False root node,
     223           0 :     head._setRight(_root);       // having root on the right.
     224             : 
     225             :     ZoneTreeNode* g = nullptr;   // Grandparent.
     226             :     ZoneTreeNode* p = nullptr;   // Parent.
     227             :     ZoneTreeNode* q = &head;     // Query.
     228             : 
     229             :     ZoneTreeNode* f  = nullptr;  // Found item.
     230             :     ZoneTreeNode* gf = nullptr;  // Found grandparent.
     231             :     size_t dir = 1;              // Direction (0 or 1).
     232             : 
     233             :     // Search and push a red down.
     234           0 :     while (q->hasChild(dir)) {
     235             :       size_t last = dir;
     236             : 
     237             :       // Update helpers.
     238             :       g = p;
     239             :       p = q;
     240             :       q = q->_getChild(dir);
     241           0 :       dir = cmp(*static_cast<NodeT*>(q), *static_cast<NodeT*>(node)) < 0;
     242             : 
     243             :       // Save found node.
     244           0 :       if (q == node) {
     245             :         f = q;
     246             :         gf = g;
     247             :       }
     248             : 
     249             :       // Push the red node down.
     250             :       if (!_isValidRed(q) && !_isValidRed(q->_getChild(dir))) {
     251           0 :         if (_isValidRed(q->_getChild(!dir))) {
     252           0 :           ZoneTreeNode* child = _singleRotate(q, dir);
     253             :           p->_setChild(last, child);
     254             :           p = child;
     255             :         }
     256           0 :         else if (!_isValidRed(q->_getChild(!dir)) && p->_getChild(!last)) {
     257             :           ZoneTreeNode* s = p->_getChild(!last);
     258             :           if (!_isValidRed(s->_getChild(!last)) && !_isValidRed(s->_getChild(last))) {
     259             :             // Color flip.
     260             :             p->_makeBlack();
     261             :             s->_makeRed();
     262             :             q->_makeRed();
     263             :           }
     264             :           else {
     265           0 :             ASMJIT_ASSUME(g != nullptr);
     266             :             ASMJIT_ASSUME(s != nullptr);
     267             : 
     268           0 :             size_t dir2 = g->_getRight() == p;
     269             :             ZoneTreeNode* child = g->_getChild(dir2);
     270             : 
     271             :             if (_isValidRed(s->_getChild(last))) {
     272           0 :               child = _doubleRotate(p, last);
     273             :               g->_setChild(dir2, child);
     274             :             }
     275             :             else if (_isValidRed(s->_getChild(!last))) {
     276           0 :               child = _singleRotate(p, last);
     277             :               g->_setChild(dir2, child);
     278             :             }
     279             : 
     280             :             // Ensure correct coloring.
     281             :             q->_makeRed();
     282             :             child->_makeRed();
     283             :             child->_getLeft()->_makeBlack();
     284             :             child->_getRight()->_makeBlack();
     285             :           }
     286             :         }
     287             :       }
     288             :     }
     289             : 
     290             :     // Replace and remove.
     291             :     ASMJIT_ASSERT(f != nullptr);
     292             :     ASMJIT_ASSERT(f != &head);
     293             :     ASMJIT_ASSERT(q != &head);
     294             : 
     295           0 :     p->_setChild(p->_getRight() == q,
     296             :                  q->_getChild(q->_getLeft() == nullptr));
     297             : 
     298             :     // NOTE: The original algorithm used a trick to just copy 'key/value' to `f` and mark `q` for deletion. But this
     299             :     // is unacceptable here as we really want to destroy the passed `node`. So, we have to make sure that we have
     300             :     // really removed `f` and not `q`.
     301           0 :     if (f != q) {
     302             :       ASMJIT_ASSERT(f != &head);
     303             :       ASMJIT_ASSERT(f != gf);
     304             : 
     305           0 :       ZoneTreeNode* n = gf ? gf : &head;
     306           0 :       dir = (n == &head) ? 1  : cmp(*static_cast<NodeT*>(n), *static_cast<NodeT*>(node)) < 0;
     307             : 
     308             :       for (;;) {
     309           0 :         if (n->_getChild(dir) == f) {
     310             :           n->_setChild(dir, q);
     311             :           // RAW copy, including the color.
     312           0 :           q->_rbNodeData[0] = f->_rbNodeData[0];
     313           0 :           q->_rbNodeData[1] = f->_rbNodeData[1];
     314             :           break;
     315             :         }
     316             : 
     317             :         n = n->_getChild(dir);
     318             : 
     319             :         // Cannot be true as we know that it must reach `f` in few iterations.
     320             :         ASMJIT_ASSERT(n != nullptr);
     321           0 :         dir = cmp(*static_cast<NodeT*>(n), *static_cast<NodeT*>(node)) < 0;
     322             :       }
     323             :     }
     324             : 
     325             :     // Update root and make it black.
     326           0 :     _root = static_cast<NodeT*>(head._getRight());
     327           0 :     if (_root) _root->_makeBlack();
     328           0 :   }
     329             : 
     330             :   template<typename KeyT, typename CompareT = Support::Compare<Support::SortOrder::kAscending>>
     331         100 :   inline NodeT* get(const KeyT& key, const CompareT& cmp = CompareT()) const noexcept {
     332         100 :     ZoneTreeNode* node = _root;
     333         100 :     while (node) {
     334           0 :       auto result = cmp(*static_cast<const NodeT*>(node), key);
     335         100 :       if (result == 0) break;
     336             : 
     337             :       // Go left or right depending on the `result`.
     338           0 :       node = node->_getChild(result < 0);
     339             :     }
     340         100 :     return static_cast<NodeT*>(node);
     341             :   }
     342             : 
     343             :   //! \}
     344             : 
     345             :   //! \cond INTERNAL
     346             :   //! \name Internal
     347             :   //! \{
     348             : 
     349             :   static inline bool _isValidRed(ZoneTreeNode* node) noexcept { return ZoneTreeNode::_isValidRed(node); }
     350             : 
     351             :   //! Single rotation.
     352           0 :   static inline ZoneTreeNode* _singleRotate(ZoneTreeNode* ASMJIT_NONNULL(root), size_t dir) noexcept {
     353           0 :     ZoneTreeNode* save = root->_getChild(!dir);
     354           0 :     ASMJIT_ASSUME(save != nullptr);
     355             : 
     356             :     ZoneTreeNode* saveChild = save->_getChild(dir);
     357             :     root->_setChild(!dir, saveChild);
     358             :     save->_setChild( dir, root);
     359             :     root->_makeRed();
     360             :     save->_makeBlack();
     361           0 :     return save;
     362             :   }
     363             : 
     364             :   //! Double rotation.
     365           0 :   static inline ZoneTreeNode* _doubleRotate(ZoneTreeNode* ASMJIT_NONNULL(root), size_t dir) noexcept {
     366           0 :     ZoneTreeNode* child = root->_getChild(!dir);
     367           0 :     ASMJIT_ASSUME(child != nullptr);
     368             : 
     369           0 :     root->_setChild(!dir, _singleRotate(child, !dir));
     370           0 :     return _singleRotate(root, dir);
     371             :   }
     372             : 
     373             :   //! \}
     374             :   //! \endcond
     375             : };
     376             : 
     377             : //! \}
     378             : 
     379             : ASMJIT_END_NAMESPACE
     380             : 
     381             : #endif // ASMJIT_CORE_ZONETREE_H_INCLUDED
     382             : 
     383             : #endif // __PLUMED_HAS_ASMJIT