/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* vim: set ts=8 sts=2 et sw=2 tw=80: */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "mozilla/Assertions.h"
#include "mozilla/Maybe.h"
#include "mozilla/StaticMutex.h"

#include "HashTable.h"
#include "InfallibleVector.h"
#include "ProcessRecordReplay.h"
#include "ProcessRedirect.h"
#include "ValueIndex.h"

#include "PLDHashTable.h"

namespace mozilla {
namespace recordreplay {

// Hash tables frequently incorporate pointer values into the hash numbers they
// compute, which are not guaranteed to be the same between recording and
// replaying and consequently lead to inconsistent hash numbers and iteration
// order between recording and replaying, which can in turn affect the order in
// which recorded events occur. HashTable stabilization is designed to deal
// with this problem, for specific kinds of hashtables (PLD and PL tables)
// which are based on callbacks.
//
// When the table is constructed, if we are recording/replaying then the
// callbacks are replaced with an alternate set that produces consistent hash
// numbers between recording and replay. If during replay the additions and
// removals to the tables occur in the same order that they did during
// recording, then the structure of the tables and the order in which elements
// are visited during iteration will be the same.
//
// Ensuring that hash numbers are consistent is done as follows: for each
// table, we keep track of the keys that are in the table. When computing the
// hash of an arbitrary key, we look for a matching key in the table, using
// that key's hash if found. Otherwise, a new hash is generated from an
// incrementing counter.

typedef uint32_t HashNumber;

class StableHashTableInfo
{
  // Magic number for attempting to determine whether we are dealing with an
  // actual StableHashTableInfo. Despite our best efforts some hashtables do
  // not go through stabilization (e.g. they have static constructors that run
  // before record/replay state is initialized).
  size_t mMagic;

  static const size_t MagicNumber = 0xDEADBEEFDEADBEEF;

  // Information about a key in the table: the key pointer, along with the new
  // hash number we have generated for the key.
  struct KeyInfo {
    const void* mKey;
    HashNumber mNewHash;
  };

  // Table mapping original hash numbers (produced by the table's hash
  // function) to a vector with all keys sharing that original hash number.
  struct HashInfo {
    InfallibleVector<KeyInfo> mKeys;
  };
  typedef std::unordered_map<HashNumber, UniquePtr<HashInfo>> HashToKeyMap;
  HashToKeyMap mHashToKey;

  // Table mapping key pointers to their original hash number.
  typedef std::unordered_map<const void*, HashNumber> KeyToHashMap;
  KeyToHashMap mKeyToHash;

  // The last key which the hash function was called on, and the new hash
  // number which we generated for that key.
  const void* mLastKey;
  HashNumber mLastNewHash;

  // Counter for generating new hash numbers for entries added to the table.
  // This increases monotonically, though it is fine if it overflows.
  uint32_t mHashGenerator;

  // Buffer with executable memory for use in binding functions.
  uint8_t* mCallbackStorage;
  static const size_t CallbackStorageCapacity = 4096;

  // Get an existing key in the table.
  KeyInfo* FindKeyInfo(HashNumber aOriginalHash, const void* aKey, HashInfo** aHashInfo = nullptr) {
    HashToKeyMap::iterator iter = mHashToKey.find(aOriginalHash);
    MOZ_ASSERT(iter != mHashToKey.end());

    HashInfo* hashInfo = iter->second.get();
    for (KeyInfo& keyInfo : hashInfo->mKeys) {
      if (keyInfo.mKey == aKey) {
        if (aHashInfo) {
          *aHashInfo = hashInfo;
        }
        return &keyInfo;
      }
    }
    MOZ_CRASH();
  }

public:
  StableHashTableInfo()
    : mMagic(MagicNumber)
    , mLastKey(nullptr)
    , mLastNewHash(0)
    , mHashGenerator(0)
    , mCallbackStorage(nullptr)
  {
    // Use AllocateMemory, as the result will have RWX permissions.
    mCallbackStorage = (uint8_t*) AllocateMemory(CallbackStorageCapacity, MemoryKind::Tracked);
  }

  ~StableHashTableInfo() {
    MOZ_ASSERT(mHashToKey.empty());
    DeallocateMemory(mCallbackStorage, CallbackStorageCapacity, MemoryKind::Tracked);
  }

  bool AppearsValid() {
    return mMagic == MagicNumber;
  }

  void AddKey(HashNumber aOriginalHash, const void* aKey, HashNumber aNewHash) {
    HashToKeyMap::iterator iter = mHashToKey.find(aOriginalHash);
    if (iter == mHashToKey.end()) {
      iter = mHashToKey.insert(HashToKeyMap::value_type(aOriginalHash, MakeUnique<HashInfo>())).first;
    }
    HashInfo* hashInfo = iter->second.get();

    KeyInfo key;
    key.mKey = aKey;
    key.mNewHash = aNewHash;
    hashInfo->mKeys.append(key);

    mKeyToHash.insert(KeyToHashMap::value_type(aKey, aOriginalHash));
  }

  void RemoveKey(HashNumber aOriginalHash, const void* aKey) {
    HashInfo* hashInfo;
    KeyInfo* keyInfo = FindKeyInfo(aOriginalHash, aKey, &hashInfo);
    hashInfo->mKeys.erase(keyInfo);

    if (hashInfo->mKeys.length() == 0) {
      mHashToKey.erase(aOriginalHash);
    }

    mKeyToHash.erase(aKey);
  }

  HashNumber FindKeyHash(HashNumber aOriginalHash, const void* aKey) {
    KeyInfo* info = FindKeyInfo(aOriginalHash, aKey);
    return info->mNewHash;
  }

  // Look for a key in the table with a matching original hash and for which
  // aMatch() is true for, returning its new hash number if found.
  bool HasMatchingKey(HashNumber aOriginalHash,
                      const std::function<bool(const void*)>& aMatch,
                      HashNumber* aNewHash)
  {
    HashToKeyMap::const_iterator iter = mHashToKey.find(aOriginalHash);
    if (iter != mHashToKey.end()) {
      HashInfo* hashInfo = iter->second.get();
      for (const KeyInfo& keyInfo : hashInfo->mKeys) {
        if (aMatch(keyInfo.mKey)) {
          *aNewHash = keyInfo.mNewHash;
          return true;
        }
      }
    }
    return false;
  }

  HashNumber GetOriginalHashNumber(const void* aKey) {
    KeyToHashMap::iterator iter = mKeyToHash.find(aKey);
    MOZ_ASSERT(iter != mKeyToHash.end());
    return iter->second;
  }

  class Assembler : public recordreplay::Assembler {
  public:
    explicit Assembler(StableHashTableInfo& aInfo)
      : recordreplay::Assembler(aInfo.mCallbackStorage, CallbackStorageCapacity)
    {}
  };

  // Use the callback storage buffer to create a new function T which has one
  // fewer argument than S and calls S with aArgument bound to the last
  // argument position. See BindFunctionArgument in ProcessRedirect.h
  template <typename S, typename T>
  void NewBoundFunction(Assembler& aAssembler, S aFunction, void* aArgument,
                        size_t aArgumentPosition, T* aTarget) {
    void* nfn = BindFunctionArgument(BitwiseCast<void*>(aFunction), aArgument, aArgumentPosition,
                                     aAssembler);
    BitwiseCast(nfn, aTarget);
  }

  // Set the last queried key for this table, and generate a new hash number
  // for it.
  HashNumber SetLastKey(const void* aKey) {
    // Remember the last key queried, so that if it is then added to the table
    // we know what hash number to use.
    mLastKey = aKey;
    mLastNewHash = mHashGenerator++;
    return mLastNewHash;
  }

  bool HasLastKey() {
    return !!mLastKey;
  }

  HashNumber GetLastNewHash(const void* aKey) {
    MOZ_ASSERT(aKey == mLastKey);
    return mLastNewHash;
  }

  bool IsEmpty() { return mHashToKey.empty(); }

  // Move aOther's contents into this one and clear aOther out. Callbacks for
  // the tables are left alone.
  void MoveContentsFrom(StableHashTableInfo& aOther) {
    mHashToKey = std::move(aOther.mHashToKey);
    mKeyToHash = std::move(aOther.mKeyToHash);
    mHashGenerator = aOther.mHashGenerator;

    aOther.mHashToKey.clear();
    aOther.mKeyToHash.clear();
    aOther.mHashGenerator = 0;

    mLastKey = aOther.mLastKey = nullptr;
    mLastNewHash = aOther.mLastNewHash = 0;
  }
};

///////////////////////////////////////////////////////////////////////////////
// PLHashTable Stabilization
///////////////////////////////////////////////////////////////////////////////

// For each PLHashTable in the process, a PLHashTableInfo is generated. This
// structure becomes the |allocPriv| for the table, handled by the new
// callbacks given to it.
struct PLHashTableInfo : public StableHashTableInfo
{
  // Original callbacks for the table.
  PLHashFunction mKeyHash;
  PLHashComparator mKeyCompare;
  PLHashComparator mValueCompare;
  const PLHashAllocOps* mAllocOps;

  // Original private value for the table.
  void* mAllocPrivate;

  PLHashTableInfo(PLHashFunction aKeyHash,
                  PLHashComparator aKeyCompare, PLHashComparator aValueCompare,
                  const PLHashAllocOps* aAllocOps, void* aAllocPrivate)
    : mKeyHash(aKeyHash),
      mKeyCompare(aKeyCompare),
      mValueCompare(aValueCompare),
      mAllocOps(aAllocOps),
      mAllocPrivate(aAllocPrivate)
  {}

  static PLHashTableInfo* FromPrivate(void* aAllocPrivate) {
    PLHashTableInfo* info = reinterpret_cast<PLHashTableInfo*>(aAllocPrivate);
    MOZ_RELEASE_ASSERT(info->AppearsValid());
    return info;
  }
};

static void*
WrapPLHashAllocTable(void* aAllocPrivate, PRSize aSize)
{
  PLHashTableInfo* info = PLHashTableInfo::FromPrivate(aAllocPrivate);
  return info->mAllocOps
         ? info->mAllocOps->allocTable(info->mAllocPrivate, aSize)
         : malloc(aSize);
}

static void
WrapPLHashFreeTable(void* aAllocPrivate, void* aItem)
{
  PLHashTableInfo* info = PLHashTableInfo::FromPrivate(aAllocPrivate);
  if (info->mAllocOps) {
    info->mAllocOps->freeTable(info->mAllocPrivate, aItem);
  } else {
    free(aItem);
  }
}

static PLHashEntry*
WrapPLHashAllocEntry(void* aAllocPrivate, const void* aKey)
{
  PLHashTableInfo* info = PLHashTableInfo::FromPrivate(aAllocPrivate);

  if (info->HasLastKey()) {
    uint32_t originalHash = info->mKeyHash(aKey);
    info->AddKey(originalHash, aKey, info->GetLastNewHash(aKey));
  } else {
    // A few PLHashTables are manipulated directly by Gecko code, in which case
    // the hashes are supplied directly to the table and we don't have a chance
    // to modify them. Fortunately, none of these tables are iterated in a way
    // that can cause the replay to diverge, so just punt in these cases.
    MOZ_ASSERT(info->IsEmpty());
  }

  return info->mAllocOps
         ? info->mAllocOps->allocEntry(info->mAllocPrivate, aKey)
         : (PLHashEntry*) malloc(sizeof(PLHashEntry));
}

static void
WrapPLHashFreeEntry(void *aAllocPrivate, PLHashEntry *he, PRUintn flag)
{
  PLHashTableInfo* info = PLHashTableInfo::FromPrivate(aAllocPrivate);

  // Ignore empty tables, due to the raw table manipulation described above.
  if (flag == HT_FREE_ENTRY && !info->IsEmpty()) {
    uint32_t originalHash = info->GetOriginalHashNumber(he->key);
    info->RemoveKey(originalHash, he->key);
  }

  if (info->mAllocOps) {
    info->mAllocOps->freeEntry(info->mAllocPrivate, he, flag);
  } else if (flag == HT_FREE_ENTRY) {
    free(he);
  }
}

static PLHashAllocOps gWrapPLHashAllocOps = {
  WrapPLHashAllocTable, WrapPLHashFreeTable,
  WrapPLHashAllocEntry, WrapPLHashFreeEntry
};

static uint32_t
PLHashComputeHash(void* aKey, PLHashTableInfo* aInfo)
{
  uint32_t originalHash = aInfo->mKeyHash(aKey);
  HashNumber newHash;
  if (aInfo->HasMatchingKey(originalHash,
                            [=](const void* aExistingKey) {
                              return aInfo->mKeyCompare(aKey, aExistingKey);
                            }, &newHash)) {
    return newHash;
  }
  return aInfo->SetLastKey(aKey);
}

void
GeneratePLHashTableCallbacks(PLHashFunction* aKeyHash,
                             PLHashComparator* aKeyCompare,
                             PLHashComparator* aValueCompare,
                             const PLHashAllocOps** aAllocOps,
                             void** aAllocPrivate)
{
  PLHashTableInfo* info = new PLHashTableInfo(*aKeyHash, *aKeyCompare, *aValueCompare,
                                              *aAllocOps, *aAllocPrivate);
  PLHashTableInfo::Assembler assembler(*info);
  info->NewBoundFunction(assembler, PLHashComputeHash, info, 1, aKeyHash);
  *aAllocOps = &gWrapPLHashAllocOps;
  *aAllocPrivate = info;
}

void
DestroyPLHashTableCallbacks(void* aAllocPrivate)
{
  PLHashTableInfo* info = PLHashTableInfo::FromPrivate(aAllocPrivate);
  delete info;
}

///////////////////////////////////////////////////////////////////////////////
// PLDHashTable Stabilization
///////////////////////////////////////////////////////////////////////////////

// For each PLDHashTable in the process, a PLDHashTableInfo is generated. This
// structure is supplied to its callbacks using bound functions.
struct PLDHashTableInfo : public StableHashTableInfo
{
  // Original callbacks for the table.
  const PLDHashTableOps* mOps;

  // Wrapper callbacks for the table.
  PLDHashTableOps mNewOps;

  explicit PLDHashTableInfo(const PLDHashTableOps* aOps)
    : mOps(aOps)
  {
    PodZero(&mNewOps);
  }

  static PLDHashTableInfo* MaybeFromOps(const PLDHashTableOps* aOps) {
    PLDHashTableInfo* res = reinterpret_cast<PLDHashTableInfo*>
      ((uint8_t*)aOps - offsetof(PLDHashTableInfo, mNewOps));
    return res->AppearsValid() ? res : nullptr;
  }

  static PLDHashTableInfo* FromOps(const PLDHashTableOps* aOps) {
    PLDHashTableInfo* res = MaybeFromOps(aOps);
    MOZ_RELEASE_ASSERT(res);
    return res;
  }
};

static PLDHashNumber
PLDHashTableComputeHash(const void* aKey, PLDHashTableInfo* aInfo)
{
  uint32_t originalHash = aInfo->mOps->hashKey(aKey);
  HashNumber newHash;
  if (aInfo->HasMatchingKey(originalHash,
                            [=](const void* aExistingKey) {
                              return aInfo->mOps->matchEntry((PLDHashEntryHdr*) aExistingKey, aKey);
                            }, &newHash)) {
    return newHash;
  }
  return aInfo->SetLastKey(aKey);
}

static void
PLDHashTableMoveEntry(PLDHashTable* aTable, const PLDHashEntryHdr* aFrom, PLDHashEntryHdr* aTo,
                      PLDHashTableInfo* aInfo)
{
  aInfo->mOps->moveEntry(aTable, aFrom, aTo);

  uint32_t originalHash = aInfo->GetOriginalHashNumber(aFrom);
  uint32_t newHash = aInfo->FindKeyHash(originalHash, aFrom);

  aInfo->RemoveKey(originalHash, aFrom);
  aInfo->AddKey(originalHash, aTo, newHash);
}

static void
PLDHashTableClearEntry(PLDHashTable* aTable, PLDHashEntryHdr* aEntry, PLDHashTableInfo* aInfo)
{
  aInfo->mOps->clearEntry(aTable, aEntry);

  uint32_t originalHash = aInfo->GetOriginalHashNumber(aEntry);
  aInfo->RemoveKey(originalHash, aEntry);
}

static void
PLDHashTableInitEntry(PLDHashEntryHdr* aEntry, const void* aKey, PLDHashTableInfo* aInfo)
{
  if (aInfo->mOps->initEntry) {
    aInfo->mOps->initEntry(aEntry, aKey);
  }

  uint32_t originalHash = aInfo->mOps->hashKey(aKey);
  aInfo->AddKey(originalHash, aEntry, aInfo->GetLastNewHash(aKey));
}

extern "C" {

MOZ_EXPORT const PLDHashTableOps*
RecordReplayInterface_InternalGeneratePLDHashTableCallbacks(const PLDHashTableOps* aOps)
{
  PLDHashTableInfo* info = new PLDHashTableInfo(aOps);
  PLDHashTableInfo::Assembler assembler(*info);
  info->NewBoundFunction(assembler, PLDHashTableComputeHash, info, 1, &info->mNewOps.hashKey);
  info->mNewOps.matchEntry = aOps->matchEntry;
  info->NewBoundFunction(assembler, PLDHashTableMoveEntry, info, 3, &info->mNewOps.moveEntry);
  info->NewBoundFunction(assembler, PLDHashTableClearEntry, info, 2, &info->mNewOps.clearEntry);
  info->NewBoundFunction(assembler, PLDHashTableInitEntry, info, 2, &info->mNewOps.initEntry);
  return &info->mNewOps;
}

MOZ_EXPORT const PLDHashTableOps*
RecordReplayInterface_InternalUnwrapPLDHashTableCallbacks(const PLDHashTableOps* aOps)
{
  PLDHashTableInfo* info = PLDHashTableInfo::FromOps(aOps);
  return info->mOps;
}

MOZ_EXPORT void
RecordReplayInterface_InternalDestroyPLDHashTableCallbacks(const PLDHashTableOps* aOps)
{
  // Primordial PLDHashTables used in the copy constructor might not have any ops.
  if (!aOps) {
    return;
  }

  // Note: PLDHashTables with static ctors might have been constructed before
  // record/replay state was initialized and have their normal ops. Check the
  // magic number via MaybeFromOps before destroying the info.
  PLDHashTableInfo* info = PLDHashTableInfo::MaybeFromOps(aOps);
  delete info;
}

MOZ_EXPORT void
RecordReplayInterface_InternalMovePLDHashTableContents(const PLDHashTableOps* aFirstOps,
                                                       const PLDHashTableOps* aSecondOps)
{
  PLDHashTableInfo* firstInfo = PLDHashTableInfo::FromOps(aFirstOps);
  PLDHashTableInfo* secondInfo = PLDHashTableInfo::FromOps(aSecondOps);

  secondInfo->MoveContentsFrom(*firstInfo);
}

} // extern "C"

} // namespace recordreplay
} // namespace mozilla