/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * 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 MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_CIRCULAR_POOL_H_
#define MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_CIRCULAR_POOL_H_

#include <atomic>
#include <memory>
#include <vector>
#include "minddata/dataset/util/memory_pool.h"
#include "minddata/dataset/util/arena.h"
#include "minddata/dataset/util/lock.h"

namespace mindspore {
namespace dataset {
using ListOfArenas = std::vector<std::shared_ptr<Arena>>;

// This is a dynamic memory pool built on top of memory
// segment each of which is 4G in size. Initially we start
// with one segment, and gradually add segments (not
// guaranteed contiguous) until we reach 32G in size.  There
// is an assumption about this kind of memory pool.  Allocated
// memory is not held for the whole duration of the pool and
// will be released soon. Based on this assumption, memory is
// obtained from the tail while allocated memory is returned
// to the head of the pool.
class CircularPool : public MemoryPool {
 public:
  class CircularIterator {
    friend class CircularPool;

   public:
    explicit CircularIterator(CircularPool *dp);

    ~CircularIterator() = default;

    bool has_next() const;

    ListOfArenas::iterator Next();

    void Reset();

   private:
    CircularPool *dp_;
    Arena *cur_tail_{};
    uint32_t start_{};
    uint32_t cur_{};
    bool wrap_{};
    bool has_next_{};
  };

  CircularPool(const CircularPool &) = delete;

  CircularPool &operator=(const CircularPool &) = delete;

  ~CircularPool() override;

  Status Allocate(size_t n, void **) override;

  Status Reallocate(void **, size_t old_size, size_t new_size) override;

  void Deallocate(void *) override;

  uint64_t get_max_size() const override;

  int PercentFree() const override;

  friend std::ostream &operator<<(std::ostream &os, const CircularPool &s) {
    int i = 0;
    for (auto it = s.mem_segments_.begin(); it != s.mem_segments_.end(); ++it, ++i) {
      os << "Dumping segment " << i << "\n" << *(it->get());
    }
    return os;
  }

#ifdef ENABLE_GPUQUE
  static Status CreateCircularPool(std::shared_ptr<MemoryPool> *out_pool, int max_size_in_gb = -1,
                                   int arena_size = 4096, bool create_one_arena = false, bool is_cuda_malloc = false);
#else
  static Status CreateCircularPool(std::shared_ptr<MemoryPool> *out_pool, int max_size_in_gb = -1,
                                   int arena_size = 4096, bool create_one_arena = false);
#endif

 private:
  ListOfArenas mem_segments_;
  std::atomic<Arena *> tail_{};
  bool unlimited_;
  int max_size_in_mb_;
  int arena_size_;
  int cur_size_in_mb_;
  RWLock rw_lock_;
#ifdef ENABLE_GPU
  bool is_cuda_malloc_;

  // We can take negative or 0 as input which means unlimited.
  CircularPool(int max_size_in_gb, int arena_size, bool is_cuda_malloc);
#else

  // We can take negative or 0 as input which means unlimited.
  CircularPool(int max_size_in_gb, int arena_size);
#endif

  Status AddOneArena();
};
}  // namespace dataset
}  // namespace mindspore

#endif  // MINDSPORE_CCSRC_MINDDATA_DATASET_UTIL_CIRCULAR_POOL_H_