3 files changed, 64 insertions, 11 deletions

diff --git a/gpu/command_buffer/common/id_allocator.cc b/gpu/command_buffer/common/id_allocator.cc
index 68f7972..e859282 100644
--- a/gpu/command_buffer/common/id_allocator.cc
+++ b/gpu/command_buffer/common/id_allocator.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
@@ -14,33 +14,66 @@ IdAllocator::IdAllocator() {}
 IdAllocator::~IdAllocator() {}
 
 ResourceId IdAllocator::AllocateID() {
-  ResourceId id = FindFirstFree();
+  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();
+    }
+  }
   MarkAsUsed(id);
   return id;
 }
 
 ResourceId IdAllocator::AllocateIDAtOrAbove(ResourceId desired_id) {
-  GPU_DCHECK_LT(static_cast<ResourceId>(used_ids_.size()),
-            static_cast<ResourceId>(-1));
-  for (; InUse(desired_id); ++desired_id) {}
-  MarkAsUsed(desired_id);
-  return 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();
+    }
+  }
+  MarkAsUsed(id);
+  return id;
 }
 
 bool IdAllocator::MarkAsUsed(ResourceId id) {
+  GPU_DCHECK(id);
+  free_ids_.erase(id);
   std::pair<ResourceIdSet::iterator, bool> result = used_ids_.insert(id);
   return result.second;
 }
 
 void IdAllocator::FreeID(ResourceId id) {
+  GPU_DCHECK(id);
   used_ids_.erase(id);
+  std::pair<ResourceIdSet::iterator, bool> result = free_ids_.insert(id);
+  GPU_DCHECK(result.second);
 }
 
 bool IdAllocator::InUse(ResourceId id) const {
   return id == kInvalidResource || used_ids_.find(id) != used_ids_.end();
 }
 
-ResourceId IdAllocator::FindFirstFree() const {
+ResourceId IdAllocator::LastUsedId() const {
+  if (used_ids_.empty()) {
+    return 0u;
+  } else {
+    return *used_ids_.rbegin();
+  }
+}
+
+ResourceId IdAllocator::FindFirstUnusedId() const {
   ResourceId id = 1;
   for (ResourceIdSet::const_iterator it = used_ids_.begin();
        it != used_ids_.end(); ++it) {
diff --git a/gpu/command_buffer/common/id_allocator.h b/gpu/command_buffer/common/id_allocator.h
index 05c6af07..4e80290 100644
--- a/gpu/command_buffer/common/id_allocator.h
+++ b/gpu/command_buffer/common/id_allocator.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
@@ -44,9 +44,14 @@ class IdAllocator {
   // TODO(gman): This would work much better with ranges or a hash table.
   typedef std::set<ResourceId> ResourceIdSet;
 
-  ResourceId FindFirstFree() const;
+  // The highest ID on the used list.
+  ResourceId LastUsedId() const;
+
+  // Lowest ID that isn't on the used list. This is slow, use as a last resort.
+  ResourceId FindFirstUnusedId() const;
 
   ResourceIdSet used_ids_;
+  ResourceIdSet free_ids_;
 
   DISALLOW_COPY_AND_ASSIGN(IdAllocator);
 };
diff --git a/gpu/command_buffer/common/id_allocator_test.cc b/gpu/command_buffer/common/id_allocator_test.cc
index 6869f33..69be2f0 100644
--- a/gpu/command_buffer/common/id_allocator_test.cc
+++ b/gpu/command_buffer/common/id_allocator_test.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2009 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
@@ -49,6 +49,9 @@ TEST_F(IdAllocatorTest, TestBasic) {
 TEST_F(IdAllocatorTest, TestAdvanced) {
   IdAllocator *allocator = id_allocator();
 
+  // Allocate the highest possible ID, to make life awkward.
+  allocator->AllocateIDAtOrAbove(-1);
+
   // Allocate a significant number of resources.
   const unsigned int kNumResources = 100;
   ResourceId ids[kNumResources];
@@ -95,4 +98,16 @@ TEST_F(IdAllocatorTest, AllocateIdAtOrAbove) {
   EXPECT_GT(id3, kOffset);
 }
 
+// Checks that AllocateIdAtOrAbove wraps around at the maximum 32-bit value.
+TEST_F(IdAllocatorTest, AllocateIdAtOrAboveWrapsAround) {
+  const ResourceId kMaxPossibleOffset = -1;
+  IdAllocator* allocator = id_allocator();
+  ResourceId id1 = allocator->AllocateIDAtOrAbove(kMaxPossibleOffset);
+  EXPECT_EQ(kMaxPossibleOffset, id1);
+  ResourceId id2 = allocator->AllocateIDAtOrAbove(kMaxPossibleOffset);
+  EXPECT_EQ(1u, id2);
+  ResourceId id3 = allocator->AllocateIDAtOrAbove(kMaxPossibleOffset);
+  EXPECT_EQ(2u, id3);
+}
+
 }  // namespace gpu