From 63e0064bd3a2007adf17a35d3c58894d90ddc04a Mon Sep 17 00:00:00 2001
From: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com>
Date: Tue, 25 Jun 2024 20:07:55 -0400
Subject: Fix return type address space checking. (#4465)

solves: `func-resource-result-simple.slang`, related to #4291

when getting the address space of a function return we now check the return address of the inst if the return type is a non pointer type.
---
 source/slang/slang-ir-specialize-address-space.cpp | 51 +++++++++++-----------
 1 file changed, 25 insertions(+), 26 deletions(-)

(limited to 'source')

diff --git a/source/slang/slang-ir-specialize-address-space.cpp b/source/slang/slang-ir-specialize-address-space.cpp
index 5a1874e08..b16449f2e 100644
--- a/source/slang/slang-ir-specialize-address-space.cpp
+++ b/source/slang/slang-ir-specialize-address-space.cpp
@@ -18,24 +18,8 @@ namespace Slang
         {
         }
 
-        AddressSpace getLeafInstAddressSpace(IRInst* inst)
+        AddressSpace getAddressSpaceFromVarType(IRInst* type)
         {
-            if (as<IRGroupSharedRate>(inst->getRate()))
-                return AddressSpace::GroupShared;
-            switch (inst->getOp())
-            {
-            case kIROp_RWStructuredBufferGetElementPtr:
-                return AddressSpace::Global;
-            case kIROp_Var:
-                if (as<IRBlock>(inst->getParent()))
-                    return AddressSpace::ThreadLocal;
-                break;
-            default:
-                break;
-            }
-            auto type = unwrapAttributedType(inst->getDataType());
-            if (!type)
-                return AddressSpace::Generic;
             if (as<IRUniformParameterGroupType>(type))
             {
                 return AddressSpace::Uniform;
@@ -61,6 +45,27 @@ namespace Slang
             return AddressSpace::Generic;
         }
 
+        AddressSpace getLeafInstAddressSpace(IRInst* inst)
+        {
+            if (as<IRGroupSharedRate>(inst->getRate()))
+                return AddressSpace::GroupShared;
+            switch (inst->getOp())
+            {
+            case kIROp_RWStructuredBufferGetElementPtr:
+                return AddressSpace::Global;
+            case kIROp_Var:
+                if (as<IRBlock>(inst->getParent()))
+                    return AddressSpace::ThreadLocal;
+                break;
+            default:
+                break;
+            }
+            auto type = unwrapAttributedType(inst->getDataType());
+            if (!type)
+                return AddressSpace::Generic;
+            return getAddressSpaceFromVarType(type);
+        }
+
         AddressSpace getAddrSpace(IRInst* inst)
         {
             auto addrSpace = mapInstToAddrSpace.tryGetValue(inst);
@@ -154,13 +159,7 @@ namespace Slang
         AddressSpace getFuncResultAddrSpace(IRFunc* callee)
         {
             auto funcType = as<IRFuncType>(callee->getDataType());
-            auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
-            if (!ptrResultType)
-                return AddressSpace::Generic;
-            AddressSpace resultAddrSpace = AddressSpace::Generic;
-            if (ptrResultType->hasAddressSpace())
-                resultAddrSpace = (AddressSpace)ptrResultType->getAddressSpace();
-            return resultAddrSpace;
+            return getAddressSpaceFromVarType(funcType->getResultType());
         }
 
         // Return true if the address space of the function return type is changed.
@@ -316,11 +315,11 @@ namespace Slang
                                 if (addrSpace != AddressSpace::Generic)
                                 {
                                     auto funcType = as<IRFuncType>(func->getDataType());
-                                    auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
-                                    SLANG_ASSERT(ptrResultType);
                                     AddressSpace resultAddrSpace = getFuncResultAddrSpace(func);
                                     if (resultAddrSpace != addrSpace)
                                     {
+                                        auto ptrResultType = as<IRPtrTypeBase>(funcType->getResultType());
+                                        SLANG_ASSERT(ptrResultType);
                                         IRBuilder builder(func);
                                         auto newResultType = builder.getPtrType(ptrResultType->getOp(), ptrResultType->getValueType(), addrSpace);
                                         fixUpFuncType(func, newResultType);
-- 
cgit v1.2.3