#include "slang-ir-any-value-marshalling.h"

#include "slang-ir-generics-lowering-context.h"
#include "slang-ir.h"
#include "slang-ir-insts.h"

namespace Slang
{
    // This is a subpass of generics lowering IR transformation.
    // This pass generates packing/unpacking functions for `AnyValue`s,
    // and replaces all `IRPackAnyValue` and `IRUnpackAnyValue` with calls to these
    // functions.
    struct AnyValueMarshallingContext
    {
        SharedGenericsLoweringContext* sharedContext;

        // Stores information about generated `AnyValue` struct types.
        struct AnyValueTypeInfo : RefObject
        {
            IRType* type; // The generated IR value for the `AnyValue<N>` struct type.
            List<IRStructKey*> fieldKeys; // `IRStructKey`s for the fields of the generated type.
        };

        Dictionary<IRIntegerValue, RefPtr<AnyValueTypeInfo>> generatedAnyValueTypes;

        struct MarshallingFunctionKey
        {
            IRType* originalType;
            IRIntegerValue anyValueSize;
            bool operator ==(MarshallingFunctionKey other)
            {
                return originalType == other.originalType && anyValueSize == other.anyValueSize;
            }
            HashCode getHashCode() const
            {
                return combineHash(Slang::getHashCode(originalType), Slang::getHashCode(anyValueSize));
            }
        };

        struct MarshallingFunctionSet
        {
            IRFunc* packFunc;
            IRFunc* unpackFunc;
        };

        // Stores the generated packing/unpacking functions for lookup.
        Dictionary<MarshallingFunctionKey, MarshallingFunctionSet> mapTypeMarshalingFunctions;

        AnyValueTypeInfo* ensureAnyValueType(IRAnyValueType* type)
        {
            auto size = getIntVal(type->getSize());
            if (auto typeInfo = generatedAnyValueTypes.TryGetValue(size))
                return typeInfo->Ptr();
            RefPtr<AnyValueTypeInfo> info = new AnyValueTypeInfo();
            IRBuilder builder;
            builder.sharedBuilder = &sharedContext->sharedBuilderStorage;
            builder.setInsertBefore(type);
            auto structType = builder.createStructType();
            info->type = structType;
            StringBuilder nameSb;
            nameSb << "AnyValue" << size;
            builder.addExportDecoration(structType, nameSb.getUnownedSlice());
            auto fieldCount = (size + sizeof(uint32_t) - 1) / sizeof(uint32_t);
            for (decltype(fieldCount) i = 0; i < fieldCount; i++)
            {
                auto key = builder.createStructKey();
                nameSb.Clear();
                nameSb << "field" << i;
                builder.addNameHintDecoration(key, nameSb.getUnownedSlice());
                nameSb << "_anyVal" << size;
                builder.addExportDecoration(key, nameSb.getUnownedSlice());
                builder.createStructField(structType, key, builder.getUIntType());
                info->fieldKeys.add(key);
            }
            generatedAnyValueTypes[size] = info;
            return info.Ptr();
        }

        struct TypeMarshallingContext
        {
            AnyValueTypeInfo* anyValInfo;
            uint32_t fieldOffset;
            uint32_t intraFieldOffset;
            IRType* uintPtrType;
            IRInst* anyValueVar;
            // Defines what to do with basic typed data elements.
            virtual void marshalBasicType(IRBuilder* builder, IRType* dataType, IRInst* concreteTypedVar) = 0;
            // Defines what to do with resource handle elements.
            virtual void marshalResourceHandle(IRBuilder* builder, IRType* dataType, IRInst* concreteTypedVar) = 0;
            // Validates that the type fits in the given AnyValueSize.
            // After calling emitMarshallingCode, `fieldOffset` will be increased to the required `AnyValue` size.
            // If this is larger than the provided AnyValue size, report a dianogstic. We might want to front load
            // this in a separate IR validation pass in the future, but this is the easiest way to report the
            // diagnostic now.
            void validateAnyTypeSize(DiagnosticSink* sink, IRType* concreteType)
            {
                if (fieldOffset > static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                {
                    sink->diagnose(concreteType->sourceLoc, Diagnostics::typeDoesNotFitAnyValueSize, concreteType);
                }
            }
        };

        void emitMarshallingCode(
            IRBuilder* builder,
            TypeMarshallingContext* context,
            IRInst* concreteTypedVar)
        {
            auto dataType = cast<IRPtrTypeBase>(concreteTypedVar->getDataType())->getValueType();
            switch (dataType->op)
            {
            case kIROp_IntType:
            case kIROp_FloatType:
            case kIROp_UIntType:
            case kIROp_UInt64Type:
            case kIROp_Int64Type:
            case kIROp_DoubleType:
            case kIROp_Int8Type:
            case kIROp_Int16Type:
            case kIROp_UInt8Type:
            case kIROp_UInt16Type:
            case kIROp_HalfType:
                context->marshalBasicType(builder, dataType, concreteTypedVar);
                break;
            case kIROp_VectorType:
            {
                auto vectorType = static_cast<IRVectorType*>(dataType);
                auto elementType = vectorType->getElementType();
                auto elementCount = getIntVal(vectorType->getElementCount());
                auto elementPtrType = builder->getPtrType(elementType);
                for (IRIntegerValue i = 0; i < elementCount; i++)
                {
                    auto elementAddr = builder->emitElementAddress(
                        elementPtrType,
                        concreteTypedVar,
                        builder->getIntValue(builder->getIntType(), i));
                    emitMarshallingCode(builder, context, elementAddr);
                }
                break;
            }
            case kIROp_MatrixType:
            {
                auto matrixType = static_cast<IRMatrixType*>(dataType);
                auto elementType = matrixType->getElementType();
                auto colCount = getIntVal(matrixType->getColumnCount());
                auto rowCount = getIntVal(matrixType->getRowCount());
                for (IRIntegerValue i = 0; i < colCount; i++)
                {
                    auto col = builder->emitElementAddress(
                        elementType,
                        concreteTypedVar,
                        builder->getIntValue(builder->getIntType(), i));
                    for (IRIntegerValue j = 0; j < rowCount; j++)
                    {
                        auto element = builder->emitElementExtract(
                            elementType,
                            col,
                            builder->getIntValue(builder->getIntType(), i));
                        emitMarshallingCode(builder, context, element);
                    }
                }
                break;
            }
            case kIROp_StructType:
            {
                auto structType = cast<IRStructType>(dataType);
                for (auto field : structType->getFields())
                {
                    auto fieldAddr = builder->emitFieldAddress(
                        builder->getPtrType(field->getFieldType()),
                        concreteTypedVar,
                        field->getKey());
                    emitMarshallingCode(builder, context, fieldAddr);
                }
                break;
            }
            case kIROp_ArrayType:
            {
                auto arrayType = cast<IRArrayType>(dataType);
                auto elementPtrType = builder->getPtrType(arrayType->getElementType());
                for (IRIntegerValue i = 0; i < getIntVal(arrayType->getElementCount()); i++)
                {
                    auto fieldAddr = builder->emitElementAddress(
                        elementPtrType,
                        concreteTypedVar,
                        builder->getIntValue(builder->getIntType(), i));
                    emitMarshallingCode(builder, context, fieldAddr);
                }
                break;
            }
            default:
                if (as<IRTextureTypeBase>(dataType) || as<IRSamplerStateTypeBase>(dataType))
                {
                    context->marshalResourceHandle(builder, dataType, concreteTypedVar);
                    return;
                }
                SLANG_UNIMPLEMENTED_X("Unimplemented type packing");
                break;
            }
        }

        struct TypePackingContext : TypeMarshallingContext
        {
            virtual void marshalBasicType(IRBuilder* builder, IRType* dataType, IRInst* concreteVar) override
            {
                switch (dataType->op)
                {
                case kIROp_IntType:
                case kIROp_FloatType:
                {
                    if (fieldOffset < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                    {
                        auto srcVal = builder->emitLoad(concreteVar);
                        auto dstVal = builder->emitBitCast(builder->getUIntType(), srcVal);
                        auto dstAddr = builder->emitFieldAddress(
                            uintPtrType,
                            anyValueVar,
                            anyValInfo->fieldKeys[fieldOffset]);
                        builder->emitStore(dstAddr, dstVal);
                    }
                    fieldOffset++;
                    break;
                }
                case kIROp_UIntType:
                {
                    if (fieldOffset < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                    {
                        auto srcVal = builder->emitLoad(concreteVar);
                        auto dstAddr = builder->emitFieldAddress(
                            uintPtrType,
                            anyValueVar,
                            anyValInfo->fieldKeys[fieldOffset]);
                        builder->emitStore(dstAddr, srcVal);
                    }
                    fieldOffset++;
                    break;
                }
                case kIROp_UInt64Type:
                case kIROp_Int64Type:
                case kIROp_DoubleType:
                case kIROp_Int8Type:
                case kIROp_Int16Type:
                case kIROp_UInt8Type:
                case kIROp_UInt16Type:
                case kIROp_HalfType:
                    SLANG_UNIMPLEMENTED_X("AnyValue type packing for non 32-bit elements");
                    break;
                default:
                    SLANG_UNREACHABLE("unknown basic type");
                }
            }

            virtual void marshalResourceHandle(IRBuilder* builder, IRType* dataType, IRInst* concreteVar) override
            {
                SLANG_UNUSED(dataType);
                if (fieldOffset + 1 < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                {
                    auto srcVal = builder->emitLoad(concreteVar);
                    auto uint64Val = builder->emitBitCast(builder->getUInt64Type(), srcVal);
                    auto lowBits = builder->emitConstructorInst(builder->getUIntType(), 1, &uint64Val);
                    auto shiftedBits = builder->emitShr(
                        builder->getUInt64Type(),
                        uint64Val,
                        builder->getIntValue(builder->getIntType(), 32));
                    auto highBits = builder->emitBitCast(builder->getUIntType(), shiftedBits);
                    auto dstAddr1 = builder->emitFieldAddress(
                        uintPtrType, anyValueVar, anyValInfo->fieldKeys[fieldOffset]);
                    builder->emitStore(dstAddr1, lowBits);
                    auto dstAddr2 = builder->emitFieldAddress(
                        uintPtrType, anyValueVar, anyValInfo->fieldKeys[fieldOffset + 1]);
                    builder->emitStore(dstAddr2, highBits);
                    fieldOffset += 2;
                }
            }
        };

        IRFunc* generatePackingFunc(IRType* type, IRAnyValueType* anyValueType)
        {
            IRBuilder builder;
            builder.sharedBuilder = &sharedContext->sharedBuilderStorage;
            builder.setInsertBefore(type);
            auto anyValInfo = ensureAnyValueType(anyValueType);

            auto func = builder.createFunc();

            StringBuilder nameSb;
            nameSb << "packAnyValue" << getIntVal(anyValueType->getSize());
            builder.addNameHintDecoration(func, nameSb.getUnownedSlice());
            // Currently we don't add linkage to the generated func, since we
            // do not have a way to compute mangled names from an IR entity.
            // This will leads to duplicate packing functions in linked code
            // but there won't be correctness issues.

            auto funcType = builder.getFuncType(1, &type, anyValInfo->type);
            func->setFullType(funcType);
            builder.setInsertInto(func);

            builder.emitBlock();

            auto param = builder.emitParam(type);
            auto concreteTypedVar = builder.emitVar(type);
            builder.emitStore(concreteTypedVar, param);
            auto resultVar = builder.emitVar(anyValInfo->type);

            TypePackingContext context;
            context.anyValInfo = anyValInfo;
            context.fieldOffset = context.intraFieldOffset = 0;
            context.uintPtrType = builder.getPtrType(builder.getUIntType());
            context.anyValueVar = resultVar;
            emitMarshallingCode(&builder, &context, concreteTypedVar);
            context.validateAnyTypeSize(sharedContext->sink, type);
            auto load = builder.emitLoad(resultVar);
            builder.emitReturn(load);
            return func;
        }

        struct TypeUnpackingContext : TypeMarshallingContext
        {
            virtual void marshalBasicType(IRBuilder* builder, IRType* dataType, IRInst* concreteVar) override
            {
                switch (dataType->op)
                {
                case kIROp_IntType:
                case kIROp_FloatType:
                {
                    if (fieldOffset < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                    {
                        auto srcAddr = builder->emitFieldAddress(
                            uintPtrType,
                            anyValueVar,
                            anyValInfo->fieldKeys[fieldOffset]);
                        auto srcVal = builder->emitLoad(srcAddr);
                        srcVal = builder->emitBitCast(dataType, srcVal);
                        builder->emitStore(concreteVar, srcVal);
                    }
                    fieldOffset++;
                    break;
                }
                case kIROp_UIntType:
                {
                    if (fieldOffset < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                    {
                        auto srcAddr = builder->emitFieldAddress(
                            uintPtrType,
                            anyValueVar,
                            anyValInfo->fieldKeys[fieldOffset]);
                        auto srcVal = builder->emitLoad(srcAddr);
                        builder->emitStore(concreteVar, srcVal);
                    }
                    fieldOffset++;
                    break;
                }
                case kIROp_UInt64Type:
                case kIROp_Int64Type:
                case kIROp_DoubleType:
                case kIROp_Int8Type:
                case kIROp_Int16Type:
                case kIROp_UInt8Type:
                case kIROp_UInt16Type:
                case kIROp_HalfType:
                    SLANG_UNIMPLEMENTED_X("AnyValue type packing for non 32-bit elements");
                    break;
                default:
                    SLANG_UNREACHABLE("unknown basic type");
                }
            }

            virtual void marshalResourceHandle(
                IRBuilder* builder, IRType* dataType, IRInst* concreteVar) override
            {
                if (fieldOffset + 1 < static_cast<uint32_t>(anyValInfo->fieldKeys.getCount()))
                {
                    auto srcAddr = builder->emitFieldAddress(
                        uintPtrType, anyValueVar, anyValInfo->fieldKeys[fieldOffset]);
                    auto lowBits = builder->emitLoad(srcAddr);
                    
                    auto srcAddr1 = builder->emitFieldAddress(
                        uintPtrType, anyValueVar, anyValInfo->fieldKeys[fieldOffset + 1]);
                    auto highBits = builder->emitLoad(srcAddr1);

                    auto combinedBits = builder->emitMakeUInt64(lowBits, highBits);
                    combinedBits = builder->emitBitCast(dataType, combinedBits);
                    builder->emitStore(concreteVar, combinedBits);
                    fieldOffset += 2;
                }
            }
        };

        IRFunc* generateUnpackingFunc(IRType* type, IRAnyValueType* anyValueType)
        {
            IRBuilder builder;
            builder.sharedBuilder = &sharedContext->sharedBuilderStorage;
            builder.setInsertBefore(type);
            auto anyValInfo = ensureAnyValueType(anyValueType);

            auto func = builder.createFunc();

            StringBuilder nameSb;
            nameSb << "unpackAnyValue" << getIntVal(anyValueType->getSize());
            builder.addNameHintDecoration(func, nameSb.getUnownedSlice());

            auto funcType = builder.getFuncType(1, &anyValInfo->type, type);
            func->setFullType(funcType);
            builder.setInsertInto(func);

            builder.emitBlock();

            auto param = builder.emitParam(anyValInfo->type);
            auto anyValueVar = builder.emitVar(anyValInfo->type);
            builder.emitStore(anyValueVar, param);
            auto resultVar = builder.emitVar(type);

            TypeUnpackingContext context;
            context.anyValInfo = anyValInfo;
            context.fieldOffset = context.intraFieldOffset = 0;
            context.uintPtrType = builder.getPtrType(builder.getUIntType());
            context.anyValueVar = anyValueVar;
            emitMarshallingCode(&builder, &context, resultVar);
            auto load = builder.emitLoad(resultVar);
            builder.emitReturn(load);
            return func;
        }

        // Ensures the marshalling functions between `type` and `anyValueType` are already generated.
        // Returns the generated marshalling functions.
        MarshallingFunctionSet ensureMarshallingFunc(IRType* type, IRAnyValueType* anyValueType)
        {
            auto size = getIntVal(anyValueType->getSize());
            MarshallingFunctionKey key;
            key.originalType = type;
            key.anyValueSize = size;
            MarshallingFunctionSet funcSet;
            if (mapTypeMarshalingFunctions.TryGetValue(key, funcSet))
                return funcSet;
            funcSet.packFunc = generatePackingFunc(type, anyValueType);
            funcSet.unpackFunc = generateUnpackingFunc(type, anyValueType);
            mapTypeMarshalingFunctions[key] = funcSet;
            return funcSet;
        }

        void processPackInst(IRPackAnyValue* packInst)
        {
            auto operand = packInst->getValue();
            auto func = ensureMarshallingFunc(
                operand->getDataType(),
                cast<IRAnyValueType>(packInst->getDataType()));
            IRBuilder builderStorage;
            auto builder = &builderStorage;
            builder->sharedBuilder = &sharedContext->sharedBuilderStorage;
            builder->setInsertBefore(packInst);
            auto callInst = builder->emitCallInst(packInst->getDataType(), func.packFunc, 1, &operand);
            packInst->replaceUsesWith(callInst);
            packInst->removeAndDeallocate();
        }

        void processUnpackInst(IRUnpackAnyValue* unpackInst)
        {
            auto operand = unpackInst->getValue();
            auto func = ensureMarshallingFunc(
                unpackInst->getDataType(),
                cast<IRAnyValueType>(operand->getDataType()));
            IRBuilder builderStorage;
            auto builder = &builderStorage;
            builder->sharedBuilder = &sharedContext->sharedBuilderStorage;
            builder->setInsertBefore(unpackInst);
            auto callInst = builder->emitCallInst(unpackInst->getDataType(), func.unpackFunc, 1, &operand);
            unpackInst->replaceUsesWith(callInst);
            unpackInst->removeAndDeallocate();
        }

        void processAnyValueType(IRAnyValueType* type)
        {
            auto info = ensureAnyValueType(type);
            type->replaceUsesWith(info->type);
        }

        void processInst(IRInst* inst)
        {
            if (auto packInst = as<IRPackAnyValue>(inst))
            {
                processPackInst(packInst);
            }
            else if (auto unpackInst = as<IRUnpackAnyValue>(inst))
            {
                processUnpackInst(unpackInst);
            }
        }

        void processModule()
        {
            // We start by initializing our shared IR building state,
            // since we will re-use that state for any code we
            // generate along the way.
            //
            SharedIRBuilder* sharedBuilder = &sharedContext->sharedBuilderStorage;
            sharedBuilder->module = sharedContext->module;
            sharedBuilder->session = sharedContext->module->session;

            sharedContext->addToWorkList(sharedContext->module->getModuleInst());

            while (sharedContext->workList.getCount() != 0)
            {
                IRInst* inst = sharedContext->workList.getLast();

                sharedContext->workList.removeLast();
                sharedContext->workListSet.Remove(inst);

                processInst(inst);

                for (auto child = inst->getLastChild(); child; child = child->getPrevInst())
                {
                    sharedContext->addToWorkList(child);
                }
            }

            // Finally, replace all `AnyValueType` with the actual struct type that implements it.
            for (auto inst : sharedContext->module->getModuleInst()->getChildren())
            {
                if (auto anyValueType = as<IRAnyValueType>(inst))
                    processAnyValueType(anyValueType);
            }
            // Because we replaced all `AnyValueType` uses, some old type definitions (e.g. PtrType(AnyValueType))
            // will become duplicates with new types we introduced (e.g. PtrType(AnyValueStruct)), and therefore
            // invalidates our `globalValueNumberingMap` hash map. We need to rebuild it.
            sharedContext->sharedBuilderStorage.deduplicateAndRebuildGlobalNumberingMap();
        }
    };

    void generateAnyValueMarshallingFunctions(SharedGenericsLoweringContext* sharedContext)
    {
        AnyValueMarshallingContext context;
        context.sharedContext = sharedContext;
        context.processModule();
    }
}