#pragma once
/**
 * This file encapsulates the in disused object free lists
 * that are used per slab of small objects. The implementation
 * can be configured to introduce randomness to the reallocation,
 * and also provide signing to detect free list corruption.
 *
 * # Corruption
 *
 * The corruption detection works as follows
 *
 *   FreeObject
 *   -----------------------------
 *   | next | prev_encoded | ... |
 *   -----------------------------
 * A free object contains a pointer to next object in the free list, and
 * a prev pointer, but the prev pointer is really a signature with the
 * following property
 *
 *  If n = c->next && n != 0, then n->prev_encoded = f(c,n).
 *
 * If f just returns the first parameter, then this degenerates to a doubly
 * linked list.  (Note that doing the degenerate case can be useful for
 * debugging snmalloc bugs.) By making it a function of both pointers, it
 * makes it harder for an adversary to mutate prev_encoded to a valid value.
 *
 * This provides protection against the free-list being corrupted by memory
 * safety issues.
 *
 * # Randomness
 *
 * The randomness is introduced by building two free lists simulatenously,
 * and randomly deciding which list to add an element to.
 */

#include "../ds/address.h"
#include "allocconfig.h"
#include "entropy.h"

#include <cstdint>

namespace snmalloc
{
  /**
   * This function is used to sign back pointers in the free list.
   */
  inline static uintptr_t
  signed_prev(address_t curr, address_t next, const FreeListKey& key)
  {
    auto c = curr;
    auto n = next;
    return (c + key.key1) * (n + key.key2);
  }

  class FreeObject
  {
  public:
    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    class T;

    /**
     * This "inductive step" type -- a queue-annotated pointer to a FreeObject
     * containing a queue-annotated pointer -- shows up all over the place.
     * Give it a shorter name (FreeObject::QueuePtr<BQueue>) for convenience.
     */
    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    using QueuePtr = CapPtr<FreeObject::T<BQueue>, BQueue>;

    /**
     * As with QueuePtr, but atomic.
     */
    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    using AtomicQueuePtr = AtomicCapPtr<FreeObject::T<BQueue>, BQueue>;

    /**
     * This is the "base case" of that induction.  While we can't get rid of the
     * two different type parameters (in general), we can at least get rid of a
     * bit of the clutter.  "FreeObject::HeadPtr<BView, BQueue>" looks a little
     * nicer than "CapPtr<FreeObject::T<BQueue>, BView>".
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    using HeadPtr = CapPtr<FreeObject::T<BQueue>, BView>;

    /**
     * As with HeadPtr, but atomic.
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    using AtomicHeadPtr = AtomicCapPtr<FreeObject::T<BQueue>, BView>;

    /**
     * Free objects within each slab point directly to the next.
     * There is an optional second field that is effectively a
     * back pointer in a doubly linked list, however, it is encoded
     * to prevent corruption.
     *
     * This is an inner class to avoid the need to specify BQueue when calling
     * static methods.
     *
     * Raw C++ pointers to this type are *assumed to be domesticated*.  In some
     * cases we still explicitly annotate domesticated FreeObject*-s as
     * CapPtr<>, but more often CapPtr<FreeObject::T<A>,B> will have B = A.
     *
     * TODO: Consider putting prev_encoded at the end of the object, would
     * require size to be threaded through, but would provide more OOB
     * detection.
     */
    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    class T
    {
      friend class FreeObject;

      union
      {
        QueuePtr<BQueue> next_object;
        // TODO: Should really use C++20 atomic_ref rather than a union.
        AtomicQueuePtr<BQueue> atomic_next_object;
      };
#ifdef SNMALLOC_CHECK_CLIENT
      // Encoded representation of a back pointer.
      // Hard to fake, and provides consistency on
      // the next pointers.
      address_t prev_encoded;
#endif

    public:
      template<
        SNMALLOC_CONCEPT(capptr::ConceptBound) BView = typename BQueue::
          template with_wildness<capptr::dimension::Wildness::Tame>,
        typename Domesticator>
      HeadPtr<BView, BQueue>
      atomic_read_next(const FreeListKey& key, Domesticator domesticate)
      {
        auto n_wild = FreeObject::decode_next(
          address_cast(&this->next_object),
          this->atomic_next_object.load(std::memory_order_acquire),
          key);
        auto n_tame = domesticate(n_wild);
#ifdef SNMALLOC_CHECK_CLIENT
        if (n_tame != nullptr)
        {
          n_tame->check_prev(
            signed_prev(address_cast(this), address_cast(n_tame), key));
        }
#endif
        return n_tame;
      }

      /**
       * Read the next pointer
       */
      template<
        SNMALLOC_CONCEPT(capptr::ConceptBound) BView = typename BQueue::
          template with_wildness<capptr::dimension::Wildness::Tame>,
        typename Domesticator>
      HeadPtr<BView, BQueue>
      read_next(const FreeListKey& key, Domesticator domesticate)
      {
        return domesticate(FreeObject::decode_next(
          address_cast(&this->next_object), this->next_object, key));
      }

      /**
       * Check the signature of this FreeObject
       */
      void check_prev(address_t signed_prev)
      {
        UNUSED(signed_prev);
        check_client(
          signed_prev == this->prev_encoded,
          "Heap corruption - free list corrupted!");
      }
    };

    // Note the inverted template argument order, since BView is inferable.
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView>
    static HeadPtr<BView, BQueue> make(CapPtr<void, BView> p)
    {
      return p.template as_static<FreeObject::T<BQueue>>();
    }

  private:
    /**
     * Involutive encryption with raw pointers
     */
    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    inline static FreeObject::T<BQueue>* code_next(
      address_t curr, FreeObject::T<BQueue>* next, const FreeListKey& key)
    {
      // Note we can consider other encoding schemes here.
      //   * XORing curr and next.  This doesn't require any key material
      //   * XORing (curr * key). This makes it harder to guess the underlying
      //     key, as each location effectively has its own key.
      // Curr is not used in the current encoding scheme.
      UNUSED(curr);

      if constexpr (CHECK_CLIENT && !aal_supports<StrictProvenance>)
      {
        return reinterpret_cast<FreeObject::T<BQueue>*>(
          reinterpret_cast<uintptr_t>(next) ^ key.key_next);
      }
      else
      {
        UNUSED(key);
        return next;
      }
    }

  public:
    /**
     * Encode next.  We perform two convenient little bits of type-level
     * sleight of hand here:
     *
     *  1) We convert the provided HeadPtr to a QueuePtr, forgetting BView in
     *  the result; all the callers write the result through a pointer to a
     *  QueuePtr, though, strictly, the result itself is no less domesticated
     *  than the input (even if it is obfuscated).
     *
     *  2) Speaking of obfuscation, we continue to use a CapPtr<> type even
     *  though the result is likely not safe to dereference, being an
     *  obfuscated bundle of bits (on non-CHERI architectures, anyway).  That's
     *  additional motivation to consider the result BQueue-bounded, as that is
     *  likely (but not necessarily) Wild.
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    inline static QueuePtr<BQueue> encode_next(
      address_t curr, HeadPtr<BView, BQueue> next, const FreeListKey& key)
    {
      return QueuePtr<BQueue>(code_next(curr, next.unsafe_ptr(), key));
    }

    /**
     * Decode next.  While traversing a queue, BView and BQueue here will often
     * be equal (i.e., CBAllocExportWild) rather than dichotomous.  However,
     * we do occasionally decode an actual head pointer, so be polymorphic here.
     *
     * TODO: We'd like, in some sense, to more tightly couple or integrate this
     * into to the domestication process.  We could introduce an additional
     * state in the capptr_bounds::wild taxonomy (e.g, Obfuscated) so that the
     * Domesticator-s below have to call through this function to get the Wild
     * pointer they can then make Tame.  It's not yet entirely clear what that
     * would look like and whether/how the encode_next side of things should be
     * exposed.  For the moment, obfuscation is left encapsulated within
     * FreeObject and we do not capture any of it statically.
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    inline static HeadPtr<BView, BQueue> decode_next(
      address_t curr, HeadPtr<BView, BQueue> next, const FreeListKey& key)
    {
      return HeadPtr<BView, BQueue>(code_next(curr, next.unsafe_ptr(), key));
    }

    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    static void assert_view_queue_bounds()
    {
      static_assert(
        BView::wildness == capptr::dimension::Wildness::Tame,
        "FreeObject View must be domesticated, justifying raw pointers");

      static_assert(
        std::is_same_v<
          typename BQueue::template with_wildness<
            capptr::dimension::Wildness::Tame>,
          BView>,
        "FreeObject Queue bounds must match View bounds (but may be Wild)");
    }

    /**
     * Assign next_object and update its prev_encoded if
     * SNMALLOC_CHECK_CLIENT. Static so that it can be used on reference to a
     * FreeObject.
     *
     * Returns a pointer to the next_object field of the next parameter as an
     * optimization for repeated snoc operations (in which
     * next->next_object is nullptr).
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    static QueuePtr<BQueue>* store_next(
      QueuePtr<BQueue>* curr,
      HeadPtr<BView, BQueue> next,
      const FreeListKey& key)
    {
      assert_view_queue_bounds<BView, BQueue>();

#ifdef SNMALLOC_CHECK_CLIENT
      next->prev_encoded =
        signed_prev(address_cast(curr), address_cast(next), key);
#else
      UNUSED(key);
#endif
      *curr = encode_next(address_cast(curr), next, key);
      return &(next->next_object);
    }

    template<SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    static void store_null(QueuePtr<BQueue>* curr, const FreeListKey& key)
    {
      *curr = encode_next(address_cast(curr), QueuePtr<BQueue>(nullptr), key);
    }

    /**
     * Assign next_object and update its prev_encoded if SNMALLOC_CHECK_CLIENT
     *
     * Uses the atomic view of next, so can be used in the message queues.
     */
    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    static void atomic_store_next(
      HeadPtr<BView, BQueue> curr,
      HeadPtr<BView, BQueue> next,
      const FreeListKey& key)
    {
      static_assert(BView::wildness == capptr::dimension::Wildness::Tame);

#ifdef SNMALLOC_CHECK_CLIENT
      next->prev_encoded =
        signed_prev(address_cast(curr), address_cast(next), key);
#else
      UNUSED(key);
#endif
      // Signature needs to be visible before item is linked in
      // so requires release semantics.
      curr->atomic_next_object.store(
        encode_next(address_cast(&curr->next_object), next, key),
        std::memory_order_release);
    }

    template<
      SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
      SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
    static void
    atomic_store_null(HeadPtr<BView, BQueue> curr, const FreeListKey& key)
    {
      static_assert(BView::wildness == capptr::dimension::Wildness::Tame);

      curr->atomic_next_object.store(
        encode_next(
          address_cast(&curr->next_object), QueuePtr<BQueue>(nullptr), key),
        std::memory_order_relaxed);
    }
  };

  static_assert(
    sizeof(FreeObject) <= MIN_ALLOC_SIZE,
    "Needs to be able to fit in smallest allocation.");

  /**
   * Used to iterate a free list in object space.
   *
   * Checks signing of pointers
   */
  template<
    SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
    SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue>
  class FreeListIter
  {
    FreeObject::HeadPtr<BView, BQueue> curr{nullptr};
#ifdef SNMALLOC_CHECK_CLIENT
    address_t prev{0};
#endif

  public:
    constexpr FreeListIter(
      FreeObject::HeadPtr<BView, BQueue> head, address_t prev_value)
    : curr(head)
    {
#ifdef SNMALLOC_CHECK_CLIENT
      prev = prev_value;
#endif
      UNUSED(prev_value);
    }

    constexpr FreeListIter() = default;

    /**
     * Checks if there are any more values to iterate.
     */
    bool empty()
    {
      return curr == nullptr;
    }

    /**
     * Returns current head without affecting the iterator.
     */
    FreeObject::HeadPtr<BView, BQueue> peek()
    {
      return curr;
    }

    /**
     * Moves the iterator on, and returns the current value.
     */
    FreeObject::HeadPtr<BView, BQueue> take(const FreeListKey& key)
    {
      auto c = curr;
      // TODO: Placeholder
      auto domesticate = [](FreeObject::QueuePtr<BQueue> p)
                           SNMALLOC_FAST_PATH_LAMBDA { return p; };
      auto next = curr->read_next(key, domesticate);

      Aal::prefetch(next.unsafe_ptr());
      curr = next;
#ifdef SNMALLOC_CHECK_CLIENT
      c->check_prev(prev);
      prev = signed_prev(address_cast(c), address_cast(next), key);
#else
      UNUSED(key);
#endif

      return c;
    }
  };

  /**
   * Used to build a free list in object space.
   *
   * Adds signing of pointers in the SNMALLOC_CHECK_CLIENT mode
   *
   * If RANDOM is enabled, the builder uses two queues, and
   * "randomly" decides to add to one of the two queues.  This
   * means that we will maintain a randomisation of the order
   * between allocations.
   *
   * The fields are paired up to give better codegen as then they are offset
   * by a power of 2, and the bit extract from the interleaving seed can
   * be shifted to calculate the relevant offset to index the fields.
   *
   * If RANDOM is set to false, then the code does not perform any
   * randomisation.
   */
  template<
    bool RANDOM,
    SNMALLOC_CONCEPT(capptr::ConceptBound) BView,
    SNMALLOC_CONCEPT(capptr::ConceptBound) BQueue,
    bool INIT = true>
  class FreeListBuilder
  {
    static constexpr size_t LENGTH = RANDOM ? 2 : 1;

    /*
     * We use native pointers below so that we don't run afoul of strict
     * aliasing rules.  head is a FreeObject::HeadPtr<BView, BQueue> -- that
     * is, a known-domesticated pointer to a queue of wild pointers -- and
     * it's usually the case that end is a FreeObject::QueuePtr<BQueue>* --
     * that is, a known-domesticated pointer to a wild pointer to a queue of
     * wild pointers.  However, in order to do branchless inserts, we set end
     * = &head, which breaks strict aliasing rules with the types as given.
     * Fortunately, these are private members and so we can use native
     * pointers and just expose a more strongly typed interface.
     */

    // Pointer to the first element.
    std::array<void*, LENGTH> head{nullptr};
    // Pointer to the reference to the last element.
    // In the empty case end[i] == &head[i]
    // This enables branch free enqueuing.
    std::array<void**, LENGTH> end{nullptr};

    FreeObject::QueuePtr<BQueue>* cast_end(uint32_t ix)
    {
      return reinterpret_cast<FreeObject::QueuePtr<BQueue>*>(end[ix]);
    }

    void set_end(uint32_t ix, FreeObject::QueuePtr<BQueue>* p)
    {
      end[ix] = reinterpret_cast<void**>(p);
    }

    FreeObject::HeadPtr<BView, BQueue> cast_head(uint32_t ix)
    {
      return FreeObject::HeadPtr<BView, BQueue>(
        static_cast<FreeObject::T<BQueue>*>(head[ix]));
    }

    std::array<uint16_t, RANDOM ? 2 : 0> length{};

  public:
    constexpr FreeListBuilder()
    {
      if (INIT)
      {
        init();
      }
    }

    /**
     * Checks if the builder contains any elements.
     */
    bool empty()
    {
      for (size_t i = 0; i < LENGTH; i++)
      {
        if (end[i] != &head[i])
        {
          return false;
        }
      }
      return true;
    }

    /**
     * Adds an element to the builder
     */
    void add(
      FreeObject::HeadPtr<BView, BQueue> n,
      const FreeListKey& key,
      LocalEntropy& entropy)
    {
      uint32_t index;
      if constexpr (RANDOM)
        index = entropy.next_bit();
      else
        index = 0;

      set_end(index, FreeObject::store_next(cast_end(index), n, key));
      if constexpr (RANDOM)
      {
        length[index]++;
      }
    }

    /**
     * Adds an element to the builder, if we are guaranteed that
     * RANDOM is false.  This is useful in certain construction
     * cases that do not need to introduce randomness, such as
     * during the initialisation construction of a free list, which
     * uses its own algorithm, or during building remote deallocation
     * lists, which will be randomised at the other end.
     */
    template<bool RANDOM_ = RANDOM>
    std::enable_if_t<!RANDOM_>
    add(FreeObject::HeadPtr<BView, BQueue> n, const FreeListKey& key)
    {
      static_assert(RANDOM_ == RANDOM, "Don't set template parameter");
      set_end(0, FreeObject::store_next(cast_end(0), n, key));
    }

    /**
     * Makes a terminator to a free list.
     */
    SNMALLOC_FAST_PATH void
    terminate_list(uint32_t index, const FreeListKey& key)
    {
      FreeObject::store_null(cast_end(index), key);
    }

    /**
     * Read head removing potential encoding
     *
     * Although, head does not require meta-data protection
     * as it is not stored in an object allocation. For uniformity
     * it is treated like the next_object field in a FreeObject
     * and is thus subject to encoding if the next_object pointers
     * encoded.
     */
    FreeObject::HeadPtr<BView, BQueue>
    read_head(uint32_t index, const FreeListKey& key)
    {
      return FreeObject::decode_next(
        address_cast(&head[index]), cast_head(index), key);
    }

    address_t get_fake_signed_prev(uint32_t index, const FreeListKey& key)
    {
      return signed_prev(
        address_cast(&head[index]), address_cast(read_head(index, key)), key);
    }

    /**
     * Close a free list, and set the iterator parameter
     * to iterate it.
     *
     * In the RANDOM case, it may return only part of the freelist.
     *
     * The return value is how many entries are still contained in the builder.
     */
    SNMALLOC_FAST_PATH uint16_t
    close(FreeListIter<BView, BQueue>& fl, const FreeListKey& key)
    {
      uint32_t i;
      if constexpr (RANDOM)
      {
        SNMALLOC_ASSERT(end[1] != &head[0]);
        SNMALLOC_ASSERT(end[0] != &head[1]);

        // Select longest list.
        i = length[0] > length[1] ? 0 : 1;
      }
      else
      {
        i = 0;
      }

      terminate_list(i, key);

      fl = {read_head(i, key), get_fake_signed_prev(i, key)};

      end[i] = &head[i];

      if constexpr (RANDOM)
      {
        length[i] = 0;
        return length[1 - i];
      }
      else
      {
        return 0;
      }
    }

    /**
     * Set the builder to a not building state.
     */
    constexpr void init()
    {
      for (size_t i = 0; i < LENGTH; i++)
      {
        end[i] = &head[i];
        if (RANDOM)
        {
          length[i] = 0;
        }
      }
    }

    template<bool RANDOM_ = RANDOM>
    std::enable_if_t<
      !RANDOM_,
      std::pair<
        FreeObject::HeadPtr<BView, BQueue>,
        FreeObject::HeadPtr<BView, BQueue>>>
    extract_segment(const FreeListKey& key)
    {
      static_assert(RANDOM_ == RANDOM, "Don't set SFINAE parameter!");
      SNMALLOC_ASSERT(!empty());

      auto first = read_head(0, key);
      // end[0] is pointing to the first field in the object,
      // this is doing a CONTAINING_RECORD like cast to get back
      // to the actual object.  This isn't true if the builder is
      // empty, but you are not allowed to call this in the empty case.
      auto last = FreeObject::HeadPtr<BView, BQueue>(
        reinterpret_cast<FreeObject::T<BQueue>*>(end[0]));
      init();
      return {first, last};
    }
  };
} // namespace snmalloc