/*
 * 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 NB_MEMORY_TRACKER_HELPERS_H_
#define NB_MEMORY_TRACKER_HELPERS_H_

#include <map>
#include <vector>

#include "nb/analytics/memory_tracker.h"
#include "nb/atomic.h"
#include "nb/simple_thread.h"
#include "nb/thread_local_boolean.h"
#include "nb/thread_local_pointer.h"
#include "starboard/mutex.h"
#include "starboard/thread.h"
#include "starboard/types.h"
#include "starboard/log.h"

namespace nb {
namespace analytics {

class AllocationGroup;
class AllocationRecord;

// An AllocationGroup is a collection of allocations that are logically lumped
// together, such as "Javascript" or "Graphics".
class AllocationGroup {
 public:
  explicit AllocationGroup(const std::string& name);
  ~AllocationGroup();
  const std::string& name() const { return name_; }

  void AddAllocation(int64_t num_bytes);
  void GetAggregateStats(int32_t* num_allocs, int64_t* allocation_bytes) const;

  int64_t allocation_bytes() const;
  int32_t num_allocations() const;

 private:
  const std::string name_;
  nb::atomic_int64_t allocation_bytes_;
  nb::atomic_int32_t num_allocations_;

  SB_DISALLOW_COPY_AND_ASSIGN(AllocationGroup);
};

// A self locking data structure that maps strings -> AllocationGroups. This is
// used to resolve MemoryGroup names (e.g. "Javascript") to an AllocationGroup
// which can be used to group allocations together.
class AtomicStringAllocationGroupMap {
 public:
  AtomicStringAllocationGroupMap();
  ~AtomicStringAllocationGroupMap();

  AllocationGroup* Ensure(const std::string& name);
  AllocationGroup* GetDefaultUnaccounted();

  bool Erase(const std::string& key);
  void GetAll(std::vector<const AllocationGroup*>* output) const;

 private:
  typedef std::map<std::string, AllocationGroup*> Map;
  Map group_map_;
  AllocationGroup* unaccounted_group_;
  mutable starboard::Mutex mutex_;

  SB_DISALLOW_COPY_AND_ASSIGN(AtomicStringAllocationGroupMap);
};

class AllocationGroupStack {
 public:
  AllocationGroupStack() { Push_DebugBreak(NULL); }
  ~AllocationGroupStack() {}

  void Push(AllocationGroup* group);
  void Pop();
  AllocationGroup* Peek();
  const AllocationGroup* Peek() const;

  void Push_DebugBreak(AllocationGroup* ag) { debug_stack_.push_back(ag); }
  void Pop_DebugBreak() { debug_stack_.pop_back(); }
  AllocationGroup* Peek_DebugBreak() {
    if (debug_stack_.empty()) {
      return NULL;
    }
    return debug_stack_.back();
  }
  const AllocationGroupPtrVec& data() const { return alloc_group_stack_; }

 private:
  SB_DISALLOW_COPY_AND_ASSIGN(AllocationGroupStack);
  AllocationGroupPtrVec alloc_group_stack_, debug_stack_;
};

// A per-pointer map of allocations to AllocRecords. This map is thread safe.
class AtomicAllocationMap {
 public:
  AtomicAllocationMap();
  ~AtomicAllocationMap();

  // Returns true if Added. Otherwise false means that the pointer
  // already existed.
  bool Add(const void* memory, const AllocationRecord& alloc_record);

  // Returns true if the memory exists in this set.
  bool Get(const void* memory, AllocationRecord* alloc_record) const;

  // Return true if the memory existed in this set. If true
  // then output alloc_record is written with record that was found.
  // otherwise the record is written as 0 bytes and null key.
  bool Remove(const void* memory, AllocationRecord* alloc_record);

  bool Accept(AllocationVisitor* visitor) const;

  size_t Size() const;
  bool Empty() const;
  void Clear();

 private:
  SB_DISALLOW_COPY_AND_ASSIGN(AtomicAllocationMap);
  typedef std::map<const void*, AllocationRecord> PointerMap;

  PointerMap pointer_map_;
  mutable starboard::Mutex mutex_;
};

// A per-pointer map of allocations to AllocRecords. This is a hybrid data
// structure consisting of a hashtable of maps. Each pointer that is
// stored or retrieved is hashed to a random bucket. Each bucket has it's own
// lock. This distributed pattern increases performance significantly by
// reducing contention. The top-level hashtable is of constant size and does
// not resize. Each bucket is implemented as it's own map of elements.
class ConcurrentAllocationMap {
 public:
  static const int kNumElements = 511;
  ConcurrentAllocationMap();
  ~ConcurrentAllocationMap();

  // Returns true if Added. Otherwise false means that the pointer
  // already existed.
  bool Add(const void* memory, const AllocationRecord& alloc_record);
  // Returns true if the memory exists in this set.
  bool Get(const void* memory, AllocationRecord* alloc_record) const;
  // Return true if the memory existed in this set. If true
  // then output alloc_record is written with record that was found.
  // otherwise the record is written as 0 bytes and null key.
  bool Remove(const void* memory, AllocationRecord* alloc_record);
  size_t Size() const;
  bool Empty() const;
  void Clear();

  // Provides access to all the allocations within in a thread safe manner.
  bool Accept(AllocationVisitor* visitor) const;

  AtomicAllocationMap& GetMapForPointer(const void* ptr);
  const AtomicAllocationMap& GetMapForPointer(const void* ptr) const;

 private:
  SB_DISALLOW_COPY_AND_ASSIGN(ConcurrentAllocationMap);
  // Takes a pointer and generates a hash.
  static size_t hash_ptr(const void* ptr);

  int ToIndex(const void* ptr) const;
  AtomicAllocationMap pointer_map_array_[kNumElements];
};

}  // namespace analytics
}  // namespace nb

#endif  // NB_MEMORY_TRACKER_HELPERS_H_