Canonicalize generic app member decl refs

Author: csyonghe

source/slang/slang-ast-builder.h

@@ -344,6 +344,21 @@ class ASTBuilder : public RefObject
         {
             return DeclRef<T>(getMemberDeclRef(parentMemberDeclRef->getParent(), memberDecl));
         }
+        else if (auto genericAppDeclRef = as<GenericAppDeclRef>(parent.declRefBase))
+        {
+            // For a generic `G<T>`, a direct member declared under `G` should be represented as
+            // `GenericAppDeclRef(G, T, member)` instead of
+            // `MemberDeclRef(GenericAppDeclRef(G, T, G::Inner), member)`. That keeps decl refs
+            // canonical so generic Val rewriting can remain a simple operand walk.
+            if (memberDecl->parentDecl == genericAppDeclRef->getGenericDecl())
+            {
+                return getGenericAppDeclRef(
+                           DeclRef<GenericDecl>(genericAppDeclRef->getGenericDeclRef()),
+                           genericAppDeclRef->getArgs(),
+                           memberDecl)
+                    .template as<T>();
+            }
+        }
         else if (auto lookupDeclRef = as<LookupDeclRef>(parent.declRefBase))
         {
             // Handle some specicial case rules due to the way some of our builtin decls are

source/slang/slang-ast-copy.cpp

@@ -146,103 +146,6 @@ Decl* ASTCopier::rewriteDecl(Decl* decl)
 
 Val* ASTCopier::rewriteValImpl(Val* val)
 {
-    if (auto memberDeclRef = as<MemberDeclRef>(val))
-    {
-        auto oldMemberDecl = memberDeclRef->getDecl();
-        auto newMemberDecl = rewriteDecl(oldMemberDecl);
-        if (newMemberDecl != oldMemberDecl)
-        {
-            // `MemberDeclRef` stores a substitution context as its parent operand. Checked type
-            // expressions for nested generic constraints can use an enclosing specialized type,
-            // e.g. `InlineVector<T, N>`, as that context while the constraint decl itself is being
-            // copied into a standalone generic. Rebuild through `getParent()` so the copied member
-            // is paired with the nearest rewritten declaration parent instead of a source-side
-            // context that only belonged to the checked source graph.
-            auto newParentDeclRef = rewriteDeclRef(memberDeclRef->getParent());
-            if (newParentDeclRef && newMemberDecl->isChildOf(newParentDeclRef->getDecl()))
-                return m_context.astBuilder->getMemberDeclRef(newParentDeclRef, newMemberDecl)
-                    .declRefBase;
-        }
-    }
-
-    if (auto genericAppDeclRef = as<GenericAppDeclRef>(val))
-    {
-        auto oldGenericDeclRef = genericAppDeclRef->getGenericDeclRef();
-        auto newGenericDeclRef = rewriteDeclRef(oldGenericDeclRef);
-
-        List<Val*> newArgs;
-        bool changed = newGenericDeclRef != oldGenericDeclRef;
-        for (auto arg : genericAppDeclRef->getArgs())
-        {
-            auto newArg = rewriteVal(arg);
-            changed = changed || newArg != arg;
-            newArgs.add(newArg);
-        }
-
-        auto oldInnerDecl = genericAppDeclRef->getDecl();
-        auto newInnerDecl = rewriteDecl(oldInnerDecl);
-        changed = changed || newInnerDecl != oldInnerDecl;
-
-        if (!changed)
-            return val;
-
-        if (!newGenericDeclRef || !newInnerDecl)
-        {
-            SLANG_UNEXPECTED("ASTCopier could not coherently rewrite a GenericAppDeclRef.");
-        }
-
-        if (!newInnerDecl->isChildOf(newGenericDeclRef->getDecl()))
-        {
-            auto oldGenericDecl =
-                oldGenericDeclRef ? as<GenericDecl>(oldGenericDeclRef->getDecl()) : nullptr;
-            if (oldGenericDecl && oldInnerDecl == oldGenericDecl->inner)
-            {
-                // `ITensor<Dimension>` inside an interface default method applies the original
-                // generic interface type constructor to an argument from the copied generic
-                // environment. The generic head is not itself being copied here; only the argument
-                // values need rewriting.
-                newGenericDeclRef = oldGenericDeclRef;
-                newInnerDecl = oldInnerDecl;
-            }
-            else if (auto oldGenericLookup = as<LookupDeclRef>(oldGenericDeclRef))
-            {
-                // Re-resolve looked-up generic members through the rewritten lookup source/witness.
-                // Interface default methods can reference generic requirements through a
-                // LookupDeclRef; after copying the enclosing signature, the old looked-up inner
-                // declaration is no longer a child of the rewritten generic parent, so resolve the
-                // corresponding lookup in the copied environment before rebuilding the
-                // GenericAppDeclRef.
-                auto innerLookupDeclRef = m_context.astBuilder->getLookupDeclRef(
-                    oldGenericLookup->getLookupSource(),
-                    oldGenericLookup->getWitness(),
-                    oldInnerDecl);
-                auto rewrittenInnerLookup = rewriteDeclRef(innerLookupDeclRef);
-                if (rewrittenInnerLookup.declRefBase)
-                {
-                    if (auto resolvedInnerDeclRef =
-                            as<DeclRefBase>(rewrittenInnerLookup.declRefBase->resolve()))
-                    {
-                        newInnerDecl = resolvedInnerDeclRef->getDecl();
-                    }
-                    else
-                    {
-                        newInnerDecl = rewrittenInnerLookup.getDecl();
-                    }
-                }
-            }
-        }
-
-        if (!newInnerDecl->isChildOf(newGenericDeclRef->getDecl()))
-        {
-            SLANG_UNEXPECTED("ASTCopier could not coherently rewrite a GenericAppDeclRef.");
-        }
-
-        return m_context.astBuilder->getGenericAppDeclRef(
-            DeclRef<GenericDecl>(newGenericDeclRef),
-            newArgs.getArrayView(),
-            newInnerDecl);
-    }
-
     ValNodeDesc desc;
     desc.type = val->getClass();
 
@@ -386,7 +289,6 @@ GenericSignatureCopier::GenericSignatureCopier(
     , m_destGenericDecl(destGenericDecl)
     , m_outGenericArgs(outGenericArgs)
 {
-    m_astCopier.mapDecl(sourceGenericDecl, destGenericDecl);
     m_astCopier.getContext().allowDroppingWitnessTables = true;
 }
 

source/slang/slang-ast-decl-ref.cpp

@@ -477,6 +477,13 @@ DeclRefBase* GenericAppDeclRef::_substituteImplOverride(
             substGenericDeclRef,
             substArgs.getArrayView(),
             getDecl());
+    else if (getDecl() == getGenericDecl()->inner)
+    {
+        // If substitution rewrites the generic declaration reference itself, the old primary
+        // inner declaration is no longer owned by the substituted generic. Use the substituted
+        // generic's primary inner declaration, matching `_resolveImplOverride`.
+        return astBuilder->getGenericAppDeclRef(substGenericDeclRef, substArgs.getArrayView());
+    }
     else
     {
         // If decl is no longer the child of the new parent, it's most likely due to