/* vim: set ts=4 sw=4 tw=99 et:
 *
 * Copyright (C) 2012 David Anderson
 *
 * This file is part of JITCraft.
 *
 * JITCraft 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.
 * 
 * Foobar 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
 * JITCraft. If not, see http://www.gnu.org/licenses/.
 */
#ifndef _include_jc_object_visitor_h_
#define _include_jc_object_visitor_h_

#include "Objects.h"
#include "Types.h"
#include "FixedArray.h"
#include "Code.h"
#include "Reference.h"
#include "Modules.h"
#include "Functions.h"
#include "Array.h"
#include "Strings.h"
#include "Environments.h"
#include "Structures.h"
#include "Box.h"

namespace ke {

// An ObjectVisitor is used for traversing object references in the heap. Since
// this pattern must be lightning fast, we avoid virtual dispatch by forcing
// visitors to compose themselves with the visiting logic.
template <typename T>
class ObjectVisitor : public T
{
  private:
    void visitField(Object *obj, size_t offset) {
        this->visit(reinterpret_cast<Object **>(Address(obj) + offset));
    }

    size_t iterateType(Object *obj) {
        visitField(obj, Type::offsetOfContained());
        visitField(obj, Type::offsetOfReturnType());
        visitField(obj, Type::offsetOfParameters());
        visitField(obj, Type::offsetOfName());
        visitField(obj, Type::offsetOfFields());
        visitField(obj, Type::offsetOfDefaults());
        visitField(obj, Type::offsetOfNewMap());
        return sizeof(Type);
    }

    size_t iterateFixedArray(FixedArray *array) {
        unsigned length = array->length();
        Object **ptr = reinterpret_cast<Object **>(array + 1);
        for (unsigned i = 0; i < length; i++)
            this->visit(ptr + i);
        return array->objSize();
    }

    size_t iterateByteArray(ByteArray *array) {
        return array->objSize();
    }

    size_t iterateCode(Object *obj) {
        visitField(obj, Code::offsetOfModule());
        visitField(obj, Code::offsetOfBytecode());
        visitField(obj, Code::offsetOfLocals());
        visitField(obj, Code::offsetOfObjects());
        visitField(obj, Code::offsetOfSafepoints());
        visitField(obj, Code::offsetOfSourceMap());
        return sizeof(Code);
    }

    size_t iterateReference(Object *obj) {
        visitField(obj, Reference::offsetOfBase());
        return sizeof(Reference);
    }

    size_t iterateModule(Object *obj) {
        visitField(obj, Module::offsetOfFunctions());
        visitField(obj, Module::offsetOfPublics());
        visitField(obj, Module::offsetOfNatives());
        visitField(obj, Module::offsetOfFile());
        visitField(obj, Module::offsetOfEnvironment());
        visitField(obj, Module::offsetOfCode());
        return sizeof(Module);
    }

    size_t iterateFunction(Object *obj) {
        visitField(obj, Function::offsetOfCode());
        visitField(obj, Function::offsetOfType());
        visitField(obj, Function::offsetOfName());
        return sizeof(Function);
    }

    size_t iterateDescriptor(Object *obj) {
        visitField(obj, Descriptor::offsetOfName());
        visitField(obj, Descriptor::offsetOfType());
        return sizeof(Descriptor);
    }

    size_t iterateNative(Object *obj) {
        visitField(obj, Native::offsetOfType());
        visitField(obj, Native::offsetOfName());
        return sizeof(Native);
    }

    size_t iterateString(Object *obj) {
        return String::cast(obj)->objSize();
    }

    size_t iterateArray(Array *array, ArrayMap *map) {
        assert(map->type()->isArray());
        if (map->type()->contained()->isTraceable()) {
            unsigned length = array->length();
            Object **vec = array->objVec();
            for (unsigned i = 0; i < length; i++)
                this->visit(vec + i);
        }
        return Array::SizeOf(array);
    }

    size_t iterateBox(Box *obj) {
        visitField(obj, Box::offsetOfType());
        if (obj->type()->isTraceable())
            visitField(obj, Box::offsetOfObj());
        return sizeof(Box);
    }

    size_t iterateArrayMap(ArrayMap *map) {
        visitField(map, ArrayMap::offsetOfType());
        return sizeof(ArrayMap);
    }

    size_t iterateEnvironmentSlotDescriptorArray(EnvironmentSlotDescriptorArray *env) {
        unsigned length = env->length();
        EnvironmentSlotDescriptor *desc = env->slots();
        for (unsigned i = 0; i < length; i++)
            this->visit(desc->type.addr());
        return env->objSize();
    }

    size_t iterateEnvironmentMap(EnvironmentMap *map) {
        visitField(map, EnvironmentMap::offsetOfSlots());
        return map->objSize();
    }

    size_t iterateStructMap(StructMap *map) {
        visitField(map, StructMap::offsetOfType());
        visitField(map, StructMap::offsetOfComposed());
        return map->objSize();
    }

    size_t iterateEnvironment(Environment *env) {
        EnvironmentMap *map = env->map();
        const unsigned *offsets = map->traceOffsets();
        unsigned noffsets = map->traceOffsetCount();
        for (unsigned i = 0; i < noffsets; i++)
            visitField(env, Environment::offsetOfSlotOffset(offsets[i]));
        return env->objSize();
    }

    size_t iterateStruct(Struct *obj) {
        StructMap *map = obj->map();
        const unsigned *offsets = map->traceOffsets();
        unsigned noffsets = map->traceOffsetCount();
        for (unsigned i = 0; i < noffsets; i++)
            visitField(obj, Struct::offsetOfSlotOffset(offsets[i]));
        return obj->objSize();
    }

  private:
    size_t iterateObject(Object *obj) {
        visitField(obj, Object::offsetOfMap());
        return iterate(obj, obj->map());
    }

    size_t iterateObject(Object *obj, Map *map) {
        switch (map->kind()) {
          case MapKind_Type:
            return iterateType(obj);

          case MapKind_FixedArray:
            return iterateFixedArray(FixedArray::cast(obj));

          case MapKind_ByteArray:
            return iterateByteArray(ByteArray::cast(obj));

          case MapKind_Code:
            return iterateCode(obj);

          case MapKind_Reference:
            return iterateReference(obj);

          case MapKind_Module:
            return iterateModule(obj);

          case MapKind_Function:
            return iterateFunction(obj);

          case MapKind_Descriptor:
            return iterateDescriptor(obj);

          case MapKind_String:
            return iterateString(obj);

          case MapKind_Array:
            return Align(iterateArray(Array::cast(obj), ArrayMap::cast(map)), kObjectAlignment);

          case MapKind_Box:
            return iterateBox(Box::cast(obj));

          case MapKind_Struct:
            return iterateStruct(Struct::cast(obj));

          case MapKind_EnvironmentSlotDescriptorArray:
            return iterateEnvironmentSlotDescriptorArray(EnvironmentSlotDescriptorArray::cast(obj));

          case MapKind_Environment:
            return iterateEnvironment(Environment::cast(obj));

          case MapKind_Native:
            return iterateNative(obj);

          case MapKind_MetaArray:
            return iterateArrayMap(ArrayMap::cast(obj));

          case MapKind_MetaEnvironment:
            return iterateEnvironmentMap(EnvironmentMap::cast(obj));

          case MapKind_MetaStruct:
            return iterateStructMap(StructMap::cast(obj));

          default:
            assert(map->kind() == MapKind_Meta);
            return sizeof(Map);
        }
    }

  public:
    size_t iterate(Object *obj) {
        return Align(iterateObject(obj), kObjectAlignment);
    }

    size_t iterate(Object *obj, Map *map) {
        return Align(iterateObject(obj, map), kObjectAlignment);
    }
};

static inline size_t
SizeOfObjectRaw(Object *obj, Map *map)
{
    switch (map->kind()) {
      case MapKind_Type:
        return sizeof(Type);

      case MapKind_FixedArray:
        return FixedArray::cast(obj)->objSize();

      case MapKind_ByteArray:
        return ByteArray::cast(obj)->objSize();

      case MapKind_Code:
        return sizeof(Code);

      case MapKind_Reference:
        return sizeof(Reference);

      case MapKind_Module:
        return sizeof(Module);

      case MapKind_Descriptor:
        return sizeof(Descriptor);

      case MapKind_Function:
        return sizeof(Function);

      case MapKind_String:
        return String::cast(obj)->objSize();

      case MapKind_Native:
        return sizeof(Native);

      case MapKind_Array:
        return Align(Array::SizeOf(Array::cast(obj)), kObjectAlignment);

      case MapKind_Box:
        return sizeof(Box);

      case MapKind_Struct:
        return Struct::cast(obj)->objSize();

      case MapKind_EnvironmentSlotDescriptorArray:
        return EnvironmentSlotDescriptorArray::cast(obj)->objSize();

      case MapKind_Environment:
        return Environment::cast(obj)->objSize();

      case MapKind_MetaArray:
        return sizeof(ArrayMap);

      case MapKind_MetaEnvironment:
        return EnvironmentMap::cast(obj)->objSize();

      case MapKind_MetaStruct:
        return StructMap::cast(obj)->objSize();

      case MapKind_FreeRegion:
        return *reinterpret_cast<size_t *>(reinterpret_cast<Address>(obj + 1));

      default:
        assert(map->kind() == MapKind_Meta);
        return sizeof(Map);
    }
}

static inline size_t
SizeOfObject(Object *obj, Map *map)
{
    return Align(SizeOfObjectRaw(obj, map), kObjectAlignment);
}

class ObjectSizeEvaluator
{
  public:
    size_t operator()(Object *obj, Map *map) {
        return SizeOfObject(obj, map);
    }
};

}

#endif // _include_jc_object_visitor_h_