/* 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/. */

#ifndef gc_StoreBuffer_h
#define gc_StoreBuffer_h

#include "mozilla/Assertions.h"
#include "mozilla/Attributes.h"
#include "mozilla/HashFunctions.h"
#include "mozilla/ReentrancyGuard.h"

#include <algorithm>

#include "ds/BitArray.h"
#include "ds/LifoAlloc.h"
#include "gc/Cell.h"
#include "gc/Nursery.h"
#include "js/AllocPolicy.h"
#include "js/UniquePtr.h"
#include "wasm/WasmAnyRef.h"

namespace JS {
struct GCSizes;
}  // namespace JS

namespace js {

class NativeObject;

namespace gc {

class Arena;
class ArenaCellSet;
class GCRuntime;

/*
 * BufferableRef represents an abstract reference for use in the generational
 * GC's remembered set. Entries in the store buffer that cannot be represented
 * with the simple pointer-to-a-pointer scheme must derive from this class and
 * use the generic store buffer interface.
 *
 * A single BufferableRef entry in the generic buffer can represent many entries
 * in the remembered set.  For example js::OrderedHashTableRef represents all
 * the incoming edges corresponding to keys in an ordered hash table.
 */
class BufferableRef {
 public:
  virtual void trace(JSTracer* trc) = 0;
  bool maybeInRememberedSet(const Nursery&) const { return true; }
};

using EdgeSet = HashSet<void*, PointerHasher<void*>, SystemAllocPolicy>;

/* The size of a single block of store buffer storage space. */
static const size_t LifoAllocBlockSize = 8 * 1024;

/*
 * The StoreBuffer observes all writes that occur in the system and performs
 * efficient filtering of them to derive a remembered set for nursery GC.
 */
class StoreBuffer {
  friend class mozilla::ReentrancyGuard;

  enum class PutResult { OK, AboutToOverflow };

  /*
   * This buffer holds only a single type of edge. Using this buffer is more
   * efficient than the generic buffer when many writes will be to the same
   * type of edge: e.g. Value or Cell*.
   */
  template <typename T>
  struct MonoTypeBuffer {
    /* The canonical set of stores. */
    using StoreSet = HashSet<T, typename T::Hasher, SystemAllocPolicy>;
    StoreSet stores_;
    size_t maxEntries_ = 0;

    /*
     * A one element cache in front of the canonical set to speed up
     * temporary instances of HeapPtr.
     */
    T last_ = T();

    MonoTypeBuffer() = default;

    MonoTypeBuffer(const MonoTypeBuffer& other) = delete;
    MonoTypeBuffer& operator=(const MonoTypeBuffer& other) = delete;

    inline MonoTypeBuffer(MonoTypeBuffer&& other);
    inline MonoTypeBuffer& operator=(MonoTypeBuffer&& other);

    void setSize(size_t entryCount);

    bool isEmpty() const;
    void clear();

    /* Add one item to the buffer. */
    PutResult put(const T& t) {
      PutResult r = sinkStore();
      last_ = t;
      return r;
    }

    /* Remove an item from the store buffer. */
    void unput(const T& v) {
      // Fast, hashless remove of last put.
      if (last_ == v) {
        last_ = T();
        return;
      }
      stores_.remove(v);
    }

    /* Move any buffered stores to the canonical store set. */
    PutResult sinkStore() {
      if (last_) {
        AutoEnterOOMUnsafeRegion oomUnsafe;
        if (!stores_.put(last_)) {
          oomUnsafe.crash("Failed to allocate for MonoTypeBuffer::put.");
        }
      }
      last_ = T();

      if (stores_.count() >= maxEntries_) {
        return PutResult::AboutToOverflow;
      }

      return PutResult::OK;
    }

    /* Trace the source of all edges in the store buffer. */
    void trace(TenuringTracer& mover, StoreBuffer* owner);

    size_t sizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf);
  };

  struct WholeCellBuffer {
    UniquePtr<LifoAlloc> storage_;
    size_t maxSize_ = 0;
    const Cell* last_ = nullptr;

    WholeCellBuffer() = default;

    WholeCellBuffer(const WholeCellBuffer& other) = delete;
    WholeCellBuffer& operator=(const WholeCellBuffer& other) = delete;

    inline WholeCellBuffer(WholeCellBuffer&& other);
    inline WholeCellBuffer& operator=(WholeCellBuffer&& other);

    [[nodiscard]] bool init();
    void setSize(size_t entryCount);

    bool isEmpty() const;
    void clear();

    bool isAboutToOverflow() const {
      return !storage_->isEmpty() && storage_->used() >= maxSize_;
    }

    void trace(TenuringTracer& mover, StoreBuffer* owner);

    inline void put(const Cell* cell);
    inline void putDontCheckLast(const Cell* cell);

    size_t sizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf);

    const Cell** lastBufferedPtr() { return &last_; }

   private:
    ArenaCellSet* allocateCellSet(Arena* arena);
  };

  struct GenericBuffer {
    UniquePtr<LifoAlloc> storage_;
    size_t maxSize_ = 0;

    GenericBuffer() = default;

    GenericBuffer(const GenericBuffer& other) = delete;
    GenericBuffer& operator=(const GenericBuffer& other) = delete;

    inline GenericBuffer(GenericBuffer&& other);
    inline GenericBuffer& operator=(GenericBuffer&& other);

    [[nodiscard]] bool init();
    void setSize(size_t entryCount);

    bool isEmpty() const;
    void clear();

    bool isAboutToOverflow() const {
      return !storage_->isEmpty() && storage_->used() >= maxSize_;
    }

    /* Trace all generic edges. */
    void trace(JSTracer* trc, StoreBuffer* owner);

    template <typename T>
    PutResult put(const T& t) {
      static_assert(std::is_base_of_v<BufferableRef, T>);
      MOZ_ASSERT(storage_);

      AutoEnterOOMUnsafeRegion oomUnsafe;
      unsigned size = sizeof(T);
      unsigned* sizep = storage_->pod_malloc<unsigned>();
      if (!sizep) {
        oomUnsafe.crash("Failed to allocate for GenericBuffer::put.");
      }
      *sizep = size;

      T* tp = storage_->new_<T>(t);
      if (!tp) {
        oomUnsafe.crash("Failed to allocate for GenericBuffer::put.");
      }

      if (isAboutToOverflow()) {
        return PutResult::AboutToOverflow;
      }

      return PutResult::OK;
    }

    size_t sizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf);
  };

  template <typename Edge>
  struct PointerEdgeHasher {
    using Lookup = Edge;
    static HashNumber hash(const Lookup& l) {
      return mozilla::HashGeneric(l.edge);
    }
    static bool match(const Edge& k, const Lookup& l) { return k == l; }
  };

  template <typename T>
  struct CellPtrEdge {
    T** edge = nullptr;

    CellPtrEdge() = default;
    explicit CellPtrEdge(T** v) : edge(v) {}
    bool operator==(const CellPtrEdge& other) const {
      return edge == other.edge;
    }
    bool operator!=(const CellPtrEdge& other) const {
      return edge != other.edge;
    }

    bool maybeInRememberedSet(const Nursery& nursery) const {
      MOZ_ASSERT(IsInsideNursery(*edge));
      return !nursery.isInside(edge);
    }

    void trace(TenuringTracer& mover) const;

    explicit operator bool() const { return edge != nullptr; }

    using Hasher = PointerEdgeHasher<CellPtrEdge<T>>;
  };

  using ObjectPtrEdge = CellPtrEdge<JSObject>;
  using StringPtrEdge = CellPtrEdge<JSString>;
  using BigIntPtrEdge = CellPtrEdge<JS::BigInt>;
  using GetterSetterPtrEdge = CellPtrEdge<js::GetterSetter>;

  struct ValueEdge {
    JS::Value* edge;

    ValueEdge() : edge(nullptr) {}
    explicit ValueEdge(JS::Value* v) : edge(v) {}
    bool operator==(const ValueEdge& other) const { return edge == other.edge; }
    bool operator!=(const ValueEdge& other) const { return edge != other.edge; }

    bool isGCThing() const { return edge->isGCThing(); }

    Cell* deref() const {
      return isGCThing() ? static_cast<Cell*>(edge->toGCThing()) : nullptr;
    }

    bool maybeInRememberedSet(const Nursery& nursery) const {
      MOZ_ASSERT(IsInsideNursery(deref()));
      return !nursery.isInside(edge);
    }

    void trace(TenuringTracer& mover) const;

    explicit operator bool() const { return edge != nullptr; }

    using Hasher = PointerEdgeHasher<ValueEdge>;
  };

  struct SlotsEdge {
    // These definitions must match those in HeapSlot::Kind.
    const static int SlotKind = 0;
    const static int ElementKind = 1;

    uintptr_t objectAndKind_;  // NativeObject* | Kind
    uint32_t start_;
    uint32_t count_;

    SlotsEdge() : objectAndKind_(0), start_(0), count_(0) {}
    SlotsEdge(NativeObject* object, int kind, uint32_t start, uint32_t count)
        : objectAndKind_(uintptr_t(object) | kind),
          start_(start),
          count_(count) {
      MOZ_ASSERT((uintptr_t(object) & 1) == 0);
      MOZ_ASSERT(kind <= 1);
      MOZ_ASSERT(count > 0);
      MOZ_ASSERT(start + count > start);
    }

    NativeObject* object() const {
      return reinterpret_cast<NativeObject*>(objectAndKind_ & ~1);
    }
    int kind() const { return (int)(objectAndKind_ & 1); }

    bool operator==(const SlotsEdge& other) const {
      return objectAndKind_ == other.objectAndKind_ && start_ == other.start_ &&
             count_ == other.count_;
    }
    bool operator!=(const SlotsEdge& other) const { return !(*this == other); }

    // True if this SlotsEdge range is adjacent to or overlaps with the other
    // SlotsEdge range. The adjacency case will coalesce a series of increasing
    // or decreasing single index writes 0, 1, 2, ..., N into a SlotsEdge range
    // of elements [0, N].
    bool touches(const SlotsEdge& other) const {
      if (objectAndKind_ != other.objectAndKind_) {
        return false;
      }

      if (other.start_ < start_) {
        // If the other range starts before this one, then it touches if its
        // (exclusive) end is at or after this range.
        return other.start_ + other.count_ >= start_;
      }

      // Otherwise, the other range touches if it starts before or at
      // the exclusive end of this range.
      return other.start_ <= start_ + count_;
    }

    // Destructively make this SlotsEdge range the union of the other
    // SlotsEdge range and this one. A precondition is that the ranges must
    // overlap.
    void merge(const SlotsEdge& other) {
      MOZ_ASSERT(touches(other));
      uint32_t end = std::max(start_ + count_, other.start_ + other.count_);
      start_ = std::min(start_, other.start_);
      count_ = end - start_;
      MOZ_ASSERT(count_ > 0);
      MOZ_ASSERT(start_ + count_ > start_);
    }

    bool maybeInRememberedSet(const Nursery& n) const {
      return !IsInsideNursery(reinterpret_cast<Cell*>(object()));
    }

    void trace(TenuringTracer& mover) const;

    explicit operator bool() const { return objectAndKind_ != 0; }

    struct Hasher {
      using Lookup = SlotsEdge;
      static HashNumber hash(const Lookup& l) {
        return mozilla::HashGeneric(l.objectAndKind_, l.start_, l.count_);
      }
      static bool match(const SlotsEdge& k, const Lookup& l) { return k == l; }
    };
  };

  struct WasmAnyRefEdge {
    wasm::AnyRef* edge;

    WasmAnyRefEdge() : edge(nullptr) {}
    explicit WasmAnyRefEdge(wasm::AnyRef* v) : edge(v) {}
    bool operator==(const WasmAnyRefEdge& other) const {
      return edge == other.edge;
    }
    bool operator!=(const WasmAnyRefEdge& other) const {
      return edge != other.edge;
    }

    bool isGCThing() const { return edge->isGCThing(); }

    Cell* deref() const {
      return isGCThing() ? static_cast<Cell*>(edge->toGCThing()) : nullptr;
    }

    bool maybeInRememberedSet(const Nursery& nursery) const {
      MOZ_ASSERT(IsInsideNursery(deref()));
      return !nursery.isInside(edge);
    }

    void trace(TenuringTracer& mover) const;

    explicit operator bool() const { return edge != nullptr; }

    using Hasher = PointerEdgeHasher<WasmAnyRefEdge>;
  };

#ifdef DEBUG
  void checkAccess() const;
#else
  void checkAccess() const {}
#endif

  template <typename Buffer, typename Edge>
  void unputEdge(Buffer& buffer, const Edge& edge) {
    checkAccess();
    if (!isEnabled()) {
      return;
    }

    mozilla::ReentrancyGuard g(*this);

    buffer.unput(edge);
  }

  template <typename Buffer, typename Edge>
  void putEdge(Buffer& buffer, const Edge& edge, JS::GCReason overflowReason) {
    checkAccess();
    if (!isEnabled()) {
      return;
    }

    mozilla::ReentrancyGuard g(*this);

    if (!edge.maybeInRememberedSet(nursery_)) {
      return;
    }

    PutResult r = buffer.put(edge);

    if (MOZ_UNLIKELY(r == PutResult::AboutToOverflow)) {
      setAboutToOverflow(overflowReason);
    }
  }

  template <typename Buffer, typename Edge>
  void putEdgeFromTenured(Buffer& buffer, const Edge& edge,
                          JS::GCReason overflowReason) {
    MOZ_ASSERT(edge.maybeInRememberedSet(nursery_));
    checkAccess();

    if (!isEnabled()) {
      return;
    }

    mozilla::ReentrancyGuard g(*this);
    PutResult r = buffer.put(edge);
    if (MOZ_UNLIKELY(r == PutResult::AboutToOverflow)) {
      setAboutToOverflow(overflowReason);
    }
  }

  MonoTypeBuffer<ValueEdge> bufferVal;
  MonoTypeBuffer<StringPtrEdge> bufStrCell;
  MonoTypeBuffer<BigIntPtrEdge> bufBigIntCell;
  MonoTypeBuffer<GetterSetterPtrEdge> bufGetterSetterCell;
  MonoTypeBuffer<ObjectPtrEdge> bufObjCell;
  MonoTypeBuffer<SlotsEdge> bufferSlot;
  MonoTypeBuffer<WasmAnyRefEdge> bufferWasmAnyRef;
  WholeCellBuffer bufferWholeCell;
  GenericBuffer bufferGeneric;

  GCRuntime* gc_;
  Nursery& nursery_;
  size_t entryCount_;
  double entryScaling_;

  bool aboutToOverflow_;
  bool enabled_;
  bool mayHavePointersToDeadCells_;
#ifdef DEBUG
  bool mEntered; /* For ReentrancyGuard. */
#endif

 public:
  explicit StoreBuffer(GCRuntime* gc);

  StoreBuffer(const StoreBuffer& other) = delete;
  StoreBuffer& operator=(const StoreBuffer& other) = delete;

  StoreBuffer(StoreBuffer&& other);
  StoreBuffer& operator=(StoreBuffer&& other);

  [[nodiscard]] bool enable();
  void disable();
  bool isEnabled() const { return enabled_; }

  void updateSize();

  bool isEmpty() const;
  void clear();

  const Nursery& nursery() const { return nursery_; }

  /* Get the overflowed status. */
  bool isAboutToOverflow() const { return aboutToOverflow_; }

  /*
   * Brain transplants may add whole cell buffer entires for dead cells. We must
   * evict the nursery prior to sweeping arenas if any such entries are present.
   */
  bool mayHavePointersToDeadCells() const {
    return mayHavePointersToDeadCells_;
  }

  /* Insert a single edge into the buffer/remembered set. */
  void putValue(JS::Value* vp) {
    putEdge(bufferVal, ValueEdge(vp), JS::GCReason::FULL_VALUE_BUFFER);
  }
  void unputValue(JS::Value* vp) { unputEdge(bufferVal, ValueEdge(vp)); }

  void putCell(JSString** strp) {
    putEdge(bufStrCell, StringPtrEdge(strp),
            JS::GCReason::FULL_CELL_PTR_STR_BUFFER);
  }
  void unputCell(JSString** strp) {
    unputEdge(bufStrCell, StringPtrEdge(strp));
  }

  void putCell(JS::BigInt** bip) {
    putEdge(bufBigIntCell, BigIntPtrEdge(bip),
            JS::GCReason::FULL_CELL_PTR_BIGINT_BUFFER);
  }
  void unputCell(JS::BigInt** bip) {
    unputEdge(bufBigIntCell, BigIntPtrEdge(bip));
  }

  void putCell(JSObject** strp) {
    putEdge(bufObjCell, ObjectPtrEdge(strp),
            JS::GCReason::FULL_CELL_PTR_OBJ_BUFFER);
  }
  void unputCell(JSObject** strp) {
    unputEdge(bufObjCell, ObjectPtrEdge(strp));
  }

  void putCell(js::GetterSetter** gsp) {
    putEdge(bufGetterSetterCell, GetterSetterPtrEdge(gsp),
            JS::GCReason::FULL_CELL_PTR_GETTER_SETTER_BUFFER);
  }
  void unputCell(js::GetterSetter** gsp) {
    unputEdge(bufGetterSetterCell, GetterSetterPtrEdge(gsp));
  }

  void putSlot(NativeObject* obj, int kind, uint32_t start, uint32_t count) {
    SlotsEdge edge(obj, kind, start, count);
    if (bufferSlot.last_.touches(edge)) {
      bufferSlot.last_.merge(edge);
    } else {
      putEdge(bufferSlot, edge, JS::GCReason::FULL_SLOT_BUFFER);
    }
  }

  void putWasmAnyRef(wasm::AnyRef* vp) {
    putEdge(bufferWasmAnyRef, WasmAnyRefEdge(vp),
            JS::GCReason::FULL_WASM_ANYREF_BUFFER);
  }
  void putWasmAnyRefEdgeFromTenured(wasm::AnyRef* vp) {
    putEdgeFromTenured(bufferWasmAnyRef, WasmAnyRefEdge(vp),
                       JS::GCReason::FULL_WASM_ANYREF_BUFFER);
  }
  void unputWasmAnyRef(wasm::AnyRef* vp) {
    unputEdge(bufferWasmAnyRef, WasmAnyRefEdge(vp));
  }

  static inline bool isInWholeCellBuffer(Cell* cell);
  inline void putWholeCell(Cell* cell);
  inline void putWholeCellDontCheckLast(Cell* cell);
  const void* addressOfLastBufferedWholeCell() {
    return bufferWholeCell.lastBufferedPtr();
  }

  /* Insert an entry into the generic buffer. */
  template <typename T>
  void putGeneric(const T& t) {
    putEdge(bufferGeneric, t, JS::GCReason::FULL_GENERIC_BUFFER);
  }

  void setMayHavePointersToDeadCells() { mayHavePointersToDeadCells_ = true; }

  /* Methods to trace the source of all edges in the store buffer. */
  void traceValues(TenuringTracer& mover);
  void traceCells(TenuringTracer& mover);
  void traceSlots(TenuringTracer& mover);
  void traceWasmAnyRefs(TenuringTracer& mover);
  void traceWholeCells(TenuringTracer& mover);
  void traceGenericEntries(JSTracer* trc);

  /* For use by our owned buffers and for testing. */
  void setAboutToOverflow(JS::GCReason);

  void addSizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf,
                              JS::GCSizes* sizes);

  void checkEmpty() const;
};

// A set of cells in an arena used to implement the whole cell store buffer.
// Also used to store a set of cells that need to be swept.
class ArenaCellSet {
  friend class StoreBuffer;

  // Use a 32 bit word to make it easier to access ArenaCellBits from JIT code.
  using ArenaCellBits = BitArray<MaxArenaCellIndex, uint32_t>;

  // The arena this relates to.
  Arena* arena = nullptr;

  // Pointer to next set forming a linked list. Used to form the list of cell
  // sets to sweep.
  ArenaCellSet* next = nullptr;

  // Bit vector for each possible cell start position.
  ArenaCellBits bits;

#ifdef DEBUG
  // The minor GC number when this was created. This object should not survive
  // past the next minor collection.
  const uint64_t minorGCNumberAtCreation = 0;
#endif

  // Construct the empty sentinel object.
  constexpr ArenaCellSet() = default;

 public:
  using WordT = ArenaCellBits::WordT;
  static constexpr size_t BitsPerWord = ArenaCellBits::bitsPerElement;

  explicit ArenaCellSet(Arena* arena);

  bool hasCell(const TenuredCell* cell) const {
    return hasCell(getCellIndex(cell));
  }

  void putCell(const TenuredCell* cell) { putCell(getCellIndex(cell)); }

  bool isEmpty() const { return this == &Empty; }

  bool hasCell(size_t cellIndex) const;

  void putCell(size_t cellIndex);

  void check() const;

  WordT getWord(size_t wordIndex) const { return bits.getWord(wordIndex); }

  // Sweep this set, returning whether it also needs to be swept later.
  void trace(TenuringTracer& mover);

  // Sentinel object used for all empty sets.
  //
  // We use a sentinel because it simplifies the JIT code slightly as we can
  // assume all arenas have a cell set.
  static ArenaCellSet Empty;

  static size_t getCellIndex(const TenuredCell* cell);
  static std::pair<size_t, uint32_t> getWordIndexAndMask(size_t cellIndex);

  // Attempt to trigger a minor GC if free space in the nursery (where these
  // objects are allocated) falls below this threshold.
  static const size_t NurseryFreeThresholdBytes = 64 * 1024;

  static size_t offsetOfArena() { return offsetof(ArenaCellSet, arena); }
  static size_t offsetOfBits() {
    return offsetof(ArenaCellSet, bits) + ArenaCellBits::offsetOfMap();
  }
};

// Post-write barrier implementation for GC cells.

// Implement the post-write barrier for nursery allocateable cell type |T|. Call
// this from |T::postWriteBarrier|.
template <typename T>
MOZ_ALWAYS_INLINE void PostWriteBarrierImpl(void* cellp, T* prev, T* next) {
  MOZ_ASSERT(cellp);

  // If the target needs an entry, add it.
  StoreBuffer* buffer;
  if (next && (buffer = next->storeBuffer())) {
    // If we know that the prev has already inserted an entry, we can skip
    // doing the lookup to add the new entry. Note that we cannot safely
    // assert the presence of the entry because it may have been added
    // via a different store buffer.
    if (prev && prev->storeBuffer()) {
      return;
    }
    buffer->putCell(static_cast<T**>(cellp));
    return;
  }

  // Remove the prev entry if the new value does not need it. There will only
  // be a prev entry if the prev value was in the nursery.
  if (prev && (buffer = prev->storeBuffer())) {
    buffer->unputCell(static_cast<T**>(cellp));
  }
}

template <typename T>
MOZ_ALWAYS_INLINE void PostWriteBarrier(T** vp, T* prev, T* next) {
  static_assert(std::is_base_of_v<Cell, T>);
  static_assert(!std::is_same_v<Cell, T> && !std::is_same_v<TenuredCell, T>);

  if constexpr (!GCTypeIsTenured<T>()) {
    using BaseT = typename BaseGCType<T>::type;
    PostWriteBarrierImpl<BaseT>(vp, prev, next);
    return;
  }

  MOZ_ASSERT_IF(next, !IsInsideNursery(next));
}

// Used when we don't have a specific edge to put in the store buffer.
void PostWriteBarrierCell(Cell* cell, Cell* prev, Cell* next);

}  // namespace gc
}  // namespace js

#endif /* gc_StoreBuffer_h */