Fix type checking around generic array types. (#3568)

2 files changed, 38 insertions, 23 deletions

diff --git a/source/slang/slang-ast-builder.cpp b/source/slang/slang-ast-builder.cpp
index ce13ab650..0c30366e8 100644
--- a/source/slang/slang-ast-builder.cpp
+++ b/source/slang/slang-ast-builder.cpp
@@ -338,6 +338,10 @@ ArrayExpressionType* ASTBuilder::getArrayType(Type* elementType, IntVal* element
         {
             elementCount = getIntVal(getIntType(), elementCountConstantInt->getValue());
         }
+        else
+        {
+            elementCount = getTypeCastIntVal(getIntType(), elementCount);
+        }
     }
     Val* args[] = {elementType, elementCount};
     return as<ArrayExpressionType>(getSpecializedBuiltinType(makeArrayView(args), "ArrayExpressionType"));
diff --git a/source/slang/slang-ir-specialize.cpp b/source/slang/slang-ir-specialize.cpp
index 60001661c..9de574a9b 100644
--- a/source/slang/slang-ir-specialize.cpp
+++ b/source/slang/slang-ir-specialize.cpp
@@ -1323,6 +1323,7 @@ struct SpecializationContext
         // their replacements.
         //
         IRCloneEnv cloneEnv;
+        cloneEnv.squashChildrenMapping = true;
 
         // We also need some IR building state, for any
         // new instructions we will emit.
@@ -1330,6 +1331,16 @@ struct SpecializationContext
         IRBuilder builderStorage(module);
         auto builder = &builderStorage;
 
+        // To get started, we will create the skeleton of the new
+        // specialized function, so newly created insts
+        // will be placed in a proper parent.
+        //
+
+        IRFunc* newFunc = builder->createFunc();
+
+        builder->setInsertInto(newFunc);
+        IRBlock* tempHeaderBlock = builder->emitBlock();
+
         // We will start out by determining what the parameters
         // of the specialized function should be, based on
         // the parameters of the original, and the concrete
@@ -1343,7 +1354,6 @@ struct SpecializationContext
         // block, or even a function, to insert them into.
         //
         List<IRParam*> newParams;
-        List<IRInst*> newBodyInsts;
         UInt argCounter = 0;
         for (auto oldParam : oldFunc->getParams())
         {
@@ -1387,7 +1397,6 @@ struct SpecializationContext
                 // correct existential type, and stores the right witness table).
                 //
                 auto newMakeExistential = builder->emitMakeExistential(oldParam->getFullType(), newParam, witnessTable);
-                newBodyInsts.add(newMakeExistential);
                 replacementVal = newMakeExistential;
             }
             else if (auto oldWrapExistential = as<IRWrapExistential>(arg))
@@ -1412,7 +1421,6 @@ struct SpecializationContext
                     newParam,
                     oldWrapExistential->getSlotOperandCount(),
                     oldWrapExistential->getSlotOperands());
-                newBodyInsts.add(newWrapExistential);
                 replacementVal = newWrapExistential;
             }
             else
@@ -1433,8 +1441,20 @@ struct SpecializationContext
             cloneEnv.mapOldValToNew.add(oldParam, replacementVal);
         }
 
-        // Next we will create the skeleton of the new
-        // specialized function, including its type.
+        // The above steps have accomplished the "first phase"
+        // of cloning the function (since `IRFunc`s have no
+        // operands).
+        //
+        // We can now use the shared IR cloning infrastructure
+        // to perform the second phase of cloning, which will recursively
+        // clone any nested decorations, blocks, and instructions.
+        //
+        cloneInstDecorationsAndChildren(
+            &cloneEnv,
+            builder->getModule(),
+            oldFunc,
+            newFunc);
+
         //
         // In order to construct the type of the new function, we
         // need to extract the types of all its parameters.
@@ -1448,24 +1468,9 @@ struct SpecializationContext
             newParamTypes.getCount(),
             newParamTypes.getBuffer(),
             oldFunc->getResultType());
-        IRFunc* newFunc = builder->createFunc();
         newFunc->setFullType(newFuncType);
 
-        // The above steps have accomplished the "first phase"
-        // of cloning the function (since `IRFunc`s have no
-        // operands).
-        //
-        // We can now use the shared IR cloning infrastructure
-        // to perform the second phase of cloning, which will recursively
-        // clone any nested decorations, blocks, and instructions.
-        //
-        cloneInstDecorationsAndChildren(
-            &cloneEnv,
-            builder->getModule(),
-            oldFunc,
-            newFunc);
-
-        // Now that the main body of existing isntructions have
+        // Now that the main body of existing instructions have
         // been cloned into the new function, we can go ahead
         // and insert all the parameters and body instructions
         // we built up into the function at the right place.
@@ -1474,7 +1479,7 @@ struct SpecializationContext
         // block (this was an invariant established before
         // we decided to specialize).
         //
-        auto newEntryBlock = newFunc->getFirstBlock();
+        auto newEntryBlock = as<IRBlock>(cloneEnv.mapOldValToNew[oldFunc->getFirstBlock()]);
         SLANG_ASSERT(newEntryBlock);
 
         // We expect every valid block to have at least one
@@ -1497,11 +1502,17 @@ struct SpecializationContext
         // before the first ordinary instruction (but will come
         // *after* the parameters by the order of these two loops).
         //
-        for (auto newBodyInst : newBodyInsts)
+        for (auto newBodyInst = tempHeaderBlock->getFirstChild(); newBodyInst;)
         {
+            auto next = newBodyInst->next;
             newBodyInst->insertBefore(newFirstOrdinary);
+            newBodyInst = next;
         }
 
+        // After moving all param and existential insts in tempHeaderBlock
+        // it should be empty now and we can remove it.
+        tempHeaderBlock->removeAndDeallocate();
+
         // After all this work we have a valid `newFunc` that has been
         // specialized to match the types at the call site.
         //