/*
 * Copyright 2016 Facebook, Inc.
 *
 * 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.
 */
// Copyright 2013-present Facebook. All Rights Reserved.

#include <folly/experimental/StringKeyedMap.h>
#include <folly/experimental/StringKeyedSet.h>
#include <folly/experimental/StringKeyedUnorderedMap.h>
#include <folly/experimental/StringKeyedUnorderedSet.h>

#include <list>
#include <string>

#include <glog/logging.h>
#include <gtest/gtest.h>

#include <folly/Range.h>
#include <folly/portability/GFlags.h>

using folly::StringKeyedMap;
using folly::StringKeyedSetBase;
using folly::StringKeyedUnorderedMap;
using folly::BasicStringKeyedUnorderedSet;
using folly::StringPiece;
using std::string;

static unsigned long long allocated = 0;
static unsigned long long freed = 0;

template <typename Alloc>
struct MemoryLeakCheckerAllocator {
  typedef typename Alloc::value_type value_type;
  typedef value_type *pointer;
  typedef value_type const *const_pointer;
  typedef value_type &reference;
  typedef value_type const *const_reference;

  typedef std::ptrdiff_t difference_type;
  typedef std::size_t size_type;

  explicit MemoryLeakCheckerAllocator() {
  }

  explicit MemoryLeakCheckerAllocator(Alloc alloc)
      : alloc_(alloc) {
  }

  template <class UAlloc>
  MemoryLeakCheckerAllocator(const MemoryLeakCheckerAllocator<UAlloc>& other)
      : alloc_(other.allocator()) {
  }

  value_type* allocate(size_t n, const void* hint = nullptr) {
    auto p = alloc_.allocate(n, hint);
    allocated += n * sizeof(value_type);
    return p;
  }

  void deallocate(value_type* p, size_t n) {
    alloc_.deallocate(p, n);
    freed += n * sizeof(value_type);
  }

  size_t max_size() const {
    return alloc_.max_size();
  }

  template <class... Args>
  void construct(value_type* p, Args&&... args) {
    alloc_.construct(p, std::forward<Args>(args)...);
  }

  void destroy(value_type* p) {
    alloc_.destroy(p);
  }

  template <class U>
  struct rebind {
    typedef MemoryLeakCheckerAllocator<
      typename std::allocator_traits<Alloc>::template rebind_alloc<U>
    > other;
  };

  const Alloc& allocator() const {
    return alloc_;
  }

  bool operator!=(const MemoryLeakCheckerAllocator& other) const {
    return alloc_ != other.alloc_;
  }

  bool operator==(const MemoryLeakCheckerAllocator& other) const {
    return alloc_ == other.alloc_;
  }

private:
  Alloc alloc_;
};

typedef MemoryLeakCheckerAllocator<std::allocator<char>> KeyLeakChecker;
typedef MemoryLeakCheckerAllocator<
  std::allocator<std::pair<const StringPiece, int>>> ValueLeakChecker;

typedef StringKeyedUnorderedMap<int, folly::StringPieceHash,
                                std::equal_to<StringPiece>, ValueLeakChecker>
  LeakCheckedUnorderedMap;

typedef StringKeyedSetBase<std::less<StringPiece>, ValueLeakChecker> LeakCheckedSet;

typedef StringKeyedMap<int, std::less<StringPiece>,
                       ValueLeakChecker> LeakCheckedMap;

using LeakCheckedUnorderedSet = BasicStringKeyedUnorderedSet<
  folly::StringPieceHash,
  std::equal_to<folly::StringPiece>,
  ValueLeakChecker>;

TEST(StringKeyedUnorderedMapTest, sanity) {
  LeakCheckedUnorderedMap map;
  EXPECT_TRUE(map.empty());
  EXPECT_EQ(map.size(), 0);

  {
    string s("hello");
    StringPiece piece(s, 3);
    map.insert({s, 1});
    EXPECT_FALSE(map.emplace(s, 2).second);
    EXPECT_TRUE(map.emplace(piece, 3).second);
  }

  EXPECT_EQ(map.size(), 2);

  map = map;

  EXPECT_EQ(map.find("hello")->second, 1);
  EXPECT_EQ(map.find("lo")->second, 3);

  map.erase(map.find("hello"));

  EXPECT_EQ(map.size(), 1);

  for (auto& it : map) {
    EXPECT_EQ(it.first, "lo");
  }
}

TEST(StringKeyedUnorderedMapTest, constructors) {
  LeakCheckedUnorderedMap map {
    {"hello", 1},
    {"lo", 3}
  };

  LeakCheckedUnorderedMap map2(map);
  EXPECT_EQ(map2.size(), 2);
  EXPECT_TRUE(map2 == map);

  map2.erase("lo");
  for (auto& it : map2) {
    EXPECT_EQ(it.first, "hello");
  }

  map2.clear();

  EXPECT_TRUE(map2.empty());

  map2.emplace("key1", 1);

  LeakCheckedUnorderedMap map3(std::move(map2));

  EXPECT_EQ(map3.size(), 1);
  EXPECT_EQ(map3["key1"], 1);

  EXPECT_EQ(map3["key0"], 0);
  EXPECT_EQ(map3.size(), 2);

  map3.reserve(1000);

  EXPECT_EQ(map3.size(), 2);

  LeakCheckedUnorderedMap map4 {
    {"key0", 0},
    {"key1", 1}
  };

  EXPECT_EQ(map4.erase("key0"), 1);
  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.find("key0"), map4.end());

  map3 = map4;

  EXPECT_EQ(map3.size(), 1);
  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.max_size(), map3.max_size());

  map4 = std::move(map3);

  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.at("key1"), 1);
}

TEST(StringKeyedSetTest, sanity) {
  LeakCheckedSet set;
  EXPECT_TRUE(set.empty());
  EXPECT_EQ(set.size(), 0);

  {
    string s("hello");
    StringPiece piece(s, 3);
    set.insert(s);
    EXPECT_FALSE(set.emplace(s).second);
    EXPECT_TRUE(set.emplace(piece).second);
  }

  EXPECT_EQ(set.size(), 2);

  set = set;

  EXPECT_NE(set.find(StringPiece("hello")), set.end());
  EXPECT_NE(set.find("lo"), set.end());

  auto it = set.begin();
  EXPECT_EQ(*it, "hello");
  EXPECT_EQ(*(++it), "lo");
  EXPECT_EQ(++it, set.end());

  set.erase(set.find("hello"));

  EXPECT_EQ(set.size(), 1);

  for (auto it : set) {
    EXPECT_EQ(it, "lo");
  }
}

TEST(StringKeyedSetTest, constructors) {
  LeakCheckedSet set {
    "hello",
    "lo"
  };
  LeakCheckedSet set2(set);

  EXPECT_EQ(set2.size(), 2);

  set2.erase("lo");
  for (auto it : set2) {
    EXPECT_EQ(it, "hello");
  }

  set2.clear();

  EXPECT_TRUE(set2.empty());

  set2.emplace("key1");

  LeakCheckedSet set3(std::move(set2));

  EXPECT_EQ(set3.size(), 1);
  EXPECT_EQ(set3.insert("key1").second, false);

  EXPECT_EQ(set3.emplace("key0").second, true);
  EXPECT_EQ(set3.size(), 2);

  EXPECT_EQ(set3.size(), 2);

  LeakCheckedSet set4 {
    "key0",
    "key1"
  };

  EXPECT_EQ(set4.erase("key0"), 1);
  EXPECT_EQ(set4.size(), 1);
  EXPECT_EQ(set4.find("key0"), set4.end());

  set3 = set4;

  EXPECT_EQ(set3.size(), 1);
  EXPECT_EQ(set4.size(), 1);
  EXPECT_EQ(set4.max_size(), set3.max_size());

  set4 = std::move(set3);

  EXPECT_EQ(set4.size(), 1);
  EXPECT_NE(set4.find("key1"), set4.end());
}

TEST(StringKeyedUnorderedSetTest, sanity) {
  LeakCheckedUnorderedSet set;
  EXPECT_TRUE(set.empty());
  EXPECT_EQ(set.size(), 0);

  {
    string s("hello");
    StringPiece piece(s, 3);
    set.insert(s);
    EXPECT_FALSE(set.emplace(s).second);
    EXPECT_TRUE(set.emplace(piece).second);
  }

  EXPECT_EQ(set.size(), 2);

  set = set;

  EXPECT_NE(set.find("hello"), set.end());
  EXPECT_NE(set.find("lo"), set.end());

  set.erase(set.find("hello"));

  EXPECT_EQ(set.size(), 1);

  for (auto it : set) {
    EXPECT_EQ(it, "lo");
  }
}

TEST(StringKeyedUnorderedSetTest, constructors) {
  LeakCheckedUnorderedSet s1;
  EXPECT_TRUE(s1.empty());

  LeakCheckedUnorderedSet s2(10);
  EXPECT_TRUE(s2.empty());
  EXPECT_GE(s2.bucket_count(), 10);

  std::list<StringPiece> lst { "abc", "def" };
  LeakCheckedUnorderedSet s3(lst.begin(), lst.end());
  EXPECT_EQ(s3.size(), 2);
  EXPECT_NE(s3.find("abc"), s3.end());
  EXPECT_NE(s3.find("def"), s3.end());
  EXPECT_TRUE(s3 == (LeakCheckedUnorderedSet{"abc", "def"}));

  LeakCheckedUnorderedSet s4(const_cast<LeakCheckedUnorderedSet&>(s3));
  EXPECT_TRUE(s4 == s3);

  LeakCheckedUnorderedSet s5(const_cast<LeakCheckedUnorderedSet&>(s3),
                             ValueLeakChecker());
  EXPECT_TRUE(s5 == s3);

  LeakCheckedUnorderedSet s6(std::move(s3));
  EXPECT_TRUE(s3.empty());
  EXPECT_TRUE(s6 == s5);

  LeakCheckedUnorderedSet s7(std::move(s6), s6.get_allocator());
  EXPECT_TRUE(s6.empty());
  EXPECT_TRUE(s7 == s5);

  LeakCheckedUnorderedSet s8 {
    "hello",
    "lo"
  };
  EXPECT_EQ(s8.size(), 2);
  EXPECT_NE(s8.find("hello"), s8.end());
  EXPECT_NE(s8.find("lo"), s8.end());

  LeakCheckedUnorderedSet s9({
    "hello",
    "lo"
      }, 10);
  EXPECT_EQ(s9.size(), 2);
  EXPECT_NE(s9.find("hello"), s9.end());
  EXPECT_NE(s9.find("lo"), s9.end());

  LeakCheckedUnorderedSet set2(s8);
  EXPECT_EQ(set2.size(), 2);

  set2.erase("lo");
  for (auto it : set2) {
    EXPECT_EQ(it, "hello");
  }

  set2.clear();

  EXPECT_TRUE(set2.empty());

  set2.emplace("key1");

  LeakCheckedUnorderedSet set3(std::move(set2));

  EXPECT_EQ(set3.size(), 1);
  EXPECT_EQ(set3.insert("key1").second, false);

  EXPECT_EQ(set3.emplace("key0").second, true);
  EXPECT_EQ(set3.size(), 2);

  set3.reserve(1000);

  EXPECT_EQ(set3.size(), 2);

  LeakCheckedUnorderedSet set4 {
    "key0",
    "key1"
  };

  EXPECT_EQ(set4.erase("key0"), 1);
  EXPECT_EQ(set4.size(), 1);
  EXPECT_EQ(set4.find("key0"), set4.end());

  set3 = set4;

  EXPECT_EQ(set3.size(), 1);
  EXPECT_EQ(set4.size(), 1);
  EXPECT_EQ(set4.max_size(), set3.max_size());

  set4 = std::move(set3);

  EXPECT_EQ(set4.size(), 1);
  EXPECT_NE(set4.find("key1"), set4.end());
}

TEST(StringKeyedMapTest, sanity) {
  LeakCheckedMap map;
  EXPECT_TRUE(map.empty());
  EXPECT_EQ(map.size(), 0);

  {
    string s("hello");
    StringPiece piece(s, 3);
    map.insert({s, 1});
    EXPECT_FALSE(map.emplace(s, 2).second);
    EXPECT_TRUE(map.emplace(piece, 3).second);
  }

  EXPECT_EQ(map.size(), 2);

  map = map;

  EXPECT_EQ(map.find("hello")->second, 1);
  EXPECT_EQ(map.find("lo")->second, 3);

  auto it = map.begin();
  EXPECT_EQ(it->first, "hello");
  EXPECT_EQ((++it)->first, "lo");
  EXPECT_EQ(++it, map.end());

  map.erase(map.find("hello"));

  EXPECT_EQ(map.size(), 1);

  for (auto& it : map) {
    EXPECT_EQ(it.first, "lo");
  }
}

TEST(StringKeyedMapTest, constructors) {
  LeakCheckedMap map {
    {"hello", 1},
    {"lo", 3}
  };
  LeakCheckedMap map2(map);

  EXPECT_EQ(map2.size(), 2);

  map2.erase("lo");
  for (auto& it : map2) {
    EXPECT_EQ(it.first, "hello");
  }

  map2.clear();

  EXPECT_TRUE(map2.empty());

  map2.emplace("key1", 1);

  LeakCheckedMap map3(std::move(map2));

  EXPECT_EQ(map3.size(), 1);
  EXPECT_EQ(map3["key1"], 1);

  EXPECT_EQ(map3["key0"], 0);
  EXPECT_EQ(map3.size(), 2);

  LeakCheckedMap map4 {
    {"key0", 0},
    {"key1", 1}
  };

  EXPECT_EQ(map4.erase("key0"), 1);
  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.find("key0"), map4.end());

  map3 = map4;

  EXPECT_EQ(map3.size(), 1);
  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.max_size(), map3.max_size());

  map4 = std::move(map3);

  EXPECT_EQ(map4.size(), 1);
  EXPECT_EQ(map4.at("key1"), 1);
}

int main(int argc, char **argv) {
  FLAGS_logtostderr = true;
  google::InitGoogleLogging(argv[0]);
  testing::InitGoogleTest(&argc, argv);
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  return RUN_ALL_TESTS();
}

// This MUST be the LAST test.
TEST(StringKeyed, memory_balance) {
  auto balance = allocated < freed
    ? freed - allocated
    : allocated - freed;

  LOG(INFO) << "allocated: " << allocated
    << " freed: " << freed
    << " balance: " << balance
    << (
      allocated < freed
        ? " negative (huh?)"
        : freed < allocated
          ? " positive (leak)" : ""
    );

  EXPECT_EQ(allocated, freed);
}
