/* -*- 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 "MemorySnapshot.h"

#include "ipc/ChildInternal.h"
#include "mozilla/Maybe.h"
#include "DirtyMemoryHandler.h"
#include "InfallibleVector.h"
#include "ProcessRecordReplay.h"
#include "ProcessRewind.h"
#include "SpinLock.h"
#include "SplayTree.h"
#include "Thread.h"

#include <algorithm>
#include <mach/mach.h>
#include <mach/mach_vm.h>

// Define to enable the countdown debugging thread. See StartCountdown().
//#define WANT_COUNTDOWN_THREAD 1

namespace mozilla {
namespace recordreplay {

///////////////////////////////////////////////////////////////////////////////
// Memory Snapshots Overview.
//
// Checkpoints are periodically saved, storing in memory enough information
// for the process to restore the contents of all tracked memory as it
// rewinds to earlier checkpoitns. There are two components to a saved
// checkpoint:
//
// - Stack contents for each thread are completely saved on disk at each saved
//   checkpoint. This is handled by ThreadSnapshot.cpp
//
// - Heap and static memory contents (tracked memory) are saved in memory as
//   the contents of pages modified before either the the next saved checkpoint
//   or the current execution point (if this is the last saved checkpoint).
//   This is handled here.
//
// Heap memory is only tracked when allocated with TrackedMemoryKind.
//
// Snapshots of heap/static memory is modeled on the copy-on-write semantics
// used by fork. Instead of actually forking, we use write-protected memory and
// a fault handler to perform the copy-on-write, which both gives more control
// of the snapshot process and allows snapshots to be taken on platforms
// without fork (i.e. Windows). The following example shows how snapshots are
// generated:
//
// #1 Save Checkpoint A. The initial snapshot tabulates all allocated tracked
//    memory in the process, and write-protects all of it.
//
// #2 Write pages P0 and P1. Writing to the pages trips the fault handler. The
//    handler creates copies of the initial contents of P0 and P1 (P0a and P1a)
//    and unprotects the pages.
//
// #3 Save Checkpoint B. P0a and P1a, along with any other pages modified
//    between A and B, become associated with checkpoint A. All modified pages
//    are reprotected.
//
// #4 Write pages P1 and P2. Again, writing to the pages trips the fault
//    handler and copies P1b and P2b are created and the pages are unprotected.
//
// #5 Save Checkpoint C. P1b and P2b become associated with snapshot B, and the
//    modified pages are reprotected.
//
// If we were to then rewind from C to A, we would read and restore P1b/P2b,
// followed by P0a/P1a. All data associated with checkpoints A and later is
// discarded (we can only rewind; we cannot jump forward in time).
///////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////
// Snapshot Threads Overview.
//
// After step #3 above, the main thread has created a diff snapshot with the
// copies of the original contents of pages modified between two saved
// checkpoints. These page copies are initially all in memory. It is the
// responsibility of the snapshot threads to do the following:
//
// 1. When rewinding to the last saved checkpoint, snapshot threads are used to
//    restore the original contents of pages using their in-memory copies.
//
// There are a fixed number of snapshot threads that are spawned when the
// first checkpoint is saved. Threads are each responsible for distinct sets of
// heap memory pages (see AddDirtyPageToWorklist), avoiding synchronization
// issues between different snapshot threads.
///////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////////////
// Memory Snapshot Structures
///////////////////////////////////////////////////////////////////////////////

// A region of allocated memory which should be tracked by MemoryInfo.
struct AllocatedMemoryRegion {
  uint8_t* mBase;
  size_t mSize;
  bool mExecutable;

  AllocatedMemoryRegion() : mBase(nullptr), mSize(0), mExecutable(false) {}

  AllocatedMemoryRegion(uint8_t* aBase, size_t aSize, bool aExecutable)
      : mBase(aBase), mSize(aSize), mExecutable(aExecutable) {}

  // For sorting regions by base address.
  struct AddressSort {
    typedef void* Lookup;
    static void* getLookup(const AllocatedMemoryRegion& aRegion) {
      return aRegion.mBase;
    }
    static ssize_t compare(void* aAddress,
                           const AllocatedMemoryRegion& aRegion) {
      return (uint8_t*)aAddress - aRegion.mBase;
    }
  };

  // For sorting regions by size, from largest to smallest.
  struct SizeReverseSort {
    typedef size_t Lookup;
    static size_t getLookup(const AllocatedMemoryRegion& aRegion) {
      return aRegion.mSize;
    }
    static ssize_t compare(size_t aSize, const AllocatedMemoryRegion& aRegion) {
      return aRegion.mSize - aSize;
    }
  };
};

// Information about a page which was modified between two saved checkpoints.
struct DirtyPage {
  // Base address of the page.
  uint8_t* mBase;

  // Copy of the page at the first checkpoint. Written by the dirty memory
  // handler via HandleDirtyMemoryFault if this is in the active page set,
  // otherwise accessed by snapshot threads.
  uint8_t* mOriginal;

  bool mExecutable;

  DirtyPage(uint8_t* aBase, uint8_t* aOriginal, bool aExecutable)
      : mBase(aBase), mOriginal(aOriginal), mExecutable(aExecutable) {}

  struct AddressSort {
    typedef uint8_t* Lookup;
    static uint8_t* getLookup(const DirtyPage& aPage) { return aPage.mBase; }
    static ssize_t compare(uint8_t* aBase, const DirtyPage& aPage) {
      return aBase - aPage.mBase;
    }
  };
};

// A set of dirty pages that can be searched quickly.
typedef SplayTree<DirtyPage, DirtyPage::AddressSort,
                  AllocPolicy<MemoryKind::SortedDirtyPageSet>, 4>
    SortedDirtyPageSet;

// A set of dirty pages associated with some checkpoint.
struct DirtyPageSet {
  // Checkpoint associated with this set.
  CheckpointId mCheckpoint;

  // All dirty pages in the set. Pages may be added or destroyed by the main
  // thread when all other threads are idle, by the dirty memory handler when
  // it is active and this is the active page set, and by the snapshot thread
  // which owns this set.
  InfallibleVector<DirtyPage, 256, AllocPolicy<MemoryKind::DirtyPageSet>>
      mPages;

  explicit DirtyPageSet(const CheckpointId& aCheckpoint)
      : mCheckpoint(aCheckpoint) {}
};

// Worklist used by each snapshot thread.
struct SnapshotThreadWorklist {
  // Index into gMemoryInfo->mSnapshotWorklists of the thread.
  size_t mThreadIndex;

  // Record/replay ID of the thread.
  size_t mThreadId;

  // Sets of pages in the thread's worklist. Each set is for a different diff,
  // with the oldest checkpoints first.
  InfallibleVector<DirtyPageSet, 256, AllocPolicy<MemoryKind::Generic>> mSets;
};

// Structure used to coordinate activity between the main thread and all
// snapshot threads. The workflow with this structure is as follows:
//
// 1. The main thread calls ActivateBegin(), marking the condition as active
//    and notifying each snapshot thread. The main thread blocks in this call.
//
// 2. Each snapshot thread, maybe after waking up, checks the condition, does
//    any processing it needs to (knowing the main thread is blocked) and
//    then calls WaitUntilNoLongerActive(), blocking in the call.
//
// 3. Once all snapshot threads are blocked in WaitUntilNoLongerActive(), the
//    main thread is unblocked from ActivateBegin(). It can then do whatever
//    processing it needs to (knowing all snapshot threads are blocked) and
//    then calls ActivateEnd(), blocking in the call.
//
// 4. Snapshot threads are now unblocked from WaitUntilNoLongerActive(). The
//    main thread does not unblock from ActivateEnd() until all snapshot
//    threads have left WaitUntilNoLongerActive().
//
// The intent with this class is to ensure that the main thread knows exactly
// when the snapshot threads are operating and that there is no potential for
// races between them.
class SnapshotThreadCondition {
  Atomic<bool, SequentiallyConsistent, Behavior::DontPreserve> mActive;
  Atomic<int32_t, SequentiallyConsistent, Behavior::DontPreserve> mCount;

 public:
  void ActivateBegin();
  void ActivateEnd();
  bool IsActive();
  void WaitUntilNoLongerActive();
};

static const size_t NumSnapshotThreads = 8;

// A set of free regions in the process. There are two of these, for the
// free regions in tracked and untracked memory.
class FreeRegionSet {
  // Kind of memory being managed. This also describes the memory used by the
  // set itself.
  MemoryKind mKind;

  // Lock protecting contents of the structure.
  SpinLock mLock;

  // To avoid reentrancy issues when growing the set, a chunk of pages for
  // the splay tree is preallocated for use the next time the tree needs to
  // expand its size.
  static const size_t ChunkPages = 4;
  void* mNextChunk;

  // Ensure there is a chunk available for the splay tree.
  void MaybeRefillNextChunk(AutoSpinLock& aLockHeld);

  // Get the next chunk from the free region set for this memory kind.
  void* TakeNextChunk();

  struct MyAllocPolicy {
    FreeRegionSet& mSet;

    template <typename T>
    void free_(T* aPtr, size_t aSize) {
      MOZ_CRASH();
    }

    template <typename T>
    T* pod_malloc(size_t aNumElems) {
      MOZ_RELEASE_ASSERT(sizeof(T) * aNumElems <= ChunkPages * PageSize);
      return (T*)mSet.TakeNextChunk();
    }

    explicit MyAllocPolicy(FreeRegionSet& aSet) : mSet(aSet) {}
  };

  // All memory in gMemoryInfo->mTrackedRegions that is not in use at the
  // current point in execution.
  typedef SplayTree<AllocatedMemoryRegion,
                    AllocatedMemoryRegion::SizeReverseSort, MyAllocPolicy,
                    ChunkPages>
      Tree;
  Tree mRegions;

  void InsertLockHeld(void* aAddress, size_t aSize, AutoSpinLock& aLockHeld);
  void* ExtractLockHeld(size_t aSize, AutoSpinLock& aLockHeld);

 public:
  explicit FreeRegionSet(MemoryKind aKind)
      : mKind(aKind), mRegions(MyAllocPolicy(*this)) {}

  // Get the single region set for a given memory kind.
  static FreeRegionSet& Get(MemoryKind aKind);

  // Add a free region to the set.
  void Insert(void* aAddress, size_t aSize);

  // Remove a free region of the specified size. If aAddress is specified then
  // this address will be prioritized, but a different pointer may be returned.
  // The resulting memory will be zeroed.
  void* Extract(void* aAddress, size_t aSize);

  // Return whether a memory range intersects this set at all.
  bool Intersects(void* aAddress, size_t aSize);
};

// Information about the current memory state. The contents of this structure
// are in untracked memory.
struct MemoryInfo {
  // Whether new dirty pages or allocated regions are allowed.
  bool mMemoryChangesAllowed;

  // Untracked memory regions allocated before the first checkpoint. This is
  // only accessed on the main thread, and is not a vector because of reentrancy
  // issues.
  static const size_t MaxInitialUntrackedRegions = 512;
  AllocatedMemoryRegion mInitialUntrackedRegions[MaxInitialUntrackedRegions];
  SpinLock mInitialUntrackedRegionsLock;

  // All tracked memory in the process. This may be updated by any thread while
  // holding mTrackedRegionsLock.
  SplayTree<AllocatedMemoryRegion, AllocatedMemoryRegion::AddressSort,
            AllocPolicy<MemoryKind::TrackedRegions>, 4>
      mTrackedRegions;
  InfallibleVector<AllocatedMemoryRegion, 512,
                   AllocPolicy<MemoryKind::TrackedRegions>>
      mTrackedRegionsByAllocationOrder;
  SpinLock mTrackedRegionsLock;

  // Pages from |trackedRegions| modified since the last saved checkpoint.
  // Accessed by any thread (usually the dirty memory handler) when memory
  // changes are allowed, and by the main thread when memory changes are not
  // allowed.
  SortedDirtyPageSet mActiveDirty;
  SpinLock mActiveDirtyLock;

  // All untracked memory which is available for new allocations.
  FreeRegionSet mFreeUntrackedRegions;

  // Worklists for each snapshot thread.
  SnapshotThreadWorklist mSnapshotWorklists[NumSnapshotThreads];

  // Whether snapshot threads should update memory to that when the last saved
  // diff was started.
  SnapshotThreadCondition mSnapshotThreadsShouldRestore;

  // Whether snapshot threads should idle.
  SnapshotThreadCondition mSnapshotThreadsShouldIdle;

  // Counter used by the countdown thread.
  Atomic<size_t, SequentiallyConsistent, Behavior::DontPreserve> mCountdown;

  // Information for timers.
  double mStartTime;
  uint32_t mTimeHits[(size_t)TimerKind::Count];
  double mTimeTotals[(size_t)TimerKind::Count];

  // Information for memory allocation.
  Atomic<ssize_t, Relaxed, Behavior::DontPreserve>
      mMemoryBalance[(size_t)MemoryKind::Count];

  // Recent dirty memory faults.
  void* mDirtyMemoryFaults[50];

  // Whether RecordReplayDirective may crash this process.
  bool mIntentionalCrashesAllowed;

  // Whether the CrashSoon directive has been given to this process.
  bool mCrashSoon;

  MemoryInfo()
      : mMemoryChangesAllowed(true),
        mFreeUntrackedRegions(MemoryKind::FreeRegions),
        mStartTime(CurrentTime()),
        mIntentionalCrashesAllowed(true) {
    // The singleton MemoryInfo is allocated with zeroed memory, so other
    // fields do not need explicit initialization.
  }
};

static MemoryInfo* gMemoryInfo = nullptr;

void SetMemoryChangesAllowed(bool aAllowed) {
  MOZ_RELEASE_ASSERT(gMemoryInfo->mMemoryChangesAllowed == !aAllowed);
  gMemoryInfo->mMemoryChangesAllowed = aAllowed;
}

static void EnsureMemoryChangesAllowed() {
  while (!gMemoryInfo->mMemoryChangesAllowed) {
    ThreadYield();
  }
}

void StartCountdown(size_t aCount) { gMemoryInfo->mCountdown = aCount; }

AutoCountdown::AutoCountdown(size_t aCount) { StartCountdown(aCount); }

AutoCountdown::~AutoCountdown() { StartCountdown(0); }

#ifdef WANT_COUNTDOWN_THREAD

static void CountdownThreadMain(void*) {
  while (true) {
    if (gMemoryInfo->mCountdown && --gMemoryInfo->mCountdown == 0) {
      // When debugging hangs in the child process, we can break here in lldb
      // to inspect what the process is doing.
      child::ReportFatalError(Nothing(), "CountdownThread activated");
    }
    ThreadYield();
  }
}

#endif  // WANT_COUNTDOWN_THREAD

///////////////////////////////////////////////////////////////////////////////
// Profiling
///////////////////////////////////////////////////////////////////////////////

AutoTimer::AutoTimer(TimerKind aKind) : mKind(aKind), mStart(CurrentTime()) {}

AutoTimer::~AutoTimer() {
  if (gMemoryInfo) {
    gMemoryInfo->mTimeHits[(size_t)mKind]++;
    gMemoryInfo->mTimeTotals[(size_t)mKind] += CurrentTime() - mStart;
  }
}

static const char* gTimerKindNames[] = {
#define DefineTimerKindName(aKind) #aKind,
    ForEachTimerKind(DefineTimerKindName)
#undef DefineTimerKindName
};

void DumpTimers() {
  if (!gMemoryInfo) {
    return;
  }
  Print("Times %.2fs\n", (CurrentTime() - gMemoryInfo->mStartTime) / 1000000.0);
  for (size_t i = 0; i < (size_t)TimerKind::Count; i++) {
    uint32_t hits = gMemoryInfo->mTimeHits[i];
    double time = gMemoryInfo->mTimeTotals[i];
    Print("%s: %d hits, %.2fs\n", gTimerKindNames[i], (int)hits,
          time / 1000000.0);
  }
}

///////////////////////////////////////////////////////////////////////////////
// Directives
///////////////////////////////////////////////////////////////////////////////

void SetAllowIntentionalCrashes(bool aAllowed) {
  gMemoryInfo->mIntentionalCrashesAllowed = aAllowed;
}

extern "C" {

MOZ_EXPORT void RecordReplayInterface_InternalRecordReplayDirective(
    long aDirective) {
  switch ((Directive)aDirective) {
    case Directive::CrashSoon:
      gMemoryInfo->mCrashSoon = true;
      break;
    case Directive::MaybeCrash:
      if (gMemoryInfo->mIntentionalCrashesAllowed && gMemoryInfo->mCrashSoon) {
        PrintSpew("Intentionally Crashing!\n");
        MOZ_CRASH("RecordReplayDirective intentional crash");
      }
      gMemoryInfo->mCrashSoon = false;
      break;
    case Directive::AlwaysSaveTemporaryCheckpoints:
      navigation::AlwaysSaveTemporaryCheckpoints();
      break;
    case Directive::AlwaysMarkMajorCheckpoints:
      child::NotifyAlwaysMarkMajorCheckpoints();
      break;
    default:
      MOZ_CRASH("Unknown directive");
  }
}

}  // extern "C"

///////////////////////////////////////////////////////////////////////////////
// Snapshot Thread Conditions
///////////////////////////////////////////////////////////////////////////////

void SnapshotThreadCondition::ActivateBegin() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(!mActive);
  mActive = true;
  for (size_t i = 0; i < NumSnapshotThreads; i++) {
    Thread::Notify(gMemoryInfo->mSnapshotWorklists[i].mThreadId);
  }
  while (mCount != NumSnapshotThreads) {
    Thread::WaitNoIdle();
  }
}

void SnapshotThreadCondition::ActivateEnd() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(mActive);
  mActive = false;
  for (size_t i = 0; i < NumSnapshotThreads; i++) {
    Thread::Notify(gMemoryInfo->mSnapshotWorklists[i].mThreadId);
  }
  while (mCount) {
    Thread::WaitNoIdle();
  }
}

bool SnapshotThreadCondition::IsActive() {
  MOZ_RELEASE_ASSERT(!Thread::CurrentIsMainThread());
  return mActive;
}

void SnapshotThreadCondition::WaitUntilNoLongerActive() {
  MOZ_RELEASE_ASSERT(!Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(mActive);
  if (NumSnapshotThreads == ++mCount) {
    Thread::Notify(MainThreadId);
  }
  while (mActive) {
    Thread::WaitNoIdle();
  }
  if (0 == --mCount) {
    Thread::Notify(MainThreadId);
  }
}

///////////////////////////////////////////////////////////////////////////////
// Snapshot Page Allocation
///////////////////////////////////////////////////////////////////////////////

// Get a page in untracked memory that can be used as a copy of a tracked page.
static uint8_t* AllocatePageCopy() {
  return (uint8_t*)AllocateMemory(PageSize, MemoryKind::PageCopy);
}

// Free a page allocated by AllocatePageCopy.
static void FreePageCopy(uint8_t* aPage) {
  DeallocateMemory(aPage, PageSize, MemoryKind::PageCopy);
}

///////////////////////////////////////////////////////////////////////////////
// Page Fault Handling
///////////////////////////////////////////////////////////////////////////////

void MemoryMove(void* aDst, const void* aSrc, size_t aSize) {
  MOZ_RELEASE_ASSERT((size_t)aDst % sizeof(uint32_t) == 0);
  MOZ_RELEASE_ASSERT((size_t)aSrc % sizeof(uint32_t) == 0);
  MOZ_RELEASE_ASSERT(aSize % sizeof(uint32_t) == 0);
  MOZ_RELEASE_ASSERT((size_t)aDst <= (size_t)aSrc ||
                     (size_t)aDst >= (size_t)aSrc + aSize);

  uint32_t* ndst = (uint32_t*)aDst;
  const uint32_t* nsrc = (const uint32_t*)aSrc;
  for (size_t i = 0; i < aSize / sizeof(uint32_t); i++) {
    ndst[i] = nsrc[i];
  }
}

void MemoryZero(void* aDst, size_t aSize) {
  MOZ_RELEASE_ASSERT((size_t)aDst % sizeof(uint32_t) == 0);
  MOZ_RELEASE_ASSERT(aSize % sizeof(uint32_t) == 0);

  // Use volatile here to avoid annoying clang optimizations.
  volatile uint32_t* ndst = (uint32_t*)aDst;
  for (size_t i = 0; i < aSize / sizeof(uint32_t); i++) {
    ndst[i] = 0;
  }
}

// Return whether an address is in a tracked region. This excludes memory that
// is in an active new region and is not write protected.
static bool IsTrackedAddress(void* aAddress, bool* aExecutable) {
  AutoSpinLock lock(gMemoryInfo->mTrackedRegionsLock);
  Maybe<AllocatedMemoryRegion> region =
      gMemoryInfo->mTrackedRegions.lookupClosestLessOrEqual(aAddress);
  if (region.isSome() &&
      MemoryContains(region.ref().mBase, region.ref().mSize, aAddress)) {
    if (aExecutable) {
      *aExecutable = region.ref().mExecutable;
    }
    return true;
  }
  return false;
}

bool HandleDirtyMemoryFault(uint8_t* aAddress) {
  EnsureMemoryChangesAllowed();

  bool different = false;
  for (size_t i = ArrayLength(gMemoryInfo->mDirtyMemoryFaults) - 1; i; i--) {
    gMemoryInfo->mDirtyMemoryFaults[i] = gMemoryInfo->mDirtyMemoryFaults[i - 1];
    if (gMemoryInfo->mDirtyMemoryFaults[i] != aAddress) {
      different = true;
    }
  }
  gMemoryInfo->mDirtyMemoryFaults[0] = aAddress;
  if (!different) {
    Print("WARNING: Repeated accesses to the same dirty address %p\n",
          aAddress);
  }

  // Round down to the base of the page.
  aAddress = PageBase(aAddress);

  AutoSpinLock lock(gMemoryInfo->mActiveDirtyLock);

  // Check to see if this is already an active dirty page. Once a page has been
  // marked as dirty it will be accessible until the next checkpoint is saved,
  // but it's possible for multiple threads to access the same protected memory
  // before we have a chance to unprotect it, in which case we'll end up here
  // multiple times for the page.
  if (gMemoryInfo->mActiveDirty.maybeLookup(aAddress)) {
    return true;
  }

  // Crash if this address is not in a tracked region.
  bool executable;
  if (!IsTrackedAddress(aAddress, &executable)) {
    return false;
  }

  // Copy the page's original contents into the active dirty set, and unprotect
  // it so that execution can proceed.
  uint8_t* original = AllocatePageCopy();
  MemoryMove(original, aAddress, PageSize);
  gMemoryInfo->mActiveDirty.insert(aAddress,
                                   DirtyPage(aAddress, original, executable));
  DirectUnprotectMemory(aAddress, PageSize, executable);
  return true;
}

bool MemoryRangeIsTracked(void* aAddress, size_t aSize) {
  for (uint8_t* ptr = PageBase(aAddress); ptr < (uint8_t*)aAddress + aSize;
       ptr += PageSize) {
    if (!IsTrackedAddress(ptr, nullptr)) {
      return false;
    }
  }
  return true;
}

void UnrecoverableSnapshotFailure() {
  if (gMemoryInfo) {
    AutoSpinLock lock(gMemoryInfo->mTrackedRegionsLock);
    DirectUnprotectMemory(PageBase(&errno), PageSize, false);
    for (auto region : gMemoryInfo->mTrackedRegionsByAllocationOrder) {
      DirectUnprotectMemory(region.mBase, region.mSize, region.mExecutable,
                            /* aIgnoreFailures = */ true);
    }
  }
}

///////////////////////////////////////////////////////////////////////////////
// Initial Memory Region Processing
///////////////////////////////////////////////////////////////////////////////

void AddInitialUntrackedMemoryRegion(uint8_t* aBase, size_t aSize) {
  MOZ_RELEASE_ASSERT(!HasSavedCheckpoint());

  if (gInitializationFailureMessage) {
    return;
  }

  static void* gSkippedRegion;
  if (!gSkippedRegion) {
    // We are allocating gMemoryInfo itself, and will directly call this
    // function again shortly.
    gSkippedRegion = aBase;
    return;
  }
  MOZ_RELEASE_ASSERT(gSkippedRegion == gMemoryInfo);

  AutoSpinLock lock(gMemoryInfo->mInitialUntrackedRegionsLock);

  for (AllocatedMemoryRegion& region : gMemoryInfo->mInitialUntrackedRegions) {
    if (!region.mBase) {
      region.mBase = aBase;
      region.mSize = aSize;
      return;
    }
  }

  // If we end up here then MaxInitialUntrackedRegions should be larger.
  MOZ_CRASH();
}

static void RemoveInitialUntrackedRegion(uint8_t* aBase, size_t aSize) {
  MOZ_RELEASE_ASSERT(!HasSavedCheckpoint());
  AutoSpinLock lock(gMemoryInfo->mInitialUntrackedRegionsLock);

  for (AllocatedMemoryRegion& region : gMemoryInfo->mInitialUntrackedRegions) {
    if (region.mBase == aBase) {
      MOZ_RELEASE_ASSERT(region.mSize == aSize);
      region.mBase = nullptr;
      region.mSize = 0;
      return;
    }
  }
  MOZ_CRASH();
}

// Get information about the mapped region containing *aAddress, or the next
// mapped region afterwards if aAddress is not mapped. aAddress is updated to
// the start of that region, and aSize, aProtection, and aMaxProtection are
// updated with the size and protection status of the region. Returns false if
// there are no more mapped regions after *aAddress.
static bool QueryRegion(uint8_t** aAddress, size_t* aSize,
                        int* aProtection = nullptr,
                        int* aMaxProtection = nullptr) {
  mach_vm_address_t addr = (mach_vm_address_t)*aAddress;
  mach_vm_size_t nbytes;

  vm_region_basic_info_64 info;
  mach_msg_type_number_t info_count = sizeof(vm_region_basic_info_64);
  mach_port_t some_port;
  kern_return_t rv =
      mach_vm_region(mach_task_self(), &addr, &nbytes, VM_REGION_BASIC_INFO,
                     (vm_region_info_t)&info, &info_count, &some_port);
  if (rv == KERN_INVALID_ADDRESS) {
    return false;
  }
  MOZ_RELEASE_ASSERT(rv == KERN_SUCCESS);

  *aAddress = (uint8_t*)addr;
  *aSize = nbytes;
  if (aProtection) {
    *aProtection = info.protection;
  }
  if (aMaxProtection) {
    *aMaxProtection = info.max_protection;
  }
  return true;
}

static void MarkThreadStacksAsUntracked() {
  AutoPassThroughThreadEvents pt;

  // Thread stacks are excluded from the tracked regions.
  for (size_t i = MainThreadId; i <= MaxThreadId; i++) {
    Thread* thread = Thread::GetById(i);
    if (!thread->StackBase()) {
      continue;
    }

    AddInitialUntrackedMemoryRegion(thread->StackBase(), thread->StackSize());

    // Look for a mapped region with no access permissions immediately after
    // the thread stack's allocated region, and include this in the untracked
    // memory if found. This is done to avoid confusing breakpad, which will
    // scan the allocated memory in this process and will not correctly
    // determine stack boundaries if we track these trailing regions and end up
    // marking them as readable.

    // Find the mapped region containing the end of the thread's stack.
    uint8_t* base = thread->StackBase() + thread->StackSize() - 1;
    size_t size;
    if (!QueryRegion(&base, &size)) {
      MOZ_CRASH("Could not find memory region information for thread stack");
    }

    // Sanity check the region size. Note that we don't mark this entire region
    // as untracked, since it may contain TLS data which should be tracked.
    MOZ_RELEASE_ASSERT(base + size >=
                       thread->StackBase() + thread->StackSize());

    uint8_t* trailing = base + size;
    size_t trailingSize;
    int protection;
    if (QueryRegion(&trailing, &trailingSize, &protection)) {
      if (trailing == base + size && protection == 0) {
        AddInitialUntrackedMemoryRegion(trailing, trailingSize);
      }
    }
  }
}

// Given an address region [*aAddress, *aAddress + *aSize], return true if
// there is any intersection with an excluded region
// [aExclude, aExclude + aExcludeSize], set *aSize to contain the subregion
// starting at aAddress which which is not excluded, and *aRemaining and
// *aRemainingSize to any additional subregion which is not excluded.
static bool MaybeExtractMemoryRegion(uint8_t* aAddress, size_t* aSize,
                                     uint8_t** aRemaining,
                                     size_t* aRemainingSize, uint8_t* aExclude,
                                     size_t aExcludeSize) {
  uint8_t* addrLimit = aAddress + *aSize;

  // Expand the excluded region out to the containing page boundaries.
  MOZ_RELEASE_ASSERT((size_t)aExclude % PageSize == 0);
  aExcludeSize = RoundupSizeToPageBoundary(aExcludeSize);

  uint8_t* excludeLimit = aExclude + aExcludeSize;

  if (excludeLimit <= aAddress || addrLimit <= aExclude) {
    // No intersection.
    return false;
  }

  *aSize = std::max<ssize_t>(aExclude - aAddress, 0);
  if (aRemaining) {
    *aRemaining = excludeLimit;
    *aRemainingSize = std::max<ssize_t>(addrLimit - *aRemaining, 0);
  }
  return true;
}

// Set *aSize to describe the number of bytes starting at aAddress that should
// be considered tracked memory. *aRemaining and *aRemainingSize are set to any
// remaining portion of the initial region after the first excluded portion
// that is found.
static void ExtractTrackedInitialMemoryRegion(uint8_t* aAddress, size_t* aSize,
                                              uint8_t** aRemaining,
                                              size_t* aRemainingSize) {
  // Look for the earliest untracked region which intersects the given region.
  const AllocatedMemoryRegion* earliestIntersect = nullptr;
  for (const AllocatedMemoryRegion& region :
       gMemoryInfo->mInitialUntrackedRegions) {
    size_t size = *aSize;
    if (MaybeExtractMemoryRegion(aAddress, &size, nullptr, nullptr,
                                 region.mBase, region.mSize)) {
      // There was an intersection.
      if (!earliestIntersect || region.mBase < earliestIntersect->mBase) {
        earliestIntersect = &region;
      }
    }
  }

  if (earliestIntersect) {
    if (!MaybeExtractMemoryRegion(aAddress, aSize, aRemaining, aRemainingSize,
                                  earliestIntersect->mBase,
                                  earliestIntersect->mSize)) {
      MOZ_CRASH();
    }
  } else {
    // If there is no intersection then the entire region is tracked.
    *aRemaining = aAddress + *aSize;
    *aRemainingSize = 0;
  }
}

static void AddTrackedRegion(uint8_t* aAddress, size_t aSize,
                             bool aExecutable) {
  if (aSize) {
    AutoSpinLock lock(gMemoryInfo->mTrackedRegionsLock);
    gMemoryInfo->mTrackedRegions.insert(
        aAddress, AllocatedMemoryRegion(aAddress, aSize, aExecutable));
    gMemoryInfo->mTrackedRegionsByAllocationOrder.emplaceBack(aAddress, aSize,
                                                              aExecutable);
  }
}

// Add any tracked subregions of [aAddress, aAddress + aSize].
void AddInitialTrackedMemoryRegions(uint8_t* aAddress, size_t aSize,
                                    bool aExecutable) {
  while (aSize) {
    uint8_t* remaining;
    size_t remainingSize;
    ExtractTrackedInitialMemoryRegion(aAddress, &aSize, &remaining,
                                      &remainingSize);

    AddTrackedRegion(aAddress, aSize, aExecutable);

    aAddress = remaining;
    aSize = remainingSize;
  }
}

static void UpdateNumTrackedRegionsForSnapshot();

// Fill in the set of tracked memory regions that are currently mapped within
// this process.
static void ProcessAllInitialMemoryRegions() {
  MOZ_ASSERT(!AreThreadEventsPassedThrough());

  {
    AutoPassThroughThreadEvents pt;
    for (uint8_t* addr = nullptr;;) {
      size_t size;
      int maxProtection;
      if (!QueryRegion(&addr, &size, nullptr, &maxProtection)) {
        break;
      }

      // Consider all memory regions that can possibly be written to, even if
      // they aren't currently writable.
      if (maxProtection & VM_PROT_WRITE) {
        MOZ_RELEASE_ASSERT(maxProtection & VM_PROT_READ);
        AddInitialTrackedMemoryRegions(addr, size,
                                       maxProtection & VM_PROT_EXECUTE);
      }

      addr += size;
    }
  }

  UpdateNumTrackedRegionsForSnapshot();

  // Write protect all tracked memory.
  AutoDisallowMemoryChanges disallow;
  for (const AllocatedMemoryRegion& region :
       gMemoryInfo->mTrackedRegionsByAllocationOrder) {
    DirectWriteProtectMemory(region.mBase, region.mSize, region.mExecutable);
  }
}

///////////////////////////////////////////////////////////////////////////////
// Free Region Management
///////////////////////////////////////////////////////////////////////////////

// All memory in gMemoryInfo->mTrackedRegions that is not in use at the current
// point in execution.
static FreeRegionSet gFreeRegions(MemoryKind::Tracked);

// The size of gMemoryInfo->mTrackedRegionsByAllocationOrder we expect to see
// at the point of the last saved checkpoint.
static size_t gNumTrackedRegions;

static void UpdateNumTrackedRegionsForSnapshot() {
  MOZ_ASSERT(Thread::CurrentIsMainThread());
  gNumTrackedRegions = gMemoryInfo->mTrackedRegionsByAllocationOrder.length();
}

void FixupFreeRegionsAfterRewind() {
  // All memory that has been allocated since the associated checkpoint was
  // reached is now free, and may be reused for new allocations.
  size_t newTrackedRegions =
      gMemoryInfo->mTrackedRegionsByAllocationOrder.length();
  for (size_t i = gNumTrackedRegions; i < newTrackedRegions; i++) {
    const AllocatedMemoryRegion& region =
        gMemoryInfo->mTrackedRegionsByAllocationOrder[i];
    gFreeRegions.Insert(region.mBase, region.mSize);
  }
  gNumTrackedRegions = newTrackedRegions;
}

/* static */ FreeRegionSet& FreeRegionSet::Get(MemoryKind aKind) {
  return (aKind == MemoryKind::Tracked) ? gFreeRegions
                                        : gMemoryInfo->mFreeUntrackedRegions;
}

void* FreeRegionSet::TakeNextChunk() {
  MOZ_RELEASE_ASSERT(mNextChunk);
  void* res = mNextChunk;
  mNextChunk = nullptr;
  return res;
}

void FreeRegionSet::InsertLockHeld(void* aAddress, size_t aSize,
                                   AutoSpinLock& aLockHeld) {
  mRegions.insert(aSize,
                  AllocatedMemoryRegion((uint8_t*)aAddress, aSize, true));
}

void FreeRegionSet::MaybeRefillNextChunk(AutoSpinLock& aLockHeld) {
  if (mNextChunk) {
    return;
  }

  // Look for a free region we can take the next chunk from.
  size_t size = ChunkPages * PageSize;
  gMemoryInfo->mMemoryBalance[(size_t)mKind] += size;

  mNextChunk = ExtractLockHeld(size, aLockHeld);

  if (!mNextChunk) {
    // Allocate memory from the system.
    mNextChunk = DirectAllocateMemory(nullptr, size);
    RegisterAllocatedMemory(mNextChunk, size, mKind);
  }
}

void FreeRegionSet::Insert(void* aAddress, size_t aSize) {
  MOZ_RELEASE_ASSERT(aAddress && aAddress == PageBase(aAddress));
  MOZ_RELEASE_ASSERT(aSize && aSize == RoundupSizeToPageBoundary(aSize));

  AutoSpinLock lock(mLock);

  MaybeRefillNextChunk(lock);
  InsertLockHeld(aAddress, aSize, lock);
}

void* FreeRegionSet::ExtractLockHeld(size_t aSize, AutoSpinLock& aLockHeld) {
  Maybe<AllocatedMemoryRegion> best =
      mRegions.lookupClosestLessOrEqual(aSize, /* aRemove = */ true);
  if (best.isSome()) {
    MOZ_RELEASE_ASSERT(best.ref().mSize >= aSize);
    uint8_t* res = best.ref().mBase;
    if (best.ref().mSize > aSize) {
      InsertLockHeld(res + aSize, best.ref().mSize - aSize, aLockHeld);
    }
    MemoryZero(res, aSize);
    return res;
  }
  return nullptr;
}

void* FreeRegionSet::Extract(void* aAddress, size_t aSize) {
  MOZ_RELEASE_ASSERT(aAddress == PageBase(aAddress));
  MOZ_RELEASE_ASSERT(aSize && aSize == RoundupSizeToPageBoundary(aSize));

  AutoSpinLock lock(mLock);

  if (aAddress) {
    MaybeRefillNextChunk(lock);

    // We were given a point at which to try to place the allocation. Look for
    // a free region which contains [aAddress, aAddress + aSize] entirely.
    for (typename Tree::Iter iter = mRegions.begin(); !iter.done(); ++iter) {
      uint8_t* regionBase = iter.ref().mBase;
      uint8_t* regionExtent = regionBase + iter.ref().mSize;
      uint8_t* addrBase = (uint8_t*)aAddress;
      uint8_t* addrExtent = addrBase + aSize;
      if (regionBase <= addrBase && regionExtent >= addrExtent) {
        iter.removeEntry();
        if (regionBase < addrBase) {
          InsertLockHeld(regionBase, addrBase - regionBase, lock);
        }
        if (regionExtent > addrExtent) {
          InsertLockHeld(addrExtent, regionExtent - addrExtent, lock);
        }
        MemoryZero(aAddress, aSize);
        return aAddress;
      }
    }
    // Fall through and look for a free region at another address.
  }

  // No address hint, look for the smallest free region which is larger than
  // the desired allocation size.
  return ExtractLockHeld(aSize, lock);
}

bool FreeRegionSet::Intersects(void* aAddress, size_t aSize) {
  AutoSpinLock lock(mLock);
  for (typename Tree::Iter iter = mRegions.begin(); !iter.done(); ++iter) {
    if (MemoryIntersects(iter.ref().mBase, iter.ref().mSize, aAddress, aSize)) {
      return true;
    }
  }
  return false;
}

///////////////////////////////////////////////////////////////////////////////
// Memory Management
///////////////////////////////////////////////////////////////////////////////

void RegisterAllocatedMemory(void* aBaseAddress, size_t aSize,
                             MemoryKind aKind) {
  MOZ_RELEASE_ASSERT(aBaseAddress == PageBase(aBaseAddress));
  MOZ_RELEASE_ASSERT(aSize == RoundupSizeToPageBoundary(aSize));

  uint8_t* aAddress = reinterpret_cast<uint8_t*>(aBaseAddress);

  if (aKind != MemoryKind::Tracked) {
    if (!HasSavedCheckpoint()) {
      AddInitialUntrackedMemoryRegion(aAddress, aSize);
    }
  } else if (HasSavedCheckpoint()) {
    EnsureMemoryChangesAllowed();
    DirectWriteProtectMemory(aAddress, aSize, true);
    AddTrackedRegion(aAddress, aSize, true);
  }
}

void CheckFixedMemory(void* aAddress, size_t aSize) {
  MOZ_RELEASE_ASSERT(aAddress == PageBase(aAddress));
  MOZ_RELEASE_ASSERT(aSize == RoundupSizeToPageBoundary(aSize));

  if (!HasSavedCheckpoint()) {
    return;
  }

  {
    // The memory should already be tracked. Check each page in the allocation
    // because there might be tracked regions adjacent to one another, neither
    // of which entirely contains this memory.
    AutoSpinLock lock(gMemoryInfo->mTrackedRegionsLock);
    for (size_t offset = 0; offset < aSize; offset += PageSize) {
      uint8_t* page = (uint8_t*)aAddress + offset;
      Maybe<AllocatedMemoryRegion> region =
          gMemoryInfo->mTrackedRegions.lookupClosestLessOrEqual(page);
      if (!region.isSome() ||
          !MemoryContains(region.ref().mBase, region.ref().mSize, page,
                          PageSize)) {
        MOZ_CRASH("Fixed memory is not tracked!");
      }
    }
  }

  // The memory should not be free.
  if (gFreeRegions.Intersects(aAddress, aSize)) {
    MOZ_CRASH("Fixed memory is currently free!");
  }
}

void RestoreWritableFixedMemory(void* aAddress, size_t aSize) {
  MOZ_RELEASE_ASSERT(aAddress == PageBase(aAddress));
  MOZ_RELEASE_ASSERT(aSize == RoundupSizeToPageBoundary(aSize));

  if (!HasSavedCheckpoint()) {
    return;
  }

  AutoSpinLock lock(gMemoryInfo->mActiveDirtyLock);
  for (size_t offset = 0; offset < aSize; offset += PageSize) {
    uint8_t* page = (uint8_t*)aAddress + offset;
    if (gMemoryInfo->mActiveDirty.maybeLookup(page)) {
      DirectUnprotectMemory(page, PageSize, true);
    }
  }
}

void* AllocateMemoryTryAddress(void* aAddress, size_t aSize, MemoryKind aKind) {
  MOZ_RELEASE_ASSERT(aAddress == PageBase(aAddress));
  aSize = RoundupSizeToPageBoundary(aSize);

  if (gMemoryInfo) {
    gMemoryInfo->mMemoryBalance[(size_t)aKind] += aSize;
  }

  if (HasSavedCheckpoint()) {
    if (void* res = FreeRegionSet::Get(aKind).Extract(aAddress, aSize)) {
      return res;
    }
  }

  void* res = DirectAllocateMemory(aAddress, aSize);
  RegisterAllocatedMemory(res, aSize, aKind);
  return res;
}

void* AllocateMemory(size_t aSize, MemoryKind aKind) {
  if (!IsReplaying()) {
    return DirectAllocateMemory(nullptr, aSize);
  }
  return AllocateMemoryTryAddress(nullptr, aSize, aKind);
}

void DeallocateMemory(void* aAddress, size_t aSize, MemoryKind aKind) {
  // Round the supplied region to the containing page boundaries.
  aSize += (uint8_t*)aAddress - PageBase(aAddress);
  aAddress = PageBase(aAddress);
  aSize = RoundupSizeToPageBoundary(aSize);

  if (!aAddress || !aSize) {
    return;
  }

  if (gMemoryInfo) {
    gMemoryInfo->mMemoryBalance[(size_t)aKind] -= aSize;
  }

  // Memory is returned to the system before saving the first checkpoint.
  if (!HasSavedCheckpoint()) {
    if (IsReplaying() && aKind != MemoryKind::Tracked) {
      RemoveInitialUntrackedRegion((uint8_t*)aAddress, aSize);
    }
    DirectDeallocateMemory(aAddress, aSize);
    return;
  }

  if (aKind == MemoryKind::Tracked) {
    // For simplicity, all free regions must be executable, so ignore
    // deallocated memory in regions that are not executable.
    bool executable;
    if (!IsTrackedAddress(aAddress, &executable) || !executable) {
      return;
    }
  }

  // Mark this region as free, but do not unmap it. It will become usable for
  // later allocations, but will not need to be remapped if we end up
  // rewinding to a point where this memory was in use.
  FreeRegionSet::Get(aKind).Insert(aAddress, aSize);
}

///////////////////////////////////////////////////////////////////////////////
// Snapshot Threads
///////////////////////////////////////////////////////////////////////////////

// While on a snapshot thread, restore the contents of all pages belonging to
// this thread which were modified since the last recorded diff snapshot.
static void SnapshotThreadRestoreLastDiffSnapshot(
    SnapshotThreadWorklist* aWorklist) {
  CheckpointId checkpoint = GetLastSavedCheckpoint();

  DirtyPageSet& set = aWorklist->mSets.back();
  MOZ_RELEASE_ASSERT(set.mCheckpoint == checkpoint);

  // Copy the original contents of all pages.
  for (size_t index = 0; index < set.mPages.length(); index++) {
    const DirtyPage& page = set.mPages[index];
    MOZ_RELEASE_ASSERT(page.mOriginal);
    DirectUnprotectMemory(page.mBase, PageSize, page.mExecutable);
    MemoryMove(page.mBase, page.mOriginal, PageSize);
    DirectWriteProtectMemory(page.mBase, PageSize, page.mExecutable);
    FreePageCopy(page.mOriginal);
  }

  // Remove the set from the worklist, if necessary.
  if (!aWorklist->mSets.empty()) {
    MOZ_RELEASE_ASSERT(&set == &aWorklist->mSets.back());
    aWorklist->mSets.popBack();
  }
}

// Start routine for a snapshot thread.
void SnapshotThreadMain(void* aArgument) {
  size_t threadIndex = (size_t)aArgument;
  SnapshotThreadWorklist* worklist =
      &gMemoryInfo->mSnapshotWorklists[threadIndex];
  worklist->mThreadIndex = threadIndex;

  while (true) {
    // If the main thread is waiting for us to restore the most recent diff,
    // then do so and notify the main thread we finished.
    if (gMemoryInfo->mSnapshotThreadsShouldRestore.IsActive()) {
      SnapshotThreadRestoreLastDiffSnapshot(worklist);
      gMemoryInfo->mSnapshotThreadsShouldRestore.WaitUntilNoLongerActive();
    }

    // Idle if the main thread wants us to.
    if (gMemoryInfo->mSnapshotThreadsShouldIdle.IsActive()) {
      gMemoryInfo->mSnapshotThreadsShouldIdle.WaitUntilNoLongerActive();
    }

    // Idle until notified by the main thread.
    Thread::WaitNoIdle();
  }
}

// An alternative to memcmp that can be called from any place.
static bool MemoryEquals(void* aDst, void* aSrc, size_t aSize) {
  MOZ_ASSERT((size_t)aDst % sizeof(size_t) == 0);
  MOZ_ASSERT((size_t)aSrc % sizeof(size_t) == 0);
  MOZ_ASSERT(aSize % sizeof(size_t) == 0);

  size_t* ndst = (size_t*)aDst;
  size_t* nsrc = (size_t*)aSrc;
  for (size_t i = 0; i < aSize / sizeof(size_t); i++) {
    if (ndst[i] != nsrc[i]) {
      return false;
    }
  }
  return true;
}

// Add a page to the last set in some snapshot thread's worklist. This is
// called on the main thread while the snapshot thread is idle.
static void AddDirtyPageToWorklist(uint8_t* aAddress, uint8_t* aOriginal,
                                   bool aExecutable) {
  // Distribute pages to snapshot threads using the base address of a page.
  // This guarantees that the same page will be consistently assigned to the
  // same thread as different snapshots are taken.
  MOZ_ASSERT((size_t)aAddress % PageSize == 0);
  if (MemoryEquals(aAddress, aOriginal, PageSize)) {
    FreePageCopy(aOriginal);
  } else {
    size_t pageIndex = ((size_t)aAddress / PageSize) % NumSnapshotThreads;
    SnapshotThreadWorklist* worklist =
        &gMemoryInfo->mSnapshotWorklists[pageIndex];
    MOZ_RELEASE_ASSERT(!worklist->mSets.empty());
    DirtyPageSet& set = worklist->mSets.back();
    MOZ_RELEASE_ASSERT(set.mCheckpoint == GetLastSavedCheckpoint());
    set.mPages.emplaceBack(aAddress, aOriginal, aExecutable);
  }
}

///////////////////////////////////////////////////////////////////////////////
// Snapshot Interface
///////////////////////////////////////////////////////////////////////////////

void InitializeMemorySnapshots() {
  MOZ_RELEASE_ASSERT(gMemoryInfo == nullptr);
  void* memory = AllocateMemory(sizeof(MemoryInfo), MemoryKind::Generic);
  gMemoryInfo = new (memory) MemoryInfo();

  // Mark gMemoryInfo as untracked. See AddInitialUntrackedMemoryRegion.
  AddInitialUntrackedMemoryRegion(reinterpret_cast<uint8_t*>(memory),
                                  sizeof(MemoryInfo));
}

void InitializeCountdownThread() {
#ifdef WANT_COUNTDOWN_THREAD
  Thread::SpawnNonRecordedThread(CountdownThreadMain, nullptr);
#endif
}

void TakeFirstMemorySnapshot() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(gMemoryInfo->mTrackedRegions.empty());

  // Spawn all snapshot threads.
  {
    AutoPassThroughThreadEvents pt;

    for (size_t i = 0; i < NumSnapshotThreads; i++) {
      Thread* thread =
          Thread::SpawnNonRecordedThread(SnapshotThreadMain, (void*)i);
      gMemoryInfo->mSnapshotWorklists[i].mThreadId = thread->Id();
    }
  }

  // All threads should have been created by now.
  MarkThreadStacksAsUntracked();

  // Fill in the tracked regions for the process.
  ProcessAllInitialMemoryRegions();
}

void TakeDiffMemorySnapshot() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());

  UpdateNumTrackedRegionsForSnapshot();

  AutoDisallowMemoryChanges disallow;

  // Stop all snapshot threads while we modify their worklists.
  gMemoryInfo->mSnapshotThreadsShouldIdle.ActivateBegin();

  // Add a DirtyPageSet to each snapshot thread's worklist for this snapshot.
  for (size_t i = 0; i < NumSnapshotThreads; i++) {
    SnapshotThreadWorklist* worklist = &gMemoryInfo->mSnapshotWorklists[i];
    worklist->mSets.emplaceBack(GetLastSavedCheckpoint());
  }

  // Distribute remaining active dirty pages to the snapshot thread worklists.
  for (SortedDirtyPageSet::Iter iter = gMemoryInfo->mActiveDirty.begin();
       !iter.done(); ++iter) {
    AddDirtyPageToWorklist(iter.ref().mBase, iter.ref().mOriginal,
                           iter.ref().mExecutable);
    DirectWriteProtectMemory(iter.ref().mBase, PageSize,
                             iter.ref().mExecutable);
  }

  gMemoryInfo->mActiveDirty.clear();

  // Allow snapshot threads to resume execution.
  gMemoryInfo->mSnapshotThreadsShouldIdle.ActivateEnd();
}

void RestoreMemoryToLastSavedCheckpoint() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(!gMemoryInfo->mMemoryChangesAllowed);

  // Restore all dirty regions that have been modified since the last
  // checkpoint was saved/restored.
  for (SortedDirtyPageSet::Iter iter = gMemoryInfo->mActiveDirty.begin();
       !iter.done(); ++iter) {
    MemoryMove(iter.ref().mBase, iter.ref().mOriginal, PageSize);
    FreePageCopy(iter.ref().mOriginal);
    DirectWriteProtectMemory(iter.ref().mBase, PageSize,
                             iter.ref().mExecutable);
  }
  gMemoryInfo->mActiveDirty.clear();
}

void RestoreMemoryToLastSavedDiffCheckpoint() {
  MOZ_RELEASE_ASSERT(Thread::CurrentIsMainThread());
  MOZ_RELEASE_ASSERT(!gMemoryInfo->mMemoryChangesAllowed);
  MOZ_RELEASE_ASSERT(gMemoryInfo->mActiveDirty.empty());

  // Wait while the snapshot threads restore all pages modified since the diff
  // snapshot was recorded.
  gMemoryInfo->mSnapshotThreadsShouldRestore.ActivateBegin();
  gMemoryInfo->mSnapshotThreadsShouldRestore.ActivateEnd();
}

}  // namespace recordreplay
}  // namespace mozilla