/*
 * Copyright 2014 Google, Inc
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#ifndef SkSmallAllocator_DEFINED
#define SkSmallAllocator_DEFINED

#include "SkTDArray.h"
#include "SkTypes.h"

#include <new>

/*
 *  Template class for allocating small objects without additional heap memory
 *  allocations. kMaxObjects is a hard limit on the number of objects that can
 *  be allocated using this class. After that, attempts to create more objects
 *  with this class will assert and return nullptr.
 *  kTotalBytes is the total number of bytes provided for storage for all
 *  objects created by this allocator. If an object to be created is larger
 *  than the storage (minus storage already used), it will be allocated on the
 *  heap. This class's destructor will handle calling the destructor for each
 *  object it allocated and freeing its memory.
 */
template<uint32_t kMaxObjects, size_t kTotalBytes>
class SkSmallAllocator : SkNoncopyable {
public:
    SkSmallAllocator()
    : fStorageUsed(0)
    , fNumObjects(0)
    {}

    ~SkSmallAllocator() {
        // Destruct in reverse order, in case an earlier object points to a
        // later object.
        while (fNumObjects > 0) {
            fNumObjects--;
            Rec* rec = &fRecs[fNumObjects];
            rec->fKillProc(rec->fObj);
            // Safe to do if fObj is in fStorage, since fHeapStorage will
            // point to nullptr.
            sk_free(rec->fHeapStorage);
        }
    }

    /*
     *  Create a new object of type T. Its lifetime will be handled by this
     *  SkSmallAllocator.
     *  Note: If kMaxObjects have been created by this SkSmallAllocator, nullptr
     *  will be returned.
     */
    template<typename T, typename... Args>
    T* createT(const Args&... args) {
        void* buf = this->reserveT<T>();
        if (nullptr == buf) {
            return nullptr;
        }
        return new (buf) T(args...);
    }

    /*
     *  Reserve a specified amount of space (must be enough space for one T).
     *  The space will be in fStorage if there is room, or on the heap otherwise.
     *  Either way, this class will call ~T() in its destructor and free the heap
     *  allocation if necessary.
     *  Unlike createT(), this method will not call the constructor of T.
     */
    template<typename T> void* reserveT(size_t storageRequired = sizeof(T)) {
        SkASSERT(fNumObjects < kMaxObjects);
        SkASSERT(storageRequired >= sizeof(T));
        if (kMaxObjects == fNumObjects) {
            return nullptr;
        }
        const size_t storageRemaining = SkAlign4(kTotalBytes) - fStorageUsed;
        storageRequired = SkAlign4(storageRequired);
        Rec* rec = &fRecs[fNumObjects];
        if (storageRequired > storageRemaining) {
            // Allocate on the heap. Ideally we want to avoid this situation,
            // but we're not sure we can catch all callers, so handle it but
            // assert false in debug mode.
            SkASSERT(false);
            rec->fStorageSize = 0;
            rec->fHeapStorage = sk_malloc_throw(storageRequired);
            rec->fObj = static_cast<void*>(rec->fHeapStorage);
        } else {
            // There is space in fStorage.
            rec->fStorageSize = storageRequired;
            rec->fHeapStorage = nullptr;
            SkASSERT(SkIsAlign4(fStorageUsed));
            rec->fObj = static_cast<void*>(fStorage + (fStorageUsed / 4));
            fStorageUsed += storageRequired;
        }
        rec->fKillProc = DestroyT<T>;
        fNumObjects++;
        return rec->fObj;
    }

    /*
     *  Free the memory reserved last without calling the destructor.
     *  Can be used in a nested way, i.e. after reserving A and B, calling
     *  freeLast once will free B and calling it again will free A.
     */
    void freeLast() {
        SkASSERT(fNumObjects > 0);
        Rec* rec = &fRecs[fNumObjects - 1];
        sk_free(rec->fHeapStorage);
        fStorageUsed -= rec->fStorageSize;

        fNumObjects--;
    }

private:
    struct Rec {
        size_t fStorageSize;  // 0 if allocated on heap
        void*  fObj;
        void*  fHeapStorage;
        void   (*fKillProc)(void*);
    };

    // Used to call the destructor for allocated objects.
    template<typename T>
    static void DestroyT(void* ptr) {
        static_cast<T*>(ptr)->~T();
    }

    // Number of bytes used so far.
    size_t              fStorageUsed;
    // Pad the storage size to be 4-byte aligned.
    uint32_t            fStorage[SkAlign4(kTotalBytes) >> 2];
    uint32_t            fNumObjects;
    Rec                 fRecs[kMaxObjects];
};

#endif // SkSmallAllocator_DEFINED