/* 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_->sweepWeakRoots();

    // 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();

}