Merge pull request swiftlang#24621 from aschwaighofer/opaque_specializer_fixes

Opaque specializer fixes

Author: aschwaighofer

lib/SILOptimizer/Transforms/SpecializeOpaqueArchetypes.cpp

@@ -133,18 +133,30 @@ class OpaqueSpecializerCloner
   void visitEnumInst(EnumInst *Inst) {
     getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
     SILValue opd = SILValue();
+    auto newTy = getOpType(Inst->getType());
     if (Inst->hasOperand()) {
-      SILType caseTy = Inst->getType().getEnumElementType(
-          Inst->getElement(), getBuilder().getFunction().getModule());
       opd = getOpValue(Inst->getOperand());
-      if (opd->getType() != caseTy) {
-        opd = createCast(getOpLocation(Inst->getLoc()), opd, caseTy);
-      }
+      SILType newCaseTy = newTy.getEnumElementType(
+          Inst->getElement(), getBuilder().getFunction().getModule());
+      if (opd->getType() != newCaseTy)
+        opd = createCast(getOpLocation(Inst->getLoc()), opd, newCaseTy);
     }
     recordClonedInstruction(
         Inst, getBuilder().createEnum(getOpLocation(Inst->getLoc()), opd,
-                                      Inst->getElement(),
-                                      getOpType(Inst->getType())));
+                                      Inst->getElement(), newTy));
+  }
+
+  void visitInitEnumDataAddrInst(InitEnumDataAddrInst *Inst) {
+    getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+    auto opd = getOpValue(Inst->getOperand());
+    auto caseTy = opd->getType().getEnumElementType(
+        Inst->getElement(), getBuilder().getFunction().getModule());
+    auto expectedTy = getOpType(Inst->getType());
+    if (expectedTy != caseTy)
+      expectedTy = caseTy;
+    recordClonedInstruction(Inst, getBuilder().createInitEnumDataAddr(
+                                      getOpLocation(Inst->getLoc()), opd,
+                                      Inst->getElement(), expectedTy));
   }
 
   /// Projections should not change the type if the type is not specialized.
@@ -332,24 +344,48 @@ class OpaqueSpecializerCloner
           loc, opd, type, /*withoutActuallyEscaping*/ false);
     } else if (opd->getType().isTrivial(CurFn)) {
       return getBuilder().createUncheckedTrivialBitCast(loc, opd, type);
-    } else if (opd->getType().isObject()) {
+    } else if (opd->getType().canRefCast(opd->getType(), type,
+                                         CurFn.getModule())) {
       return getBuilder().createUncheckedRefCast(loc, opd, type);
     } else {
       // This could be improved upon by recursively recomposing the type.
       auto *stackLoc = getBuilder().createAllocStack(loc, type);
-      auto *addr =
-          getBuilder().createUncheckedAddrCast(loc, stackLoc, opd->getType());
-      getBuilder().createTrivialStoreOr(loc, addr, opd,
-                                        StoreOwnershipQualifier::Init);
+      auto *addr = getBuilder().createUncheckedAddrCast(
+          loc, stackLoc, opd->getType().getAddressType());
+      getBuilder().createTrivialStoreOr(loc, opd, addr,
+                                        StoreOwnershipQualifier::Init, true);
       SILValue res = getBuilder().createTrivialLoadOr(
-          loc, addr, LoadOwnershipQualifier::Take);
+          loc, stackLoc, LoadOwnershipQualifier::Take, true);
       getBuilder().createDeallocStack(loc, stackLoc);
       return res;
     }
   }
 
+  void replaceBlockArgumentType(SILLocation loc, SILBasicBlock *destBlock,
+                                SILType withType) {
+    assert(destBlock->getArguments().size() == 1);
+
+    auto origType = (*destBlock->args_begin())->getType();
+    auto origPhi = destBlock->getPhiArguments()[0];
+    SILValue undef = SILUndef::get(origType, getBuilder().getFunction());
+    SmallVector<Operand *, 8> useList(origPhi->use_begin(), origPhi->use_end());
+    for (auto *use : useList) {
+      use->set(undef);
+    }
+
+    auto *newPhi =
+        destBlock->replacePhiArgument(0, withType, origPhi->getOwnershipKind());
+
+    getBuilder().setInsertionPoint(destBlock->begin());
+    auto cast = createCast(loc, newPhi, origType);
+    for (auto *use : useList) {
+      use->set(cast);
+    }
+  }
+
   void fixUp(SILFunction *) {
-    for (auto &BB : getBuilder().getFunction()) {
+    auto &clonedFunction = getBuilder().getFunction();
+    for (auto &BB : clonedFunction) {
       for (auto &cloned : BB) {
         // Fix up the type of try_apply successor block arguments.
         if (auto *tryApply = dyn_cast<TryApplyInst>(&cloned)) {
@@ -360,22 +396,25 @@ class OpaqueSpecializerCloner
           auto normalBBType = (*normalBB->args_begin())->getType();
           auto applyResultType = calleeConv.getSILResultType();
           if (normalBBType != calleeConv.getSILResultType()) {
-            auto origPhi = normalBB->getPhiArguments()[0];
-            SILValue undef =
-                SILUndef::get(normalBBType, getBuilder().getFunction());
-            SmallVector<Operand *, 8> useList(origPhi->use_begin(),
-                                              origPhi->use_end());
-            for (auto *use : useList) {
-              use->set(undef);
-            }
-
-            auto *newPhi = normalBB->replacePhiArgument(
-                0, applyResultType, origPhi->getOwnershipKind());
+            replaceBlockArgumentType(tryApply->getLoc(), normalBB, applyResultType);
+          }
+        }
+        // Fix up the type of switch_enum successor block arguments.
+        if (auto *switchEnum = dyn_cast<SwitchEnumInst>(&cloned)) {
+          SILType enumTy = switchEnum->getOperand()->getType();
+          for (unsigned i = 0, e = switchEnum->getNumCases(); i < e; ++i) {
+            EnumElementDecl *elt;
+            SILBasicBlock *dest;
+            std::tie(elt, dest) = switchEnum->getCase(i);
+
+            if (elt->hasAssociatedValues() &&
+                dest->getArguments().size() == 1) {
+              SILType eltArgTy =
+                  enumTy.getEnumElementType(elt, clonedFunction.getModule());
+              SILType bbArgTy = dest->getArguments()[0]->getType();
+              if (eltArgTy != bbArgTy)
+                replaceBlockArgumentType(switchEnum->getLoc(), dest, eltArgTy);
 
-            getBuilder().setInsertionPoint(normalBB->begin());
-            auto cast = createCast(tryApply->getLoc(), newPhi, normalBBType);
-            for (auto *use : useList) {
-              use->set(cast);
             }
           }
         }

test/SILOptimizer/specialize_opaque_type_archetypes.swift

@@ -478,3 +478,64 @@ extension PA {
     useP(p.1(5))
   }
 }
+
+public struct Foo {
+  var id : Int = 0
+  var p : Int64 = 1
+}
+
+struct Test : RandomAccessCollection {
+    struct Index : Comparable, Hashable {
+        var identifier: AnyHashable?
+        var offset: Int
+
+        static func < (lhs: Index, rhs: Index) -> Bool {
+            return lhs.offset < rhs.offset
+        }
+
+        func hash(into hasher: inout Hasher) {
+            hasher.combine(identifier)
+            hasher.combine(offset)
+        }
+    }
+
+    let foos: [Foo]
+    let ids: [AnyHashable]
+
+    init(foos: [Foo]) {
+        self.foos = foos
+        self.ids = foos.map { $0.id }
+    }
+
+    func _index(atOffset n: Int) -> Index {
+        return Index(identifier: ids.isEmpty ? nil : ids[n], offset: n)
+    }
+
+    var startIndex: Index {
+        return _index(atOffset: 0)
+    }
+
+    var endIndex: Index {
+        return Index(identifier: nil, offset: ids.endIndex)
+    }
+
+    func index(after i: Index) -> Index {
+        return _index(atOffset: i.offset + 1)
+    }
+
+    func index(before i: Index) -> Index {
+        return _index(atOffset: i.offset - 1)
+    }
+
+    func distance(from start: Index, to end: Index) -> Int {
+        return end.offset - start.offset
+    }
+
+    func index(_ i: Index, offsetBy n: Int) -> Index {
+        return _index(atOffset: i.offset + n)
+    }
+
+   subscript(i: Index) -> some P {
+        return foos[i.offset].p
+    }
+}