From 071f1b6062b459928ebfd6f2f60a8d6ad021112b Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 14 Aug 2024 18:41:48 -0700
Subject: Variadic Generics Part 1: parsing and type checking. (#4833)

---
 source/slang/slang-ast-builder.cpp | 95 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)

(limited to 'source/slang/slang-ast-builder.cpp')

diff --git a/source/slang/slang-ast-builder.cpp b/source/slang/slang-ast-builder.cpp
index ce4c32c3a..faf15470f 100644
--- a/source/slang/slang-ast-builder.cpp
+++ b/source/slang/slang-ast-builder.cpp
@@ -533,12 +533,65 @@ TypeType* ASTBuilder::getTypeType(Type* type)
     return getOrCreate<TypeType>(type);
 }
 
+Type* ASTBuilder::getEachType(Type* baseType)
+{
+    // each expand T ==> T
+    if (auto expandType = as<ExpandType>(baseType))
+    {
+        return expandType->getPatternType();
+    }
+
+    // each Tuple<X> ==> each X, because we know that Tuple type must be in the form of Tuple<ConcreteTypePack<...>>.
+    if (auto tupleType = as<TupleType>(baseType))
+    {
+        return getEachType(tupleType->getTypePack());
+    }
+    SLANG_ASSERT(!as<EachType>(baseType));
+    return getOrCreate<EachType>(baseType);
+}
+
+Type* ASTBuilder::getExpandType(Type* pattern, ArrayView<Type*> capturedPacks)
+{
+    // expand each T ==> T
+    if (auto eachType = as<EachType>(pattern))
+    {
+        return eachType->getElementType();
+    }
+    return getOrCreate<ExpandType>(pattern, capturedPacks);
+}
+
+ConcreteTypePack* ASTBuilder::getTypePack(ArrayView<Type*> types)
+{
+    return getOrCreate<ConcreteTypePack>(types);
+}
+
 TypeEqualityWitness* ASTBuilder::getTypeEqualityWitness(
     Type* type)
 {
     return getOrCreate<TypeEqualityWitness>(type, type);
 }
 
+TypePackSubtypeWitness* ASTBuilder::getSubtypeWitnessPack(
+    Type* subType, Type* superType, ArrayView<SubtypeWitness*> witnesses)
+{
+    return getOrCreate<TypePackSubtypeWitness>(subType, superType, witnesses);
+}
+
+SubtypeWitness* ASTBuilder::getExpandSubtypeWitness(
+    Type* subType, Type* superType, SubtypeWitness* patternWitness)
+{
+    if (auto eachWitness = as<EachSubtypeWitness>(patternWitness))
+        return eachWitness->getPatternTypeWitness();
+    return getOrCreate<ExpandSubtypeWitness>(subType, superType, patternWitness);
+}
+
+SubtypeWitness* ASTBuilder::getEachSubtypeWitness(
+    Type* subType, Type* superType, SubtypeWitness* patternWitness)
+{
+    if (auto expandWitness = as<ExpandSubtypeWitness>(patternWitness))
+        return expandWitness->getPatternTypeWitness();
+    return getOrCreate<EachSubtypeWitness>(subType, superType, patternWitness);
+}
 
 DeclaredSubtypeWitness* ASTBuilder::getDeclaredSubtypeWitness(
     Type*                   subType,
@@ -651,6 +704,48 @@ top:
             indexOfCInConjunction);
     }
 
+    // If left hand is a TypePackSubtypeWitness, then we should also return a TypePackSubtypeWitness
+    // where each witness in the pack is the transitive subtype witness of the corresponding
+    // witness in the original pack.
+    //
+    if (auto witnessPack = as<TypePackSubtypeWitness>(aIsSubtypeOfBWitness))
+    {
+        List<SubtypeWitness*> newWitnesses;
+        for (Index i = 0; i < witnessPack->getCount(); i++)
+        {
+            newWitnesses.add(getTransitiveSubtypeWitness(witnessPack->getWitness(i), bIsSubtypeOfCWitness));
+        }
+        return getSubtypeWitnessPack(
+            aType,
+            cType,
+            newWitnesses.getArrayView());
+    }
+
+    // If left hand is a ExpandSubtypeWitness, then we want to perform the transitive lookup
+    // on the pattern witness, and then form a new ExpandSubtypeWitness with the result.
+    //
+    if (auto expandWitness = as<ExpandSubtypeWitness>(aIsSubtypeOfBWitness))
+    {
+        auto innerTransitiveWitness = getTransitiveSubtypeWitness(expandWitness->getPatternTypeWitness(), bIsSubtypeOfCWitness);
+        return getExpandSubtypeWitness(expandWitness->getSub(), cType, innerTransitiveWitness);
+    }
+
+    // If left hand is a DeclaredWitness for a type pack parameter T, then we want to perform the
+    // transitive lookup on `each T`, and then form a new ExpandSubtypeWitness with the result.
+    //
+    if (auto declaredWitness = as<DeclaredSubtypeWitness>(aIsSubtypeOfBWitness))
+    {
+        if (auto declRefType = as<DeclRefType>(declaredWitness->getSub()))
+        {
+            if (declRefType->getDeclRef().as<GenericTypePackParamDecl>())
+            {
+                auto newLeftHandWitness = getEachSubtypeWitness(getEachType(declaredWitness->getSub()), declaredWitness->getSup(), declaredWitness);
+                auto transitiveWitness = getTransitiveSubtypeWitness(newLeftHandWitness, bIsSubtypeOfCWitness);
+                return getExpandSubtypeWitness(aType, cType, transitiveWitness);
+            }
+        }
+    }
+
     // If none of the above special cases applied, then we are just going to create
     // a `TransitiveSubtypeWitness` directly.
     //
-- 
cgit v1.2.3