// slang-legalize-types.h
#ifndef SLANG_LEGALIZE_TYPES_H_INCLUDED
#define SLANG_LEGALIZE_TYPES_H_INCLUDED

// This file and `legalize-types.cpp` implement the core
// logic for taking a `Type` as produced by the front-end,
// and turning it into a suitable representation for use
// on a particular back-end.
//
// The main work applies to aggregate (e.g., `struct`) types,
// since various targets have rules about what is and isn't
// allowed in an aggregate (or where aggregates are allowed
// to be used).
//
// We might completely replace an aggregate `Type` with a
// "pseudo-type" that is just the enumeration of its field
// types (sort of a tuple type) so that a variable declared
// with the original type should be transformed into a
// bunch of individual variables.
//
// Alternatively, we might replace an aggregate type, where
// only *some* of the fields are illegal with a combination
// of an aggregate (containing the legal/legalized fields),
// and some extra tuple-ified fields.

#include "../core/slang-basic.h"
#include "slang-ir-insts.h"
#include "slang-syntax.h"
#include "slang-type-layout.h"

#include "../compiler-core/slang-name.h"

namespace Slang
{

struct IRBuilder;

struct LegalTypeImpl : RefObject
{
};
struct ImplicitDerefType;
struct TuplePseudoType;
struct PairPseudoType;
struct PairInfo;
struct LegalElementWrapping;
struct WrappedBufferPseudoType;

    /// A flavor for types or values that arise during legalization.
enum class LegalFlavor
{
        /// Nothing: an empty type or value. Equivalent to `void`.
    none,

        /// A simple type/value that can be represented as an `IRType*` or `IRInst*`
    simple,

        /// Logically, a pointer-like type/value, but represented as the type/value being pointed type.
    implicitDeref,

        /// A compound type/value made up of the constituent fields of some original value.
    tuple,

        /// A type/value that was split into "ordinary" and "special" parts.
    pair,

        /// A type/value that represents, e.g., `ConstantBuffer<T>` where `T` needed legalization.
    wrappedBuffer,
};

struct LegalType
{
    typedef LegalFlavor Flavor;

    Flavor              flavor = Flavor::none;
    RefPtr<RefObject>   obj;
    IRType*             irType = nullptr;

    static LegalType simple(IRType* type)
    {
        LegalType result;
        result.flavor = Flavor::simple;
        result.irType = type;
        return result;
    }

    IRType* getSimple() const
    {
        SLANG_ASSERT(flavor == Flavor::simple);
        return irType;
    }

    static LegalType implicitDeref(
        LegalType const& valueType);

    RefPtr<ImplicitDerefType> getImplicitDeref() const
    {
        SLANG_ASSERT(flavor == Flavor::implicitDeref);
        return obj.as<ImplicitDerefType>();
    }

    static LegalType tuple(
        RefPtr<TuplePseudoType>   tupleType);

    RefPtr<TuplePseudoType> getTuple() const
    {
        SLANG_ASSERT(flavor == Flavor::tuple);
        return obj.as<TuplePseudoType>();
    }

    static LegalType pair(
        RefPtr<PairPseudoType> pairType);

    static LegalType pair(
        LegalType const&    ordinaryType,
        LegalType const&    specialType,
        RefPtr<PairInfo>    pairInfo);

    RefPtr<PairPseudoType> getPair() const
    {
        SLANG_ASSERT(flavor == Flavor::pair);
        return obj.as<PairPseudoType>();
    }

    static LegalType makeWrappedBuffer(
        IRType*             simpleType,
        LegalElementWrapping const&   elementInfo);

    RefPtr<WrappedBufferPseudoType> getWrappedBuffer() const
    {
        SLANG_ASSERT(flavor == Flavor::wrappedBuffer);
        return obj.as<WrappedBufferPseudoType>();
    }
};

struct LegalElementWrappingObj : RefObject
{
};

struct SimpleLegalElementWrappingObj;
struct ImplicitDerefLegalElementWrappingObj;
struct PairLegalElementWrappingObj;
struct TupleLegalElementWrappingObj;

    /// Information on how the element type of a buffer needs to be wrapped.
struct LegalElementWrapping
{
    typedef LegalFlavor Flavor;

    Flavor flavor;
    RefPtr<LegalElementWrappingObj> obj;

    static LegalElementWrapping makeVoid();
    static LegalElementWrapping makeSimple(IRStructKey* key, IRType* type);
    static LegalElementWrapping makeImplicitDeref(LegalElementWrapping const& field);
    static LegalElementWrapping makePair(
        LegalElementWrapping const& ordinary,
        LegalElementWrapping const& special,
        PairInfo* pairInfo);
    static LegalElementWrapping makeTuple(TupleLegalElementWrappingObj* obj);

    RefPtr<SimpleLegalElementWrappingObj> getSimple() const;
    RefPtr<ImplicitDerefLegalElementWrappingObj> getImplicitDeref() const;
    RefPtr<PairLegalElementWrappingObj> getPair() const;
    RefPtr<TupleLegalElementWrappingObj> getTuple() const;
};

struct SimpleLegalElementWrappingObj : LegalElementWrappingObj
{
    IRStructKey*    key;
    IRType*         type;
};

struct ImplicitDerefLegalElementWrappingObj : LegalElementWrappingObj
{
    LegalElementWrapping field;
};

struct PairLegalElementWrappingObj : LegalElementWrappingObj
{
    LegalElementWrapping ordinary;
    LegalElementWrapping special;
    RefPtr<PairInfo> pairInfo;
};

struct TupleLegalElementWrappingObj : LegalElementWrappingObj
{
    struct Element
    {
        IRStructKey*    key;
        LegalElementWrapping      field;
    };

    List<Element> elements;
};

// Represents the pseudo-type of a type that is pointer-like
// (and thus requires dereferencing, even if implicit), but
// was legalized to just use the type of the pointed-type value.
//
// The two cases where this comes up are:
//
//  1. When we have a type like `ConstantBuffer<Texture2D>` that
//  implies a level of indirection, but need to legalize it to just
//  `Texture2D`, which eliminates that indirection.
//
//  2. When we have a type like `ExistentialBox<Foo>` that will
//  become just a `Foo` field, but which needs to be allocated
//  out-of-line from the rest of its enclosing type.
//
struct ImplicitDerefType : LegalTypeImpl
{
    LegalType valueType;
};

// Represents the pseudo-type for a compound type
// that had to be broken apart because it contained
// one or more fields of types that shouldn't be
// allowed in aggregates.
//
// A tuple pseduo-type will have an element for
// each field of the original type, that represents
// the legalization of that field's type.
//
// It optionally also contains an "ordinary" type
// that packs together any per-field data that
// itself has (or contains) an ordinary type.
struct TuplePseudoType : LegalTypeImpl
{
    // Represents one element of the tuple pseudo-type
    struct Element
    {
        // The field that this element replaces
        IRStructKey*    key;

        // The legalized type of the element
        LegalType       type;
    };

    // All of the elements of the tuple pseduo-type.
    List<Element>   elements;
};

struct IRStructKey;

struct PairInfo : RefObject
{
    typedef unsigned int Flags;
    enum
    {
        kFlag_hasOrdinary = 0x1,
        kFlag_hasSpecial = 0x2,
    };


    struct Element
    {
        // The original field the element represents
        IRStructKey* key;

        // The conceptual type of the field.
        // If both the `hasOrdinary` and
        // `hasSpecial` bits are set, then
        // this is expected to be a
        // `LegalType::Flavor::pair`
        LegalType   type;

        // Is the value represented on
        // the ordinary side, the special
        // side, or both?
        Flags       flags;

        // If the type of this element is
        // itself a pair type (that is,
        // it both `hasOrdinary` and `hasSpecial`)
        // then this is the `PairInfo` for that
        // pair type:
        RefPtr<PairInfo> fieldPairInfo;
    };

    // For a pair type or value, we need to track
    // which fields are on which side(s).
    List<Element> elements;

    Element* findElement(IRStructKey* key)
    {
        for (auto& ee : elements)
        {
            if(ee.key == key)
                return &ee;
        }
        return nullptr;
    }
};

struct PairPseudoType : LegalTypeImpl
{
    // Any field(s) with ordinary types will
    // get captured here, usually as a single
    // `simple` or `implicitDeref` type.
    LegalType ordinaryType;

    // Any fields with "special" (not ordinary)
    // types will get captured here (usually
    // with a tuple).
    LegalType specialType;

    // The `pairInfo` field helps to tell us which members
    // of the original aggregate type appear on which side(s)
    // of the new pair type.
    RefPtr<PairInfo> pairInfo;
};


struct WrappedBufferPseudoType : LegalTypeImpl
{
    // The actual IR type that was used for the buffer.
    IRType*     simpleType;

    // Adjustments that need to be made when fetching
    // an element from this buffer type.
    //
    LegalElementWrapping  elementInfo;
};

//

IRTypeLayout* getDerefTypeLayout(
    IRTypeLayout* typeLayout);

IRVarLayout* getFieldLayout(
    IRTypeLayout*     typeLayout,
    IRInst*         fieldKey);

    /// Represents a "chain" of variables leading to some leaf field.
    ///
    /// Consider code like:
    ///
    ///     struct Branch { int leaf; }
    ///     struct Tree { Branch left; Branch right; }
    ///     cbuffer Forest
    ///     {
    ///         int maxTreeHeight;
    ///         Tree tree;
    ///     }
    ///
    /// If we ask "what is the offset of `leaf`" the simple answer is zero,
    /// but sometimes we are talking about `Forest.tree.right.leaf` which
    /// will have a very different offset. In Slang parameters can consume
    /// various (and multiple) resource kinds, so a single offset can't
    /// be tunneled down through most recursive procedures.
    ///
    /// Instead we use a "chain" that works up through the stack, and
    /// records the path from leaf field like `leaf` up to whatever
    /// variable is the root for the curent operation.
    ///
    /// Operations like computing an offset can then be encoded by
    /// starting with zero and then walking up the chain and adding in
    /// offsets as encountered.
    ///
struct SimpleLegalVarChain
{
    // The next link up the chain, or null if this is the end.
    SimpleLegalVarChain*    next = nullptr;

    // The layout for the variable at this link in thain.
    IRVarLayout*            varLayout = nullptr;
};

    /// A "chain" of variable declarations that can handle both primary and "pending" data.
    ///
    /// In the presence of interface-type fields, a single variable may
    /// have data that sits in two distinct allocations, and may have
    /// `VarLayout`s that represent offseting into each of those
    /// allocations.
    ///
    /// A `LegalVarChain` tracks two distinct `SimpleVarChain`s: one for
    /// the primary/ordinary data allocation, and one for any pending
    /// data.
    ///
    /// It is okay if the primary/pending chains have different numbers
    /// of links in them.
    ///
    /// Offsets for particular resource kinds in the primary or pending
    /// data allocation can be queried on the appropriate sub-chain.
    ///
struct LegalVarChain
{
    // The chain of variables that represents the primary allocation.
    SimpleLegalVarChain*    primaryChain = nullptr;

    // The chain of variables that represents the pending allocation.
    SimpleLegalVarChain*    pendingChain = nullptr;
};

    /// RAII type for adding a link to a `LegalVarChain` as needed.
    ///
    /// This type handles the bookkeeping for creating a `LegalVarChain`
    /// that links in one more variable. It will add a link to each of
    /// the primary and pending sub-chains if and only if there is non-null
    /// layout information for the primary/pending case.
    ///
    /// Typical usage in a recursive function is:
    ///
    ///     void someRecursiveFunc(LegalVarChain const& outerChain, ...)
    ///     {
    ///         if(auto subVar = needToRecurse(...))
    ///         {
    ///             LegalVarChainLink subChain(outerChain, subVar);
    ///             someRecursiveFunc(subChain, ...);
    ///         }
    ///         ...
    ///     }
    ///
struct LegalVarChainLink : LegalVarChain
{
        /// Default constructor: yields an empty chain.
    LegalVarChainLink()
    {
    }

        /// Copy constructor: yields a copy of the `parent` chain.
    LegalVarChainLink(LegalVarChain const& parent)
        : LegalVarChain(parent)
    {}

        /// Construct a chain that extends `parent` with `varLayout`, if it is non-null.
    LegalVarChainLink(LegalVarChain const& parent, IRVarLayout* varLayout)
        : LegalVarChain(parent)
    {
        if( varLayout )
        {
            primaryLink.next = parent.primaryChain;
            primaryLink.varLayout = varLayout;
            primaryChain = &primaryLink;

            if( auto pendingVarLayout = varLayout->getPendingVarLayout() )
            {
                pendingLink.next = parent.pendingChain;
                pendingLink.varLayout = pendingVarLayout;
                pendingChain = &pendingLink;
            }
        }
    }

    SimpleLegalVarChain primaryLink;
    SimpleLegalVarChain pendingLink;
};

IRVarLayout* createVarLayout(
    IRBuilder*              irBuilder,
    LegalVarChain const&    varChain,
    IRTypeLayout*           typeLayout);

IRVarLayout* createSimpleVarLayout(
    IRBuilder*              irBuilder,
    SimpleLegalVarChain*    varChain,
    IRTypeLayout*           typeLayout);

//
// The result of legalizing an IR value will be
// represented with the `LegalVal` type. It is exposed
// in this header (rather than kept as an implementation
// detail, because the AST-based legalization logic needs
// a way to find the post-legalization version of a
// global name).
//
// TODO: We really shouldn't have this structure exposed,
// and instead should really be constructing AST-side
// `LegalExpr` values on-demand whenever we legalize something
// in the IR that will need to be used by the AST, and then
// store *those* in a map indexed in mangled names.
//

struct LegalValImpl : RefObject
{
};
struct TuplePseudoVal;
struct PairPseudoVal;
struct WrappedBufferPseudoVal;

struct LegalVal
{
    typedef LegalFlavor Flavor;

    Flavor              flavor = Flavor::none;
    RefPtr<RefObject>   obj;
    IRInst*             irValue = nullptr;

    static LegalVal simple(IRInst* irValue)
    {
        LegalVal result;
        result.flavor = Flavor::simple;
        result.irValue = irValue;
        return result;
    }

    IRInst* getSimple() const
    {
        SLANG_ASSERT(flavor == Flavor::simple);
        return irValue;
    }

    static LegalVal tuple(RefPtr<TuplePseudoVal> tupleVal);

    RefPtr<TuplePseudoVal> getTuple() const
    {
        SLANG_ASSERT(flavor == Flavor::tuple);
        return obj.as<TuplePseudoVal>();
    }

    static LegalVal implicitDeref(LegalVal const& val);
    LegalVal getImplicitDeref() const;

    static LegalVal pair(RefPtr<PairPseudoVal> pairInfo);
    static LegalVal pair(
        LegalVal const&     ordinaryVal,
        LegalVal const&     specialVal,
        RefPtr<PairInfo>    pairInfo);

    RefPtr<PairPseudoVal> getPair() const
    {
        SLANG_ASSERT(flavor == Flavor::pair);
        return obj.as<PairPseudoVal>();
    }

    static LegalVal wrappedBuffer(
        LegalVal const& baseVal,
        LegalElementWrapping const& elementInfo);

    RefPtr<WrappedBufferPseudoVal> getWrappedBuffer() const
    {
        SLANG_ASSERT(flavor == Flavor::wrappedBuffer);
        return obj.as<WrappedBufferPseudoVal>();
    }
};

struct TuplePseudoVal : LegalValImpl
{
    struct Element
    {
        IRStructKey*    key;
        LegalVal        val;
    };

    List<Element>   elements;
};

struct PairPseudoVal : LegalValImpl
{
    LegalVal ordinaryVal;
    LegalVal specialVal;

    // The info to tell us which fields
    // are on which side(s)
    RefPtr<PairInfo>  pairInfo;
};

struct ImplicitDerefVal : LegalValImpl
{
    LegalVal val;
};

struct WrappedBufferPseudoVal : LegalValImpl
{
    LegalVal                base;
    LegalElementWrapping    elementInfo;
};

//

    /// Information about a function that has been legalized
    ///
    /// This type is used to track any information about the function
    /// and its signature that might be relevant to the legalization
    /// of instructions inside the function body.
    ///
struct LegalFuncInfo : RefObject
{
        /// Any parameters that were added to the function signature
        /// to represent the function result after legalization.
        ///
        /// It is possible that the result type of a function needed
        /// to be split into multiple types, and as a result a single
        /// function result couldn't return all of them.
        ///
        /// This array is a list of `out` parameters created to represent
        /// additional function results. Because they are `out` parameters,
        /// each is a *pointer* to a value of the relevant type.
        ///
    List<IRInst*> resultParamVals;
};

//

    /// Context that drives type legalization
    ///
    /// This type is an abstract base class, and there are
    /// customization points that a concrete pass needs to
    /// override (e.g., to specify what needs to be legalized).
struct IRTypeLegalizationContext
{
    Session*    session;
    IRModule*   module;
    IRBuilder*  builder;

    IRBuilder builderStorage;

    IRTypeLegalizationContext(
        IRModule* inModule);

    // When inserting new globals, put them before this one.
    IRInst* insertBeforeGlobal = nullptr;

    // When inserting new parameters, put them before this one.
    IRParam* insertBeforeParam = nullptr;

    Dictionary<IRInst*, LegalVal> mapValToLegalVal;

    IRVar* insertBeforeLocalVar = nullptr;

    // store instructions that have been replaced here, so we can free them
    // when legalization has done
    OrderedHashSet<IRInst*> replacedInstructions;

    Dictionary<IRType*, LegalType> mapTypeToLegalType;

        /// Map a function to information about how it was legalized.
        ///
        /// Note that entries are only created if there is somehting for them
        /// to represent, so many functions may lack entries in this map even
        /// after legalization.
        ///
    Dictionary<IRFunc*, RefPtr<LegalFuncInfo>> mapFuncToInfo;

        /// 
        /// Special handling for pointer types. If we have a situation where 
        /// a type could end up in a loop pointing to itself, the activePointerValues
        /// stack records which pointer value types (ie the thing being pointed to)
        /// are "active". The usedCount is to indicate how many times the type was
        /// used whilst active. If it's !=0, we should check the assumption about what 
        /// should have been produced.
        /// 
    struct PointerValue
    {
        IRType* type = nullptr;
        Index usedCount = 0;
    };

    List<PointerValue> activePointerValues;

    IRBuilder* getBuilder() { return builder; }

        /// Customization point to decide what types are "special."
        ///
        /// When legalizing a `struct` type, any fields that have "special"
        /// types will get moved out of the `struc` itself.
    virtual bool isSpecialType(IRType* type) = 0;

        /// Customization point to decide what types are "simple."
        ///
        /// When a type is "simple" it means that it should not be changed
        /// during legalization.
    virtual bool isSimpleType(IRType* type) = 0;

        /// Customization point to construct uniform-buffer/block types.
        ///
        /// This function will only be called if `legalElementType` is
        /// somehow non-trivial.
        ///
    virtual LegalType createLegalUniformBufferType(
        IROp        op,
        LegalType   legalElementType) = 0;
};

// This typedef exists to support pre-existing code from when
// `IRTypeLegalizationContext` and `TypeLegalizationContext` were
// two different types that had to coordinate.
typedef struct IRTypeLegalizationContext TypeLegalizationContext;

LegalType legalizeType(
    TypeLegalizationContext*    context,
    IRType*                     type);

/// Try to find the module that (recursively) contains a given declaration.
ModuleDecl* findModuleForDecl(
    Decl*   decl);

    /// Create a uniform buffer type suitable for resource legalization.
    ///
    /// This will allocate a real buffer for the ordinary data (if any),
    /// and leave the special data (if any) as a tuple.
    ///
LegalType createLegalUniformBufferTypeForResources(
    TypeLegalizationContext*    context,
    IROp                        op,
    LegalType                   legalElementType);

    /// Create a uniform buffer type suitable for existential legalization.
    ///
    /// This will allocate a real uniform buffer for *all* the data, by
    /// declaring an intermediate `struct` type to hold the ordinary and
    /// special (existential-box) fields, if required.
    ///
LegalType createLegalUniformBufferTypeForExistentials(
    TypeLegalizationContext*    context,
    IROp                        op,
    LegalType                   legalElementType);




void legalizeExistentialTypeLayout(
    IRModule*       module,
    DiagnosticSink* sink);

void legalizeResourceTypes(
    IRModule*       module,
    DiagnosticSink* sink);

void legalizeEmptyTypes(
    IRModule* module,
    DiagnosticSink* sink);

bool isResourceType(IRType* type);


}

#endif