Refactor lower-generics pass into separate subpasses. (#1442)

1 files changed, 193 insertions, 0 deletions

diff --git a/source/slang/slang-ir-lower-generic-var.cpp b/source/slang/slang-ir-lower-generic-var.cpp
new file mode 100644
index 000000000..0c45e8e38
--- /dev/null
+++ b/source/slang/slang-ir-lower-generic-var.cpp
@@ -0,0 +1,193 @@
+// slang-ir-lower-generic-function.cpp
+#include "slang-ir-lower-generic-function.h"
+
+#include "slang-ir-generics-lowering-context.h"
+#include "slang-ir.h"
+#include "slang-ir-clone.h"
+#include "slang-ir-insts.h"
+
+namespace Slang
+{
+    // This is a subpass of generics lowering IR transformation.
+    // This pass lowers all generic function types and function definitions, including
+    // the function types used in interface types, to ordinary functions that takes
+    // raw pointers in place of generic types.
+    struct GenericVarLoweringContext
+    {
+        SharedGenericsLoweringContext* sharedContext;
+
+        void processVarInst(IRInst* varInst)
+        {
+            // We process only var declarations that have type
+            // `Ptr<IRParam>`.
+            // Due to the processing of `lowerGenericFunction`,
+            // A local variable of generic type now appears as
+            // `var X:Ptr<irParam:Ptr<RTTIType>>`
+            // We match this pattern and turn this inst into
+            // `X:RawPtr = alloca(rtti_extract_size(irParam))`
+            auto varTypeInst = varInst->getDataType();
+            if (!varTypeInst)
+                return;
+            auto ptrType = as<IRPtrType>(varTypeInst);
+            if (!ptrType)
+                return;
+
+            // `varTypeParam` represents a pointer to the RTTI object.
+            auto varTypeParam = ptrType->getValueType();
+            if (varTypeParam->op != kIROp_Param)
+                return;
+            if (!varTypeParam->getDataType())
+                return;
+            if (varTypeParam->getDataType()->op != kIROp_PtrType)
+                return;
+            if (as<IRPtrType>(varTypeParam->getDataType())->getValueType()->op != kIROp_RTTIType)
+                return;
+
+
+            // A local variable of generic type has a type that is an IRParam.
+            // This parameter represents the RTTI that tells us the size of the type.
+            // We need to transform the variable into an `alloca` call to allocate its
+            // space based on the provided RTTI object.
+
+            // Initialize IRBuilder for emitting instructions.
+            IRBuilder builderStorage;
+            auto builder = &builderStorage;
+            builder->sharedBuilder = &sharedContext->sharedBuilderStorage;
+            builder->setInsertBefore(varInst);
+
+            // The result of `alloca` is an RTTIPointer(rttiObject).
+            auto type = builder->getRTTIPointerType(varTypeParam);
+            auto newVarInst = builder->emitAlloca(type, varTypeParam);
+            varInst->replaceUsesWith(newVarInst);
+            varInst->removeAndDeallocate();
+        }
+
+        void processStoreInst(IRStore* storeInst)
+        {
+            auto rttiType = as<IRRTTIPointerType>(storeInst->ptr.get()->getDataType());
+            if (!rttiType)
+                return;
+            // All stores of generic typed variables needs to be translated
+            // to `IRCopy`s.
+            auto valPtr = storeInst->val.get();
+            if (valPtr->getDataType()->op == kIROp_RTTIPointerType)
+            {
+                // If `value` of the store is from another generic variable, it should
+                // have already been replaced with the pointer to that variable by now.
+                // So we don't need to do anything here.
+            }
+            else
+            {
+                // If value does not come from another generic variable, then it must be
+                // a param. In this case, the parameter is a bitCast of the parameter to an
+                // RTTIPointer type, so we just use the original parameter pointer and get
+                // rid of the bitcast.
+                SLANG_ASSERT(valPtr->op == kIROp_BitCast);
+                valPtr = valPtr->getOperand(0);
+                SLANG_ASSERT(valPtr->op == kIROp_Param);
+            }
+            IRBuilder builderStorage;
+            auto builder = &builderStorage;
+            builder->sharedBuilder = &sharedContext->sharedBuilderStorage;
+            builder->setInsertBefore(storeInst);
+            auto copy = builder->emitCopy(
+                storeInst->ptr.get(),
+                valPtr,
+                rttiType->getRTTIOperand());
+            storeInst->replaceUsesWith(copy);
+            storeInst->removeAndDeallocate();
+        }
+
+        void processLoadInst(IRLoad* loadInst)
+        {
+            auto rttiType = as<IRRTTIPointerType>(loadInst->ptr.get()->getDataType());
+            if (!rttiType)
+                return;
+            // There are only two possible uses of a load(genericVar):
+            // 1. store(x, load(genVar)), which will be handled by processStoreInst.
+            // 2. call(f, load(genVar)) when calling a generic function or a member function
+            //    via an interface witness lookup. In this case, we need to replace with
+            //    just `genVar`, since that function has already been lowered to take
+            //    raw pointers.
+            // In both cases, we can simply replace the use side with a pointer instead
+            // and never need to represent a "value" typed object explicitly.
+            // However, to preserve the ordering, we must make a copy of every load so
+            // we don't change the meaning of the code if there are `store`s between the
+            // `load` and the use site.
+
+            IRBuilder builderStorage;
+            auto builder = &builderStorage;
+            builder->sharedBuilder = &sharedContext->sharedBuilderStorage;
+            builder->setInsertBefore(loadInst);
+
+            // Allocate a copy of the value.
+            auto allocaInst = builder->emitAlloca(rttiType, rttiType->getRTTIOperand());
+            builder->emitCopy(allocaInst, loadInst->ptr.get(), rttiType->getRTTIOperand());
+
+            // Here we replace all uses of load to just the pointer to the copy.
+            // After this, all arguments in `call`s will be in its correct form.
+            // All `store`s will become `store(x, genVar)`, and still need
+            // to be translated into another `copy`, we leave that step when we get to
+            // process the `store` inst.
+            loadInst->replaceUsesWith(allocaInst);
+            loadInst->removeAndDeallocate();
+        }
+
+        void processInst(IRInst* inst)
+        {
+            if (inst->op == kIROp_Var || inst->op == kIROp_undefined)
+            {
+                processVarInst(inst);
+            }
+            else if (inst->op == kIROp_Load)
+            {
+                processLoadInst(cast<IRLoad>(inst));
+            }
+            else if (inst->op == kIROp_Store)
+            {
+                processStoreInst(cast<IRStore>(inst));
+            }
+        }
+
+        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)
+            {
+                // We will then iterate until our work list goes dry.
+                //
+                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);
+                    }
+                }
+            }
+        }
+    };
+
+    void lowerGenericVar(SharedGenericsLoweringContext* sharedContext)
+    {
+        GenericVarLoweringContext context;
+        context.sharedContext = sharedContext;
+        context.processModule();
+    }
+}
+