From 6f6103c4dbc77d5bceae7c8e766ec3cabc293364 Mon Sep 17 00:00:00 2001
From: Julius Ikkala <julius.ikkala@gmail.com>
Date: Sat, 3 May 2025 23:27:03 +0300
Subject: Add IREnumType to distinguish enums from ints and each other (#6973)

* Add IREnumType to distinguish enums from ints and each other

* Add issue example as test

* format code

* Add expected test output

* Fix peephole optimization hanging

No idea why this PR triggered this, but there seems to have been a clear bug
here anyway, so may just as well fix it now.

* Move enum lowering later

* Add linkage decoration to enum type

* Use filecheck-buffer instead of expected.txt

* Fix comment

* Make enum casts actually use IR enum casts

They were all BuiltinCasts by accident

* Lower enum type before VM

* Deal with rate-qualified types in enum cast

* Allow any value marshalling for enum types

* Handle new enum instructions in a couple more switches

* Fix formatting

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
---
 source/slang/slang-ir-lower-enum-type.cpp | 149 ++++++++++++++++++++++++++++++
 1 file changed, 149 insertions(+)
 create mode 100644 source/slang/slang-ir-lower-enum-type.cpp

(limited to 'source/slang/slang-ir-lower-enum-type.cpp')

diff --git a/source/slang/slang-ir-lower-enum-type.cpp b/source/slang/slang-ir-lower-enum-type.cpp
new file mode 100644
index 000000000..548c29a51
--- /dev/null
+++ b/source/slang/slang-ir-lower-enum-type.cpp
@@ -0,0 +1,149 @@
+// slang-ir-lower-enum-type.cpp
+
+#include "slang-ir-lower-enum-type.h"
+
+#include "slang-ir-insts.h"
+#include "slang-ir-util.h"
+#include "slang-ir.h"
+
+namespace Slang
+{
+struct EnumTypeLoweringContext
+{
+    IRModule* module;
+    DiagnosticSink* sink;
+
+    InstWorkList workList;
+    InstHashSet workListSet;
+
+    IRGeneric* genericOptionalStructType = nullptr;
+    IRStructKey* valueKey = nullptr;
+    IRStructKey* hasValueKey = nullptr;
+
+    EnumTypeLoweringContext(IRModule* inModule)
+        : module(inModule), workList(inModule), workListSet(inModule)
+    {
+    }
+
+    struct LoweredEnumTypeInfo : public RefObject
+    {
+        IRType* enumType = nullptr;
+        IRType* loweredType = nullptr;
+    };
+    Dictionary<IRInst*, RefPtr<LoweredEnumTypeInfo>> loweredEnumTypes;
+
+    void addToWorkList(IRInst* inst)
+    {
+        if (workListSet.contains(inst))
+            return;
+
+        workList.add(inst);
+        workListSet.add(inst);
+    }
+
+    LoweredEnumTypeInfo* getLoweredEnumType(IRInst* type)
+    {
+        if (auto loweredInfo = loweredEnumTypes.tryGetValue(type))
+            return loweredInfo->Ptr();
+
+        if (!type)
+            return nullptr;
+
+        if (type->getOp() != kIROp_EnumType)
+            return nullptr;
+
+        RefPtr<LoweredEnumTypeInfo> info = new LoweredEnumTypeInfo();
+        auto enumType = cast<IREnumType>(type);
+        auto valueType = enumType->getTagType();
+        info->enumType = (IRType*)type;
+        info->loweredType = valueType;
+        loweredEnumTypes[type] = info;
+        return info.Ptr();
+    }
+
+    void processEnumType(IREnumType* inst)
+    {
+        auto loweredEnumTypeInfo = getLoweredEnumType(inst);
+        SLANG_ASSERT(loweredEnumTypeInfo);
+        SLANG_UNUSED(loweredEnumTypeInfo);
+    }
+
+    void processEnumCast(IRInst* inst)
+    {
+        IRBuilder builderStorage(module);
+        auto builder = &builderStorage;
+        builder->setInsertBefore(inst);
+
+        auto value = inst->getOperand(0);
+        if (auto enumType = getLoweredEnumType(value->getDataType()))
+        {
+            auto rate = value->getRate();
+            auto type = enumType->loweredType;
+            if (rate)
+            {
+                type = builder->getRateQualifiedType(rate, type);
+            }
+
+            value->setFullType(type);
+        }
+
+        auto type = inst->getDataType();
+        if (auto enumType = getLoweredEnumType(type))
+        { // Cast was into enum, so use tag type instead.
+            type = enumType->loweredType;
+        }
+
+        auto cast = builder->emitCast(type, value);
+
+        inst->replaceUsesWith(cast);
+        inst->removeAndDeallocate();
+    }
+
+    void processInst(IRInst* inst)
+    {
+        switch (inst->getOp())
+        {
+        case kIROp_EnumType:
+            processEnumType((IREnumType*)inst);
+            break;
+        case kIROp_CastEnumToInt:
+        case kIROp_CastIntToEnum:
+        case kIROp_EnumCast:
+            processEnumCast(inst);
+            break;
+        default:
+            break;
+        }
+    }
+
+    void processModule()
+    {
+        addToWorkList(module->getModuleInst());
+
+        while (workList.getCount() != 0)
+        {
+            IRInst* inst = workList.getLast();
+            workList.removeLast();
+            workListSet.remove(inst);
+
+            processInst(inst);
+
+            for (auto child = inst->getLastChild(); child; child = child->getPrevInst())
+            {
+                addToWorkList(child);
+            }
+        }
+
+        // Replace all enum types with their lowered equivalent types.
+        for (const auto& [key, value] : loweredEnumTypes)
+            key->replaceUsesWith(value->loweredType);
+    }
+};
+
+void lowerEnumType(IRModule* module, DiagnosticSink* sink)
+{
+    EnumTypeLoweringContext context(module);
+    context.sink = sink;
+    context.processModule();
+}
+} // namespace Slang
-- 
cgit v1.2.3