// slang-ir-wrap-structured-buffers.cpp
#include "slang-ir-wrap-structured-buffers.h"

#include "slang-ir.h"
#include "slang-ir-insts.h"

namespace Slang
{

// Suppose a programmer wrote HLSL/Slang code like this:
//
//      StructuredBuffer<float4x4> gBuffer;
//
//      ... gBuffer[index] ...
//      ... gBuffer.Load(index) ...
//
// Further suppose that they specified that row-major
// matrix layout should be used. It might surprise such
// a user that the fxc and dxc compilers both use
// column-major layout for the matrices in `gBuffer`,
// and there is no way to change that fact.
//
// In contrast, we want Slang to respect the matrix layout
// setting given by the user in all cases, so we need to
// find a way to force fxc and dxc to do what we want.
//
// Fortunately, fxc and dxc *do* respect the row-major
// layout setting for code like the following:
//
//      struct Wrapper { float4x4 wrapped; }
//      StructuredBuffer<Wrapper> gBuffer;
//
//      ... gBuffer[index].wrapped ...
//      ... gBuffer.Load(index).wrapped ...
//
// The role of the pass in this file is to transform IR
// for code like the first case into IR for code like the
// second case.
//
// The main thing that makes this pass challenging is
// recognizing calls to `StructuredBuffer<T>` builtins
// so that we can transform them to append the `.wrapped`
// field reference.

// As is typical, we will wrap up our IR pass in a
// "context" structure type.
//
struct WrapStructuredBuffersContext
{
    IRModule* m_module = nullptr;

    // We process a module by processing all its instructions, recursively.
    //
    void processModule()
    {
        processInstRec(m_module->getModuleInst());
    }

    void processInstRec(IRInst* inst)
    {
        processInst(inst);

        for(auto child : inst->getChildren())
            processInstRec(child);
    }

    void processInst(IRInst* inst)
    {
        // We will frame our processing as a search for types
        // of the form `*StructuredBuffer<matrixType>` where
        // `matrixType` is any matrix type.
        //
        // If the instruction we are looking at doesn't have
        // the right form, then we will skip it.
        //
        auto oldStructuredBufferType = as<IRHLSLStructuredBufferTypeBase>(inst);
        if(!oldStructuredBufferType)
            return;
        auto matrixType = as<IRMatrixType>(oldStructuredBufferType->getElementType());
        if(!matrixType)
            return;

        // Having found a `*StructuredBuffer<M>` we will now
        // need an IR builder to help us construct the wrapper code.
        //
        IRBuilder builderStorage(m_module);
        auto builder = &builderStorage;

        // We begin by constructing a structure type that wraps
        // our `matrixType`, into something like:
        //
        //      struct Wrapper { matrixType wrapped; }
        //
        builder->setInsertBefore(oldStructuredBufferType);
        auto wrappedFieldKey = getWrappedFieldKey(builder);
        auto wrapperStruct = getWrapperStruct(builder, matrixType);

        // Now that we have our wrapper struct, we can construct a type
        // that is `*StructuredBuffer<wrapperStruct>` and use it to
        // replace `*StructuredBuffer<matrixType>`
        //
        // Note: we are constructing a *new* type instead of modifying
        // the old one in-place because eventually when we do deduplication
        // of types/constants more consistently it might cause problems
        // to modify a tyep in a way that changes its hash code.
        //
        // TODO: the above statement is still slippery, since we are still
        // replacing one type A with another type B globally, and doing
        // so could affect any type that in turn referenced A...
        //
        auto newStructuredBufferType = builder->getType(oldStructuredBufferType->getOp(), wrapperStruct);
        oldStructuredBufferType->replaceUsesWith(newStructuredBufferType);

        // Any values that used our old structured bufer type
        // now have the new structured buffer type, but that
        // means that operations on them might be getting
        // applied incorrectly.
        //
        // For example, if we had a call like:
        //
        //      float4x4 m = gBuffer.Load(someIndex);
        //
        // the result of the `Load` call is now `wrapperStruct`
        // and we can't assign that to a matrix-type variable.
        //
        // we need to invetigate values of our structured
        // buffer type, and then investigate operations that
        // are using those values, to see if we can find the
        // ones we need to rewrite.
        //
        // We can find values of `newStructuredBufferType` by
        // scanning through its IR uses, since values of that
        // type are using it as a (type) operand.
        //
        traverseUses(newStructuredBufferType, [&](IRUse* typeUse)
        {
            // There might be uses of `newStructuredBufferType` where
            // it isn't being used as the type of a value, so we
            // start by weeding out the ones we don't care about.
            //
            auto valueOfStructuredBufferType = typeUse->getUser();
            if(valueOfStructuredBufferType->getFullType() != newStructuredBufferType)
                return;

            // Now we have some `valueOfStructuredBufferType`. In our running
            // example, this might be `gBuffer`, which is an `IRGlobalParam`.
            //
            // We don't need to change anything about `gBuffer` itself, since
            // replacing `oldStructuredBufferType` with `newStructuredBufferType`
            // already replaced the type of `gBuffer`.
            //
            // Instead, we want to look for instructions that *use* the buffer,
            // because these could be calls to intrinsic functions like
            // `RWStructuredBuffer.Load`
            //
            traverseUses(valueOfStructuredBufferType, [&](IRUse* valueUse)
            {
                // we are only interested in instructions that are calls,
                // with at least one argument, where the first argument
                // is our `valueOfStructuredBufferType`. These
                // are calls that could potentially be intrinsic
                // operations on `*StructuredBuffer`.
                //
                auto user = valueUse->getUser();
                switch (user->getOp())
                {
                case kIROp_StructuredBufferLoad:
                case kIROp_StructuredBufferLoadStatus:
                case kIROp_RWStructuredBufferStore:
                case kIROp_RWStructuredBufferLoadStatus:
                case kIROp_RWStructuredBufferGetElementPtr:
                    break;
                default:
                    return;
                }

                builder->setInsertAfter(user);
                auto oldResultType = user->getDataType();

                // First we care about the case for `Load`, which
                // will return the element type, which would be
                // a matrix type.
                //
                if( as<IRMatrixType>(oldResultType) )
                {
                    // We know that the call to `Load` should now
                    // return our wrapper struct type, so we will
                    // go ahead and modify its type to be correct.
                    //
                    auto newResultType = wrapperStruct;
                    builder->setDataType(user, newResultType);

                    // Next, we need to make sure to extract the
                    // field from the wrapper struct, so that
                    // we get back to a value of the expected
                    // type.
                    //
                    // This logic takes something like:
                    //
                    //      WrapperStruct call = gBuffer.Load(index);
                    //
                    // and follows it with:
                    //
                    //      float4x4 newVal = call.wrapped;
                    //
                    auto newVal = builder->emitFieldExtract(oldResultType, user, wrappedFieldKey);

                    // Any code that used the value of `call` should
                    // now use `newVal` instead...
                    //
                    user->replaceUsesWith(newVal);
                    //
                    // ... except for one important gotcha, which is
                    // that `newVal` itself used `call`, and replacing
                    // `call` with `newVal` results in `newVal` using
                    // itself as one of its operands.
                    //
                    // It is a bit of a kludge, but we fix the situation
                    // by just setting the appropriate operand again.
                    //
                    // TODO: it might be helpful to have a variant
                    // of `replaceUsesWith` that can handle cases like
                    // this.
                    //
                    newVal->setOperand(0, user);
                }
                //
                // The second interesting case is the `ref` accessor
                // in `operator[]` for a `RWStructuredBuffer`, which
                // at the IR level returns a *pointer* to the buffer
                // element type.
                //
                else if(auto oldPtrType = as<IRPtrTypeBase>(oldResultType))
                {
                    auto pointeeType = oldPtrType->getValueType();
                    if( as<IRMatrixType>(pointeeType) )
                    {
                        // At this point we know that the intrinsic
                        // operation returned a pointer to a matrix,
                        // which seems like a good indications that
                        // it is our `operator[]` and it should now
                        // return a pointer to the wrapper struct
                        // instead.
                        //
                        // The logic here is almost identical to the
                        // non-pointer case above, so please refer
                        // there if you want the comments.

                        auto newResultType = builder->getPtrType(oldPtrType->getOp(), wrapperStruct);
                        builder->setDataType(user, newResultType);

                        auto newVal = builder->emitFieldAddress(oldResultType, user, wrappedFieldKey);
                        user->replaceUsesWith(newVal);
                        newVal->setOperand(0, user);
                    }
                }
            });
        });
    }

        /// Get the struture field "key" to use for generated wrappers
    IRStructKey* getWrappedFieldKey(IRBuilder* builder)
    {
        // We will re-use the same field key for all of the
        // wrapper structs we might create, so that all of
        // the new field access operations will use the same
        // field name to make their purpose more clear.
        //
        // TODO: It might be useful to give the field key
        // a name that is indicative of its purpose so
        // that people won't be confused why their code
        // has been transformed and now there is this
        // `._S2` in the middle of their expressions.

        if( !m_wrappedFieldKey )
        {
            m_wrappedFieldKey = builder->createStructKey();
        }
        return m_wrappedFieldKey;
    }

        /// Lazily created and cached field "key" to use for wrapper structs.
    IRStructKey* m_wrappedFieldKey = nullptr;

        /// Get the wrapper struct to use for a particular `matrixType`.
    IRStructType* getWrapperStruct(IRBuilder* builder, IRMatrixType* matrixType)
    {
        // TODO: Because our type de-duplication isn't perfect right now,
        // it is possible that there could be more than one equivalent
        // matrix type, and thus more than one equivalent structured buffer
        // type using that matrix type. As a result, if we just generate
        // one new `struct` per IR structured-buffer type, we could end up
        // with conflicts when a buffer with one IR type gets passed to
        // a function with another (equivalent) IR type.
        //
        // The right fix here is to cache and look up these structure
        // tyeps based on the scalar type (opcode) and row/column counts
        // of the given `matrixType`.
        //
        // For now I will ignore this issue and hope that we can address
        // the more fundamental issue of type deduplication before this
        // choice comes back to bite me.

        auto wrappedFieldKey = getWrappedFieldKey(builder);

        auto wrapperStruct = builder->createStructType();
        builder->createStructField(wrapperStruct, wrappedFieldKey, matrixType);
        return wrapperStruct;
    }
};

void wrapStructuredBuffersOfMatrices(
    IRModule*                           module)
{
    WrapStructuredBuffersContext context;
    context.m_module = module;
    context.processModule();
}

}