/* vim: set ts=4 sw=4 tw=99 et:
 *
 * Copyright (C) 2012 David Anderson
 *
 * This file is part of SourcePawn.
 *
 * SourcePawn is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or (at your option)
 * any later version.
 * 
 * SourcePawn is distributed in the hope that it will be useful, but WITHOUT ANY
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with
 * SourcePawn. If not, see http://www.gnu.org/licenses/.
 */
#include "Heap.h"
#include "Heap-inl.h"
#include "Spaces-inl.h"
#include "VirtualObjectVisitor.h"
#include "ObjectVisitor.h"
#include "Zone.h"
#include "StringPool.h"

using namespace ke;

static const size_t kMarkStackSize = 256 * kKB;

Heap::Heap(Zone *zone)
  : zone_(zone),
    mapSpace_(zone),
    oldSpace_(zone),
    markStack_(NULL)
{
}

bool
Heap::initialize()
{
    if (!lifoSpace_.initialize())
        return false;
    if (!oldSpace_.initialize())
        return false;
    if (!mapSpace_.initialize())
        return false;
    if (!largeSpace_.initialize())
        return false;

    markStack_ = malloc(kMarkStackSize);
    if (!markStack_)
        return false;

    if ((metaMap_ = allocateMap(MapKind_Meta, sizeof(Map))) == NULL)
        return false;
    metaMap_->init(metaMap_);

    return true;
}

Heap::~Heap()
{
    free(markStack_);
}

Map *
Heap::allocateMap(MapKind kind, size_t bytes)
{
    size_t actualBytes = Align(bytes, kObjectAlignment);

    Address addr = mapSpace_.allocate(actualBytes);
    if (!addr) {
        reportOutOfMemory();
        return NULL;
    }

    // Note that we use a NULL type initially, and reinterpret-cast to get
    // the Type pointer, so we can allocate the self-referencing meta type.
    Object *obj = InitAddressIntoObject(addr, NULL);
    Map *map = reinterpret_cast<Map *>(obj);
    map->setKind(kind);
    return map;
}

// To avoid marking recursively, we use a worklist to manage the marking
// process. Algorithm:
//
// (1) While the mark stack is not empty:
//  (a) Remove object |o| from the mark stack.
//  (b) For each reference |p| in object |o|,
//   (i)  Mark |p| as Black.
//   (ii) Push |p| onto the mark stack.
//   (iv) If the mark stack overflowde, mark |p| as Grey.
//
// We run this algorithm after pushing every root object onto the mark stack.
// After all roots have been processed, it may be that the mark stack
// overflowed. In this case, we use the following algorithm:
//
//  (1) while the mark stack has overflowed:
//   (a) Search the heap for grey objects, adding as many as possible to the
//       mark stack.
//   (b) Run the mark-stack-emptying algorithm above.
//
// This algorithm is guaranteed to terminate, because even if there are more
// grey objects than mark stack capacity, we eventually all objects will
// become grey, and grey objects are not re-pushed onto the mark stack. Thus,
// eventually, all objects will become black.
//
namespace ke {
class MarkStack
{
    Object **base_;
    Object **top_;
    Object **end_;
    bool overflowed_;

  public:
    MarkStack(Object **base, Object **end)
      : base_(base),
        top_(base),
        end_(end),
        overflowed_(false)
    {
    }

    void pushBlack(Object *ptr) {
        assert(HeapChunk::IsBlack(HeapChunk::GetMarkBit(ptr)));
        if (top_ >= end_) {
            MarkBit bit = HeapChunk::GetMarkBit(ptr);
            HeapChunk::BlackToGrey(bit);
            overflowed_ = true;
            return;
        }
        *top_++ = ptr;
    }

    Object *pop() {
        assert(!empty());
        return *--top_;
    }
    bool overflowed() const {
        return overflowed_;
    }
    bool empty() const {
        return base_ == top_;
    }
    void reset() {
        assert(empty());
        overflowed_ = true;
    }
};
} // namespace ke

// Iterates over an object's body, pushing its referenced objects onto the
// mark stack.
class MarkingVisitor : public ObjectVisitor<VirtualObjectVisitor>
{
    MarkStack &stack_;

  public:
    MarkingVisitor(MarkStack &stack)
      : stack_(stack)
    {
    }

    void visit(Object **ptr) {
        if (!*ptr)
            return;

        Object *obj = *ptr;

        // If the object is grey, we'll find it later when searching the heap.
        // If the object is black, it's already been scsanned.
        MarkBit bit = HeapChunk::GetMarkBit(obj);
        if (HeapChunk::IsBlackOrGrey(bit))
            return;

        // Place the object in the stack and optimistically mark it black. If
        // the stack overflows, it'll become grey instead.
        HeapChunk::WhiteToBlack(bit);
        stack_.pushBlack(obj);
    }
};

static void
ProcessMarkStack(MarkStack &stack)
{
    MarkingVisitor visitor(stack);

    while (!stack.empty()) {
        Object *obj = stack.pop();
        assert(HeapChunk::IsBlack(HeapChunk::GetMarkBit(obj)));

        visitor.iterate(obj);
    }
}

class RootMarkingVisitor : public VirtualObjectVisitor
{
    MarkStack &stack_;
    MarkingVisitor visitor_;

  public:
    RootMarkingVisitor(MarkStack &stack)
      : stack_(stack),
        visitor_(stack)
    {
    }

    void visit(Object **ptr) {
        // Flush the stack after every new root. Note that the root gets pushed
        // as black; that's okay. If the mark stack overflows, its children are
        // guaranteed to be at least grey.
        visitor_.visit(ptr);
        ProcessMarkStack(stack_);
    }
};

typedef ChunkObjectIteratorBase<ObjectSizeEvaluator> ChunkObjectIterator;

static inline bool
EnqueueGreyObject(MarkStack &stack, Object *obj)
{
    assert(!stack.overflowed());
    MarkBit bit = HeapChunk::GetMarkBit(obj);
    if (HeapChunk::IsGrey(bit)) {
        HeapChunk::GreyToBlack(bit);
        stack.pushBlack(obj);
        return !stack.overflowed();
    }
    return true;
}

static void
EnqueueGreyObjectsInChunks(MarkStack &stack, ChunkObjectIterator iter)
{
    for (Object *obj = iter.next(); obj; obj = iter.next()) {
        if (!EnqueueGreyObject(stack, obj))
            return;
    }
}

static void
EnqueueGreyObjectsInLargeChunks(MarkStack &stack, HeapChunk *first)
{
    for (HeapChunk *chunk = first; chunk; chunk = chunk->next()) {
        Object *obj = reinterpret_cast<Object *>(chunk->first());
        if (!EnqueueGreyObject(stack, obj))
            return;
    }
}

void
Heap::findGreyObjects(MarkStack &stack)
{
    EnqueueGreyObjectsInChunks(stack, lifoSpace_.begin(ObjectSizeEvaluator()));
    if (stack.overflowed())
        return;

    EnqueueGreyObjectsInChunks(stack, mapSpace_.begin(ObjectSizeEvaluator()));
    if (stack.overflowed())
        return;
    
    EnqueueGreyObjectsInChunks(stack, oldSpace_.begin(ObjectSizeEvaluator()));
    if (stack.overflowed())
        return;

    EnqueueGreyObjectsInLargeChunks(stack, largeSpace_.first());
    if (stack.overflowed())
        return;
}

void
Heap::mark()
{
    Object **qbase = reinterpret_cast<Object **>(markStack_);
    Object **qend = qbase + (kMarkStackSize / sizeof(Object *));
    MarkStack stack(qbase, qend);
    RootMarkingVisitor visitor(stack);

    // Attempt to mark all objects reachable from the root set.
    zone_->accept(&visitor);
    assert(stack.empty());

    // If the stack overflowed, we need to look through the entire heap for
    // grey objects to rescan. We keep doing this until there are no grey
    // objects left.
    while (stack.overflowed()) {
        stack.reset();
        findGreyObjects(stack);
        ProcessMarkStack(stack);
    }
}

// Use the mark bits to scan for live objects, and mark the regions in
// between them as free.
static void
SweepOldSpaceChunk(OldSpace *space, HeapChunk *chunk)
{
    uint32 *markbits = chunk->markBits();
    uint32 *markbitsEnd = markbits + HeapChunk::kMarkWords;
    Address freeStart = chunk->first();

    // Scoot up to the first usable markword. Since memory is reserved for
    // mark bits, not all addresses are allocatable, so quite a few of the
    // initial mark bits are wasted. (1KB of mark bytes on x86.)
    markbits = HeapChunk::GetMarkBit(reinterpret_cast<Object *>(chunk->first())).word();

    uint32 bits = *markbits;
    for (;;) {
        while (bits) {
            uint32 bit = FindRightmostBit(bits);
            Address object = HeapChunk::FromMarkBit(chunk, markbits, bit);

            assert(object >= freeStart);
            if (freeStart != object)
                space->markRegionFree(freeStart, object - freeStart);

            Object *obj = reinterpret_cast<Object *>(object);
            freeStart = object + SizeOfObject(obj, obj->map());

            // Since black is 10, we never need to clear the next bit.
            bits &= ~(1 << bit);
        }
        markbits++;
        bits = *markbits;
        if (markbits >= markbitsEnd)
            break;
    }

    if (freeStart == chunk->base())
        chunk->setEmpty();
    else if (freeStart != chunk->end())
        space->markRegionFree(freeStart, chunk->end() - freeStart);
}

static void
SweepOldSpace(OldSpace *space)
{
    space->resetFreeList();
    for (HeapChunk *chunk = space->first(); chunk; chunk = chunk->next()) {
        SweepOldSpaceChunk(space, chunk);
        chunk->resetMarkBits();
    }
    space->freeUnusedPages();
}

void
Heap::sweep()
{
    // Before destroying mark bits, sweep weak roots.
    zone_->symbols()->sweep();

    // Note: we don't bother sweeping the LIFO space.
    SweepOldSpace(&mapSpace_);
    SweepOldSpace(&oldSpace_);
    largeSpace_.sweep();
}

void
Heap::reclaim()
{
    mark();
    sweep();
}

void
Heap::reportOutOfMemory()
{
    zone_->reportOutOfMemory();
}