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

#include "gc/Barrier.h"
#include "gc/Tracer.h"
#include "js/GCHashTable.h"
#include "js/GCPolicyAPI.h"
#include "js/GCVector.h"
#include "js/Utility.h"

namespace js {

namespace detail {

// This class only handles the incremental case and does not deal with nursery
// pointers. The only users should be for NurseryAwareHashMap; it is defined
// externally because we need a GCPolicy for its use in the contained map.

DEFINE_BARRIERED_PTR(UnsafeBareWeakHeapPtr, gc::BarrierOption_ReadBarrier);

}  // namespace detail

enum : bool { DuplicatesNotPossible, DuplicatesPossible };

// The "nursery aware" hash map is a special case of GCHashMap that is able to
// treat nursery allocated members weakly during a minor GC: e.g. it allows for
// nursery allocated objects to be collected during nursery GC where a normal
// hash table treats such edges strongly.
//
// Doing this requires some strong constraints on what can be stored in this
// table and how it can be accessed. At the moment, this table assumes that all
// values contain a strong reference to the key. This limits its usefulness to
// the CrossCompartmentMap at the moment, but might serve as a useful base for
// other tables in future.
template <typename Key, typename Value, typename AllocPolicy = TempAllocPolicy,
          bool AllowDuplicates = DuplicatesNotPossible>
class NurseryAwareHashMap {
  using MapKey = UnsafeBarePtr<Key>;
  using MapValue = detail::UnsafeBareWeakHeapPtr<Value>;
  using HashPolicy = DefaultHasher<MapKey>;
  using MapType = GCRekeyableHashMap<MapKey, MapValue, HashPolicy, AllocPolicy>;
  MapType map;

  // Keep a list of all keys for which key->isTenured() is false. This lets us
  // avoid a full traversal of the map on each minor GC, keeping the minor GC
  // times proportional to the nursery heap size.
  using KeyVector = GCVector<Key, 0, AllocPolicy>;
  KeyVector nurseryEntries;

 public:
  using Lookup = typename MapType::Lookup;
  using Ptr = typename MapType::Ptr;
  using Iterator = typename MapType::Iterator;
  using ModIterator = typename MapType::ModIterator;
  using Entry = typename MapType::Entry;

  explicit NurseryAwareHashMap(AllocPolicy a = AllocPolicy())
      : map(a), nurseryEntries(std::move(a)) {}
  explicit NurseryAwareHashMap(size_t length) : map(length) {}
  NurseryAwareHashMap(AllocPolicy a, size_t length)
      : map(a, length), nurseryEntries(std::move(a)) {}

  bool empty() const { return map.empty(); }
  Ptr lookup(const Lookup& l) const { return map.lookup(l); }
  void remove(Ptr p) { map.remove(p); }
  Iterator iter() const { return map.iter(); }
  ModIterator modIter() { return map.modIter(); }
  struct Enum : public MapType::Enum {
    explicit Enum(NurseryAwareHashMap& namap) : MapType::Enum(namap.map) {}
  };
  size_t sizeOfExcludingThis(mozilla::MallocSizeOf mallocSizeOf) const {
    return map.shallowSizeOfExcludingThis(mallocSizeOf) +
           nurseryEntries.sizeOfExcludingThis(mallocSizeOf);
  }
  size_t sizeOfIncludingThis(mozilla::MallocSizeOf mallocSizeOf) const {
    return mallocSizeOf(this) + sizeOfExcludingThis(mallocSizeOf);
  }

  [[nodiscard]] bool put(const Key& key, const Value& value) {
    if ((!key->isTenured() || !value->isTenured()) &&
        !nurseryEntries.append(key)) {
      return false;
    }

    auto p = map.lookupForAdd(key);
    if (p) {
      p->value() = value;
      return true;
    }

    return map.add(p, key, value);
  }

  void sweepAfterMinorGC(JSTracer* trc) {
    nurseryEntries.mutableEraseIf([this, trc](Key& key) {
      auto p = map.lookup(key);
      if (!p) {
        return true;
      }

      // Drop the entry if the value is not marked.
      if (!JS::GCPolicy<MapValue>::traceWeak(trc, &p->value())) {
        map.remove(p);
        return true;
      }

      // Update and relocate the key, if the value is still needed.
      //
      // Non-string Values will contain a strong reference to Key, as per its
      // use in the CrossCompartmentWrapperMap, so the key will never be dying
      // here. Strings do *not* have any sort of pointer from wrapper to
      // wrappee, as they are just copies. The wrapper map entry is merely used
      // as a cache to avoid re-copying the string, and currently that entire
      // cache is flushed on major GC.
      //
      // Since |key| is a reference, this updates the content of the
      // nurseryEntries vector.
      Key prior = key;
      if (!TraceManuallyBarrieredWeakEdge(trc, &key,
                                          "NurseryAwareHashMap key")) {
        map.remove(p);
        return true;
      }

      bool valueIsTenured = p->value().unbarrieredGet()->isTenured();

      if constexpr (AllowDuplicates) {
        // Drop duplicated keys.
        //
        // A key can be forwarded to another place. In this case, rekey the
        // item. If two or more different keys are forwarded to the same new
        // key, simply drop the later ones.
        if (key == prior) {
          // No rekey needed.
        } else if (map.has(key)) {
          // Key was forwarded to the same place that another key was already
          // forwarded to.
          map.remove(p);
          return true;
        } else {
          map.rekeyAs(prior, key, key);
        }
      } else {
        MOZ_ASSERT(key == prior || !map.has(key));
        map.rekeyIfMoved(prior, key);
      }

      return key->isTenured() && valueIsTenured;
    });

    checkNurseryEntries();
  }

  void checkNurseryEntries() const {
#ifdef DEBUG
    AutoEnterOOMUnsafeRegion oomUnsafe;
    HashSet<Key, DefaultHasher<Key>, SystemAllocPolicy> set;
    for (const auto& key : nurseryEntries) {
      if (!set.put(key)) {
        oomUnsafe.crash("NurseryAwareHashMap::checkNurseryEntries");
      }
    }

    for (auto i = map.iter(); !i.done(); i.next()) {
      Key key = i.get().key().get();
      MOZ_ASSERT(gc::IsCellPointerValid(key));
      MOZ_ASSERT_IF(IsInsideNursery(key), set.has(key));

      Value value = i.get().value().unbarrieredGet();
      MOZ_ASSERT(gc::IsCellPointerValid(value));
      MOZ_ASSERT_IF(IsInsideNursery(value), set.has(key));
    }
#endif
  }

  void traceWeak(JSTracer* trc) { map.traceWeak(trc); }

  void clear() {
    map.clear();
    nurseryEntries.clear();
  }

  bool hasNurseryEntries() const { return !nurseryEntries.empty(); }
};

}  // namespace js

namespace JS {

template <typename T>
struct GCPolicy<js::detail::UnsafeBareWeakHeapPtr<T>>
    : public GCPolicyBase<js::detail::UnsafeBareWeakHeapPtr<T>> {
  static void trace(JSTracer* trc, js::detail::UnsafeBareWeakHeapPtr<T>* thingp,
                    const char* name) {
    js::TraceEdge(trc, thingp, name);
  }
  static bool traceWeak(JSTracer* trc,
                        js::detail::UnsafeBareWeakHeapPtr<T>* thingp) {
    return js::TraceWeakEdge(trc, thingp, "UnsafeBareWeakHeapPtr");
  }
};

}  // namespace JS

namespace mozilla {}  // namespace mozilla

#endif  // gc_NurseryAwareHashMap_h