// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef COBALT_MEDIA_MEDIA_BUFFER_ALLOCATOR_H_
#define COBALT_MEDIA_MEDIA_BUFFER_ALLOCATOR_H_

#include "nb/memory_pool.h"
#include "starboard/common/scoped_ptr.h"

namespace cobalt {
namespace media {

class MediaBufferAllocator {
 public:
  MediaBufferAllocator(void* pool, size_t main_pool_size,
                       size_t small_allocation_pool_size,
                       size_t small_allocation_threshold)
      : pool_(reinterpret_cast<uint8*>(pool)),
        main_pool_size_(main_pool_size),
        small_allocation_pool_size_(small_allocation_pool_size),
        small_allocation_threshold_(small_allocation_threshold),
        main_pool_(starboard::make_scoped_ptr(new nb::MemoryPool(
            pool_, main_pool_size_, true /* verify_full_capacity */))) {
    if (small_allocation_pool_size_ > 0u) {
      DCHECK_GT(small_allocation_threshold_, 0u);
      small_allocation_pool_.set(starboard::make_scoped_ptr(new nb::MemoryPool(
          pool_ + main_pool_size_, small_allocation_pool_size_,
          true /* verify_full_capacity */)));
    } else {
      DCHECK_EQ(small_allocation_pool_size_, 0u);
      DCHECK_EQ(small_allocation_threshold_, 0u);
    }
  }

  void* Allocate(size_t size, size_t alignment) {
    void* p = NULL;
    if (size < small_allocation_threshold_ &&
        small_allocation_pool_->is_valid()) {
      p = small_allocation_pool_->Allocate(size, alignment);
    }
    if (!p) {
      p = main_pool_->Allocate(size, alignment);
    }
    if (!p && small_allocation_pool_) {
      p = small_allocation_pool_->Allocate(size, alignment);
    }
    return p;
  }

  void Free(void* p) {
    if (p >= pool_ + main_pool_size_ && small_allocation_pool_->is_valid()) {
      DCHECK_LT(p, pool_ + main_pool_size_ + small_allocation_pool_size_);
      small_allocation_pool_->Free(p);
      return;
    }
    DCHECK_GE(p, pool_);
    DCHECK_LT(p, pool_ + main_pool_size_);
    main_pool_->Free(p);
  }

 private:
  uint8* pool_;
  size_t main_pool_size_;
  size_t small_allocation_pool_size_;
  size_t small_allocation_threshold_;

  starboard::LockedPtr<nb::MemoryPool> main_pool_;
  starboard::LockedPtr<nb::MemoryPool> small_allocation_pool_;

  DISALLOW_COPY_AND_ASSIGN(MediaBufferAllocator);
};

}  // namespace media
}  // namespace cobalt

#endif  // COBALT_MEDIA_MEDIA_BUFFER_ALLOCATOR_H_