command_buffer: Implement IdAllocator by recording ranges in an AVL tree

gpu/command_buffer/common/id_allocator.cc

Index: gpu/command_buffer/common/id_allocator.cc
diff --git a/gpu/command_buffer/common/id_allocator.cc b/gpu/command_buffer/common/id_allocator.cc
index 61d2c4a91135381921171a4d1fd0f418fbf77d06..29ec66ae2d92443a8832efd17d794283a1825b23 100644
--- a/gpu/command_buffer/common/id_allocator.cc
+++ b/gpu/command_buffer/common/id_allocator.cc
@@ -6,155 +6,401 @@
 
 #include "gpu/command_buffer/common/id_allocator.h"
 
+#include <limits>
+
 #include "base/logging.h"
 
 namespace gpu {
 
-IdAllocator::IdAllocator() {}
+// IdAllocator stores the list of allocated ids in an AVL tree. The
+// tree node |IdAllocator::Node| represents maximal closed interval
+// (min, max) of allocated ids. The tree is ordered by the range
+// minimum value.
 
-IdAllocator::~IdAllocator() {}
+struct IdAllocator::Node {
+  ResourceId min;
+  ResourceId max;
+  Node* left;
+  Node* right;
+  int height;
 
-ResourceId IdAllocator::AllocateID() {
-  ResourceId id;
-  ResourceIdSet::iterator iter = free_ids_.begin();
-  if (iter != free_ids_.end()) {
-    id = *iter;
-  } else {
-    id = LastUsedId() + 1;
-    if (!id) {
-      // We wrapped around to 0.
-      id = FindFirstUnusedId();
-    }
+  Node(ResourceId min, ResourceId max)
+      : min(min), max(max), left(NULL), right(NULL), height(0) {}
+};
+
+IdAllocator::IdAllocator() : root_(NULL) {
+}
+
+IdAllocator::~IdAllocator() {
+  if (root_ != NULL) {
+    Destroy(root_);
   }
-  MarkAsUsed(id);
-  return id;
+}
+
+ResourceId IdAllocator::AllocateID() {
+  return AllocateIDRange(1);
 }
 
 ResourceId IdAllocator::AllocateIDAtOrAbove(ResourceId desired_id) {
-  ResourceId id;
-  ResourceIdSet::iterator iter = free_ids_.lower_bound(desired_id);
-  if (iter != free_ids_.end()) {
-    id = *iter;
-  } else if (LastUsedId() < desired_id) {
-    id = desired_id;
-  } else {
-    id = LastUsedId() + 1;
-    if (!id) {
-      // We wrapped around to 0.
-      id = FindFirstUnusedId();
-    }
+  if (desired_id == 0) {
+    desired_id = 1;
+  }
+
+  if (root_ == NULL) {
+    root_ = Insert(desired_id, desired_id, root_);
+    return desired_id;
+  }
+
+  ResourceId result = 0;
+  root_ = AddRange(1,
+                   &result,
+                   desired_id,
+                   std::numeric_limits<ResourceId>::max(),
+                   root_,
+                   NULL,
+                   NULL);
+
+  if (result == 0) {
+    root_ = AddRange(1,
+                     &result,
+                     1,
+                     std::numeric_limits<ResourceId>::max(),
+                     root_,
+                     NULL,
+                     NULL);
   }
-  MarkAsUsed(id);
-  return id;
+
+  return result;
 }
 
 ResourceId IdAllocator::AllocateIDRange(uint32_t range) {
-  DCHECK(range > 0);
-  if (range == 1) {
-    return AllocateID();
-  }
-
-  ResourceId first_id = kInvalidResource;
-
-  {
-    // Try to re-use ids from free_ids.
-    ResourceIdSet::const_iterator iter = free_ids_.begin();
-    ResourceIdSet::const_iterator end = free_ids_.end();
-    const ResourceIdSet::size_type size = free_ids_.size();
-    if (iter != end && size >= range) {
-      ResourceId first = *iter;
-      ResourceIdSet::size_type needed_size = range;
-      ResourceId prev = first;
-      ++iter;
-      for (; iter != end && size >= needed_size; ++iter) {
-        if (prev + 1 == *iter) {
-          if (first + range - 1 == *iter) {
-            first_id = first;
-            break;
-          }
-        } else {
-          needed_size += prev - first + 1;
-          first = *iter;
-        }
-        prev = *iter;
-      }
-    }
+  if (root_ == NULL) {
+    root_ = Insert(1, range, NULL);
+    return 1;
   }
 
-  if (first_id == kInvalidResource) {
-    // Allocate after the last id.
-    ResourceId first = LastUsedId() + 1;
-    if (first != 0) {
-      ResourceId last = first + range - 1;
-      if (first < last) {
-        first_id = first;
-      }
+  ResourceId result = 0;
+  root_ = AddRange(range,
+                   &result,
+                   1,
+                   std::numeric_limits<ResourceId>::max(),
+                   root_,
+                   NULL,
+                   NULL);
+  return result;
+}
+
+bool IdAllocator::MarkAsUsed(ResourceId id) {
+  if (root_ == NULL) {
+    root_ = Insert(id, id, NULL);
+    return true;
+  }
+
+  ResourceId result = 0;
+  root_ = AddRange(1, &result, id, id, root_, NULL, NULL);
+  return result != 0;
+}
+
+void IdAllocator::FreeID(ResourceId id) {
+  FreeIDRange(id, 1);
+}
+
+void IdAllocator::FreeIDRange(ResourceId first_id, uint32_t range) {
+  if (root_ == NULL) {
+    return;
+  }
+
+  root_ = RemoveRange(first_id, first_id + range - 1u, root_);
+  DCHECK(root_ == NULL || root_->min > 0);
+}
+
+bool IdAllocator::InUse(ResourceId id) const {
+  if (id == 0) {
+    return false;
+  }
+
+  Node* pos = root_;
+  while (pos != NULL) {
+    if (id >= pos->min && id <= pos->max) {
+      return true;
+    }
+
+    if (id < pos->min) {
+      pos = pos->left;
+    } else {
+      pos = pos->right;
     }
   }
 
-  if (first_id == kInvalidResource) {
-    // Scan the gaps between used ids.
-    ResourceId prev = 0;
-    for (auto id : used_ids_) {
-      if (id - prev > range) {
-        first_id = prev + 1;
-        break;
-      }
-      prev = id;
+  return false;
+}
+
+inline int IdAllocator::GetHeight(Node* node) {
+  return node ? node->height : 0;
+}
+
+inline void IdAllocator::UpdateHeight(Node* node) {
+  node->height = std::max(GetHeight(node->left), GetHeight(node->right)) + 1;
+}
+
+inline int IdAllocator::GetBalance(Node* node) {
+  return node ? GetHeight(node->left) - GetHeight(node->right) : 0;
+}
+
+inline IdAllocator::Node* IdAllocator::RotateLeft(Node* node) {
+  Node* new_root = node->right;
+  node->right = new_root->left;
+  new_root->left = node;
+  UpdateHeight(node);
+  UpdateHeight(new_root);
+  return new_root;
+}
+
+inline IdAllocator::Node* IdAllocator::RotateRight(Node* node) {
+  Node* new_root = node->left;
+  node->left = new_root->right;
+  new_root->right = node;
+  UpdateHeight(node);
+  UpdateHeight(new_root);
+  return new_root;
+}
+
+IdAllocator::Node* IdAllocator::Balance(Node* node) {
+  UpdateHeight(node);
+  int balance = GetBalance(node);
+  if (balance > 1) {
+    if (GetBalance(node->left) >= 0) {
+      return RotateRight(node);
     }
+    node->left = RotateLeft(node->left);
+    return RotateRight(node);
   }
 
-  if (first_id != kInvalidResource) {
-    for (uint32 ii = 0; ii < range; ++ii) {
-      MarkAsUsed(first_id + ii);
+  if (balance < -1) {
+    if (GetBalance(node->right) <= 0) {
+      return RotateLeft(node);
     }
+    node->right = RotateRight(node->right);
+    return RotateLeft(node);
   }
 
-  return first_id;
+  return node;
 }
 
-bool IdAllocator::MarkAsUsed(ResourceId id) {
-  DCHECK(id);
-  free_ids_.erase(id);
-  std::pair<ResourceIdSet::iterator, bool> result = used_ids_.insert(id);
-  return result.second;
+// Finds the largest node of |max_subtree| and stores its
+// contents to |to_replace|. Deletes the found node.
+IdAllocator::Node* IdAllocator::ReplaceMax(Node* to_replace,
+                                           Node* max_subtree) {
+  if (max_subtree->right != NULL) {
+    max_subtree->right = ReplaceMax(to_replace, max_subtree->right);
+    return Balance(max_subtree);
+  }
+
+  to_replace->min = max_subtree->min;
+  to_replace->max = max_subtree->max;
+  Node* result = max_subtree->left;
+  delete max_subtree;
+  return result;
 }
 
-void IdAllocator::FreeID(ResourceId id) {
-  if (id) {
-    used_ids_.erase(id);
-    free_ids_.insert(id);
+IdAllocator::Node* IdAllocator::Delete(Node* node) {
+  if (node->left != NULL && node->right != NULL) {
+    node->left = ReplaceMax(node, node->left);
+    return Balance(node);
   }
+
+  Node* result = node->left != NULL ? node->left : node->right;
+  delete node;
+  return result;
 }
 
-void IdAllocator::FreeIDRange(ResourceId first_id, uint32 range) {
-  for (uint32 ii = 0; ii < range; ++ii) {
-    FreeID(first_id + ii);
+IdAllocator::Node* IdAllocator::Insert(ResourceId min,
+                                       ResourceId max,
+                                       Node* pos) {
+  if (pos == NULL) {
+    return new Node(min, max);
   }
-}
 
-bool IdAllocator::InUse(ResourceId id) const {
-  return id == kInvalidResource || used_ids_.find(id) != used_ids_.end();
+  if (min < pos->min) {
+    pos->left = Insert(min, max, pos->left);
+  } else {
+    pos->right = Insert(min, max, pos->right);
+  }
+
+  return Balance(pos);
 }
 
-ResourceId IdAllocator::LastUsedId() const {
-  if (used_ids_.empty()) {
-    return 0u;
+// Allocates |range| amount of ids and stores the first id allocated
+// in |first_id_out|. In case of failed allocation, |first_id_out|
+// remains 0. The allocated id range will start at the smallest possible
+// value between closed range [|first_id_min|, |first_id_max|].
+
+// Works by recursing through the tree. When taking the left branch
+// (searching the smaller id space), |right_limit| is updated. When taking
+// the right branch (searching the larger id space), |left_limit| is updated.
+// When the id space is found, uses |left_limit| or |right_limit| to identify
+// whether the new id range needs to be merged at the end of the limiter or the
+// current position in the tree. Additional option the special case where passed
+// in
+// |first_id_min| causes a new range to be created.
+IdAllocator::Node* IdAllocator::AddRange(uint32_t range,
+                                         ResourceId* first_id_out,
+                                         ResourceId first_id_min,
+                                         ResourceId first_id_max,
+                                         Node* node,
+                                         Node* left_limit,
+                                         Node* right_limit) {
+  DCHECK(*first_id_out == 0);
+
+  ResourceId new_range_max = first_id_min + range - 1u;
+  if (new_range_max < first_id_min) {
+    return node;
+  }
+
+  if (node->min > new_range_max) {
+    if (node->left != NULL) {
+      node->left = AddRange(range,
+                            first_id_out,
+                            first_id_min,
+                            first_id_max,
+                            node->left,
+                            left_limit,
+                            node);
+      if (*first_id_out != 0) {
+        return Balance(node);
+      }
+    } else {
+      *first_id_out = first_id_min;
+      bool merges_node = node->min - new_range_max == 1u;
+      bool merges_left_limit =
+          left_limit != NULL && (first_id_min - left_limit->max == 1u);
+      if (merges_node && merges_left_limit) {
+        left_limit->max = node->max;
+        return Delete(node);
+      } else if (merges_node) {
+        node->min = first_id_min;
+        return node;
+      } else if (merges_left_limit) {
+        left_limit->max = new_range_max;
+        return node;
+      }
+      return Insert(first_id_min, new_range_max, node);
+    }
+  }
+
+  if (first_id_min <= node->max) {
+    first_id_min = node->max + 1u;
+    if (first_id_min < node->max) {
+      return node;
+    }
+  }
+
+  if (first_id_min > first_id_max) {
+    return node;
+  }
+
+  if (node->right != NULL) {
+    node->right = AddRange(range,
+                           first_id_out,
+                           first_id_min,
+                           first_id_max,
+                           node->right,
+                           node,
+                           right_limit);
+    if (*first_id_out != 0) {
+      return Balance(node);
+    }
+    return node;
+  }
+
+  new_range_max = first_id_min + range - 1u;
+  if (new_range_max < first_id_min) {
+    return node;
+  }
+
+  if (right_limit != NULL) {
+    if (new_range_max < right_limit->min) {
+      *first_id_out = first_id_min;
+      if (right_limit->min - new_range_max == 1u) {
+        right_limit->min = node->min;
+        return Delete(node);
+      }
+      node->max = new_range_max;
+      return node;
+    }
   } else {
-    return *used_ids_.rbegin();
+    *first_id_out = first_id_min;
+    if (first_id_min - node->max == 1u) {
+      node->max = new_range_max;
+      return node;
+    }
+    return Insert(first_id_min, new_range_max, node);
   }
+
+  return node;
 }
 
-ResourceId IdAllocator::FindFirstUnusedId() const {
-  ResourceId id = 1;
-  for (ResourceIdSet::const_iterator it = used_ids_.begin();
-       it != used_ids_.end(); ++it) {
-    if ((*it) != id) {
-      return id;
+IdAllocator::Node* IdAllocator::RemoveRange(ResourceId first_id,
+                                            ResourceId last_id,
+                                            Node* pos) {
+  if (pos == NULL) {
+    return NULL;
+  }
+
+  if (first_id < pos->min) {
+    pos->left = RemoveRange(first_id, std::min(pos->min, last_id), pos->left);
+  }
+
+  if (last_id > pos->max) {
+    pos->right = RemoveRange(std::max(pos->max, first_id), last_id, pos->right);
+  }
+
+  if (first_id <= pos->min) {
+    if (last_id >= pos->min) {
+      ResourceId old_max = pos->max;
+      if (last_id < old_max) {
+        pos->min = last_id + 1u;
+        pos->max = old_max;
+      } else {
+        pos = Delete(pos);
+      }
+    }
+  } else if (first_id <= pos->max) {
+    ResourceId old_max = pos->max;
+    pos->max = first_id - 1u;
+    if (last_id < old_max) {
+      pos = Insert(last_id + 1u, old_max, pos);
+    }
+  }
+
+  if (pos == NULL) {
+    return NULL;
+  }
+
+  // The nodes |pos->left| and |pos->right| are both
+  // balanced. However, |pos| is potentially unbalanced. Since
+  // multiple operations might be done during update of |pos->left| and
+  // |pos->right|, we must balance |pos| multiple times.
+  while (true) {
+    int balance = GetBalance(pos);
+    if (balance >= -1 && balance <= 1) {
+      break;
     }
-    ++id;
+    pos = Balance(pos);
   }
-  return id;
+
+  return pos;
+}
+
+void IdAllocator::Destroy(Node* pos) {
+  do {
+    if (pos->left != NULL) {
+      Destroy(pos->left);
+    }
+    Node* right = pos->right;
+    delete pos;
+    pos = right;
+  } while (pos != NULL);
 }
 
 }  // namespace gpu