1 files changed, 73 insertions, 0 deletions
diff --git a/libcxx/test/support/test_allocator.h b/libcxx/test/support/test_allocator.h
index 394d1ccd450..41473371500 100644
--- a/libcxx/test/support/test_allocator.h
+++ b/libcxx/test/support/test_allocator.h
@@ -302,5 +302,78 @@ operator!=(const TaggingAllocator<T>&, const TaggingAllocator<U>&)
 { return false; }
 #endif
 
+template <std::size_t MaxAllocs>
+struct limited_alloc_handle {
+  std::size_t outstanding_;
+  void* last_alloc_;
+
+  limited_alloc_handle() : outstanding_(0), last_alloc_(nullptr) {}
+
+  template <class T>
+  T *allocate(std::size_t N) {
+    if (N + outstanding_ > MaxAllocs)
+      TEST_THROW(std::bad_alloc());
+    last_alloc_ = ::operator new(N*sizeof(T));
+    outstanding_ += N;
+    return static_cast<T*>(last_alloc_);
+  }
+
+  void deallocate(void* ptr, std::size_t N) {
+    if (ptr == last_alloc_) {
+      last_alloc_ = nullptr;
+      assert(outstanding_ >= N);
+      outstanding_ -= N;
+    }
+    ::operator delete(ptr);
+  }
+};
+
+template <class T, std::size_t N>
+class limited_allocator
+{
+    typedef limited_alloc_handle<N> BuffT;
+    std::shared_ptr<BuffT> handle_;
+public:
+    typedef T                 value_type;
+    typedef value_type*       pointer;
+    typedef const value_type* const_pointer;
+    typedef value_type&       reference;
+    typedef const value_type& const_reference;
+    typedef std::size_t       size_type;
+    typedef std::ptrdiff_t    difference_type;
+
+    template <class U> struct rebind { typedef limited_allocator<U, N> other; };
+
+    limited_allocator() : handle_(new BuffT) {}
+
+    limited_allocator(limited_allocator const& other) : handle_(other.handle_) {}
+
+    template <class U>
+    explicit limited_allocator(limited_allocator<U, N> const& other)
+        : handle_(other.handle_) {}
+
+private:
+    limited_allocator& operator=(const limited_allocator&);// = delete;
+
+public:
+    pointer allocate(size_type n) { return handle_->template allocate<T>(n); }
+    void deallocate(pointer p, size_type n) { handle_->deallocate(p, n); }
+    size_type max_size() const {return N;}
+
+    BuffT* getHandle() const { return handle_.get(); }
+};
+
+template <class T, class U, std::size_t N>
+inline bool operator==(limited_allocator<T, N> const& LHS,
+                       limited_allocator<U, N> const& RHS) {
+  return LHS.getHandle() == RHS.getHandle();
+}
+
+template <class T, class U, std::size_t N>
+inline bool operator!=(limited_allocator<T, N> const& LHS,
+                       limited_allocator<U, N> const& RHS) {
+  return !(LHS == RHS);
+}
+
 
 #endif  // TEST_ALLOCATOR_H