michael@0: // Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
michael@0: // Use of this source code is governed by a BSD-style license that can be
michael@0: // found in the LICENSE file.
michael@0: 
michael@0: #ifndef BASE_STACK_CONTAINER_H_
michael@0: #define BASE_STACK_CONTAINER_H_
michael@0: 
michael@0: #include <string>
michael@0: #include <vector>
michael@0: 
michael@0: #include "base/basictypes.h"
michael@0: 
michael@0: // This allocator can be used with STL containers to provide a stack buffer
michael@0: // from which to allocate memory and overflows onto the heap. This stack buffer
michael@0: // would be allocated on the stack and allows us to avoid heap operations in
michael@0: // some situations.
michael@0: //
michael@0: // STL likes to make copies of allocators, so the allocator itself can't hold
michael@0: // the data. Instead, we make the creator responsible for creating a
michael@0: // StackAllocator::Source which contains the data. Copying the allocator
michael@0: // merely copies the pointer to this shared source, so all allocators created
michael@0: // based on our allocator will share the same stack buffer.
michael@0: //
michael@0: // This stack buffer implementation is very simple. The first allocation that
michael@0: // fits in the stack buffer will use the stack buffer. Any subsequent
michael@0: // allocations will not use the stack buffer, even if there is unused room.
michael@0: // This makes it appropriate for array-like containers, but the caller should
michael@0: // be sure to reserve() in the container up to the stack buffer size. Otherwise
michael@0: // the container will allocate a small array which will "use up" the stack
michael@0: // buffer.
michael@0: template<typename T, size_t stack_capacity>
michael@0: class StackAllocator : public std::allocator<T> {
michael@0:  public:
michael@0:   typedef typename std::allocator<T>::pointer pointer;
michael@0:   typedef typename std::allocator<T>::size_type size_type;
michael@0: 
michael@0:   // Backing store for the allocator. The container owner is responsible for
michael@0:   // maintaining this for as long as any containers using this allocator are
michael@0:   // live.
michael@0:   struct Source {
michael@0:     Source() : used_stack_buffer_(false) {
michael@0:     }
michael@0: 
michael@0:     // Casts the buffer in its right type.
michael@0:     T* stack_buffer() { return reinterpret_cast<T*>(stack_buffer_); }
michael@0:     const T* stack_buffer() const {
michael@0:       return reinterpret_cast<const T*>(stack_buffer_);
michael@0:     }
michael@0: 
michael@0:     //
michael@0:     // IMPORTANT: Take care to ensure that stack_buffer_ is aligned
michael@0:     // since it is used to mimic an array of T.
michael@0:     // Be careful while declaring any unaligned types (like bool)
michael@0:     // before stack_buffer_.
michael@0:     //
michael@0: 
michael@0:     // The buffer itself. It is not of type T because we don't want the
michael@0:     // constructors and destructors to be automatically called. Define a POD
michael@0:     // buffer of the right size instead.
michael@0:     char stack_buffer_[sizeof(T[stack_capacity])];
michael@0: 
michael@0:     // Set when the stack buffer is used for an allocation. We do not track
michael@0:     // how much of the buffer is used, only that somebody is using it.
michael@0:     bool used_stack_buffer_;
michael@0:   };
michael@0: 
michael@0:   // Used by containers when they want to refer to an allocator of type U.
michael@0:   template<typename U>
michael@0:   struct rebind {
michael@0:     typedef StackAllocator<U, stack_capacity> other;
michael@0:   };
michael@0: 
michael@0:   // For the straight up copy c-tor, we can share storage.
michael@0:   StackAllocator(const StackAllocator<T, stack_capacity>& rhs)
michael@0:       : source_(rhs.source_) {
michael@0:   }
michael@0: 
michael@0:   // ISO C++ requires the following constructor to be defined,
michael@0:   // and std::vector in VC++2008SP1 Release fails with an error
michael@0:   // in the class _Container_base_aux_alloc_real (from <xutility>)
michael@0:   // if the constructor does not exist.
michael@0:   // For this constructor, we cannot share storage; there's
michael@0:   // no guarantee that the Source buffer of Ts is large enough
michael@0:   // for Us.
michael@0:   // TODO: If we were fancy pants, perhaps we could share storage
michael@0:   // iff sizeof(T) == sizeof(U).
michael@0:   template<typename U, size_t other_capacity>
michael@0:   StackAllocator(const StackAllocator<U, other_capacity>& other)
michael@0:       : source_(NULL) {
michael@0:   }
michael@0: 
michael@0:   explicit StackAllocator(Source* source) : source_(source) {
michael@0:   }
michael@0: 
michael@0:   // Actually do the allocation. Use the stack buffer if nobody has used it yet
michael@0:   // and the size requested fits. Otherwise, fall through to the standard
michael@0:   // allocator.
michael@0:   pointer allocate(size_type n, void* hint = 0) {
michael@0:     if (source_ != NULL && !source_->used_stack_buffer_
michael@0:         && n <= stack_capacity) {
michael@0:       source_->used_stack_buffer_ = true;
michael@0:       return source_->stack_buffer();
michael@0:     } else {
michael@0:       return std::allocator<T>::allocate(n, hint);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   // Free: when trying to free the stack buffer, just mark it as free. For
michael@0:   // non-stack-buffer pointers, just fall though to the standard allocator.
michael@0:   void deallocate(pointer p, size_type n) {
michael@0:     if (source_ != NULL && p == source_->stack_buffer())
michael@0:       source_->used_stack_buffer_ = false;
michael@0:     else
michael@0:       std::allocator<T>::deallocate(p, n);
michael@0:   }
michael@0: 
michael@0:  private:
michael@0:   Source* source_;
michael@0: };
michael@0: 
michael@0: // A wrapper around STL containers that maintains a stack-sized buffer that the
michael@0: // initial capacity of the vector is based on. Growing the container beyond the
michael@0: // stack capacity will transparently overflow onto the heap. The container must
michael@0: // support reserve().
michael@0: //
michael@0: // WATCH OUT: the ContainerType MUST use the proper StackAllocator for this
michael@0: // type. This object is really intended to be used only internally. You'll want
michael@0: // to use the wrappers below for different types.
michael@0: template<typename TContainerType, int stack_capacity>
michael@0: class StackContainer {
michael@0:  public:
michael@0:   typedef TContainerType ContainerType;
michael@0:   typedef typename ContainerType::value_type ContainedType;
michael@0:   typedef StackAllocator<ContainedType, stack_capacity> Allocator;
michael@0: 
michael@0:   // Allocator must be constructed before the container!
michael@0:   StackContainer() : allocator_(&stack_data_), container_(allocator_) {
michael@0:     // Make the container use the stack allocation by reserving our buffer size
michael@0:     // before doing anything else.
michael@0:     container_.reserve(stack_capacity);
michael@0:   }
michael@0: 
michael@0:   // Getters for the actual container.
michael@0:   //
michael@0:   // Danger: any copies of this made using the copy constructor must have
michael@0:   // shorter lifetimes than the source. The copy will share the same allocator
michael@0:   // and therefore the same stack buffer as the original. Use std::copy to
michael@0:   // copy into a "real" container for longer-lived objects.
michael@0:   ContainerType& container() { return container_; }
michael@0:   const ContainerType& container() const { return container_; }
michael@0: 
michael@0:   // Support operator-> to get to the container. This allows nicer syntax like:
michael@0:   //   StackContainer<...> foo;
michael@0:   //   std::sort(foo->begin(), foo->end());
michael@0:   ContainerType* operator->() { return &container_; }
michael@0:   const ContainerType* operator->() const { return &container_; }
michael@0: 
michael@0: #ifdef UNIT_TEST
michael@0:   // Retrieves the stack source so that that unit tests can verify that the
michael@0:   // buffer is being used properly.
michael@0:   const typename Allocator::Source& stack_data() const {
michael@0:     return stack_data_;
michael@0:   }
michael@0: #endif
michael@0: 
michael@0:  protected:
michael@0:   typename Allocator::Source stack_data_;
michael@0:   Allocator allocator_;
michael@0:   ContainerType container_;
michael@0: 
michael@0:   DISALLOW_EVIL_CONSTRUCTORS(StackContainer);
michael@0: };
michael@0: 
michael@0: // StackString
michael@0: template<size_t stack_capacity>
michael@0: class StackString : public StackContainer<
michael@0:     std::basic_string<char,
michael@0:                       std::char_traits<char>,
michael@0:                       StackAllocator<char, stack_capacity> >,
michael@0:     stack_capacity> {
michael@0:  public:
michael@0:   StackString() : StackContainer<
michael@0:       std::basic_string<char,
michael@0:                         std::char_traits<char>,
michael@0:                         StackAllocator<char, stack_capacity> >,
michael@0:       stack_capacity>() {
michael@0:   }
michael@0: 
michael@0:  private:
michael@0:   DISALLOW_EVIL_CONSTRUCTORS(StackString);
michael@0: };
michael@0: 
michael@0: // StackWString
michael@0: template<size_t stack_capacity>
michael@0: class StackWString : public StackContainer<
michael@0:     std::basic_string<wchar_t,
michael@0:                       std::char_traits<wchar_t>,
michael@0:                       StackAllocator<wchar_t, stack_capacity> >,
michael@0:     stack_capacity> {
michael@0:  public:
michael@0:   StackWString() : StackContainer<
michael@0:       std::basic_string<wchar_t,
michael@0:                         std::char_traits<wchar_t>,
michael@0:                         StackAllocator<wchar_t, stack_capacity> >,
michael@0:       stack_capacity>() {
michael@0:   }
michael@0: 
michael@0:  private:
michael@0:   DISALLOW_EVIL_CONSTRUCTORS(StackWString);
michael@0: };
michael@0: 
michael@0: // StackVector
michael@0: //
michael@0: // Example:
michael@0: //   StackVector<int, 16> foo;
michael@0: //   foo->push_back(22);  // we have overloaded operator->
michael@0: //   foo[0] = 10;         // as well as operator[]
michael@0: template<typename T, size_t stack_capacity>
michael@0: class StackVector : public StackContainer<
michael@0:     std::vector<T, StackAllocator<T, stack_capacity> >,
michael@0:     stack_capacity> {
michael@0:  public:
michael@0:   StackVector() : StackContainer<
michael@0:       std::vector<T, StackAllocator<T, stack_capacity> >,
michael@0:       stack_capacity>() {
michael@0:   }
michael@0: 
michael@0:   // We need to put this in STL containers sometimes, which requires a copy
michael@0:   // constructor. We can't call the regular copy constructor because that will
michael@0:   // take the stack buffer from the original. Here, we create an empty object
michael@0:   // and make a stack buffer of its own.
michael@0:   StackVector(const StackVector<T, stack_capacity>& other)
michael@0:       : StackContainer<
michael@0:             std::vector<T, StackAllocator<T, stack_capacity> >,
michael@0:             stack_capacity>() {
michael@0:     this->container().assign(other->begin(), other->end());
michael@0:   }
michael@0: 
michael@0:   StackVector<T, stack_capacity>& operator=(
michael@0:       const StackVector<T, stack_capacity>& other) {
michael@0:     this->container().assign(other->begin(), other->end());
michael@0:     return *this;
michael@0:   }
michael@0: 
michael@0:   // Vectors are commonly indexed, which isn't very convenient even with
michael@0:   // operator-> (using "->at()" does exception stuff we don't want).
michael@0:   T& operator[](size_t i) { return this->container().operator[](i); }
michael@0:   const T& operator[](size_t i) const {
michael@0:     return this->container().operator[](i);
michael@0:   }
michael@0: };
michael@0: 
michael@0: #endif  // BASE_STACK_CONTAINER_H_