Merge pull request #34633 from gottesmm/pr-200203a48964c3efb090a7f4f2389b115a88b607

[ownership] Change switch_enum to be a true forwarding instructions.

Author: gottesmm

include/swift/SIL/SILInstruction.h

@@ -40,6 +40,7 @@
 #include "swift/Strings.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/ilist.h"
 #include "llvm/ADT/ilist_node.h"
 #include "llvm/Support/TrailingObjects.h"
@@ -7318,6 +7319,20 @@ class TermInst : public NonValueInstruction {
   }
 };
 
+class OwnershipForwardingTermInst : public TermInst {
+  ValueOwnershipKind ownershipKind;
+
+protected:
+  OwnershipForwardingTermInst(SILInstructionKind kind,
+                              SILDebugLocation debugLoc,
+                              ValueOwnershipKind ownershipKind)
+      : TermInst(kind, debugLoc), ownershipKind(ownershipKind) {}
+
+public:
+  ValueOwnershipKind getOwnershipKind() const { return ownershipKind; }
+  void setOwnershipKind(ValueOwnershipKind newKind) { ownershipKind = newKind; }
+};
+
 /// UnreachableInst - Position in the code which would be undefined to reach.
 /// These are always implicitly generated, e.g. when falling off the end of a
 /// function or after a no-return function call.
@@ -7806,9 +7821,9 @@ class SwitchValueInst final
   }
 };
 
-/// Common implementation for the switch_enum and
-/// switch_enum_addr instructions.
-class SwitchEnumInstBase : public TermInst {
+/// Common implementation for the switch_enum and switch_enum_addr instructions.
+template <typename BaseTy>
+class SwitchEnumInstBase : public BaseTy {
   FixedOperandList<1> Operands;
 
   // Tail-allocated after the SwitchEnumInst record are:
@@ -7837,48 +7852,97 @@ class SwitchEnumInstBase : public TermInst {
   }
 
 protected:
+  template <typename... Rest>
   SwitchEnumInstBase(
       SILInstructionKind Kind, SILDebugLocation DebugLoc, SILValue Operand,
       SILBasicBlock *DefaultBB,
       ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
-      Optional<ArrayRef<ProfileCounter>> Counts, ProfileCounter DefaultCount);
+      Optional<ArrayRef<ProfileCounter>> Counts, ProfileCounter DefaultCount,
+      Rest &&... rest)
+      : BaseTy(Kind, DebugLoc, std::forward<Rest>(rest)...),
+        Operands(this, Operand) {
+    SILInstruction::Bits.SEIBase.HasDefault = bool(DefaultBB);
+    SILInstruction::Bits.SEIBase.NumCases = CaseBBs.size();
+    // Initialize the case and successor arrays.
+    auto *cases = getCaseBuf();
+    auto *succs = getSuccessorBuf();
+    for (unsigned i = 0, size = CaseBBs.size(); i < size; ++i) {
+      cases[i] = CaseBBs[i].first;
+      if (Counts) {
+        ::new (succs + i)
+            SILSuccessor(this, CaseBBs[i].second, Counts.getValue()[i]);
+      } else {
+        ::new (succs + i) SILSuccessor(this, CaseBBs[i].second);
+      }
+    }
+
+    if (hasDefault()) {
+      ::new (succs + getNumCases()) SILSuccessor(this, DefaultBB, DefaultCount);
+    }
+  }
 
-  template <typename SWITCH_ENUM_INST>
+  template <typename SWITCH_ENUM_INST, typename... RestTys>
   static SWITCH_ENUM_INST *createSwitchEnum(
       SILDebugLocation DebugLoc, SILValue Operand, SILBasicBlock *DefaultBB,
       ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
       SILFunction &F, Optional<ArrayRef<ProfileCounter>> Counts,
-      ProfileCounter DefaultCount);
+      ProfileCounter DefaultCount, RestTys &&... restArgs);
 
 public:
   /// Clean up tail-allocated successor records for the switch cases.
-  ~SwitchEnumInstBase();
+  ~SwitchEnumInstBase() {
+    // Destroy the successor records to keep the CFG up to date.
+    auto *succs = getSuccessorBuf();
+    for (unsigned i = 0, end = getNumCases() + hasDefault(); i < end; ++i) {
+      succs[i].~SILSuccessor();
+    }
+  }
 
   SILValue getOperand() const { return Operands[0].get(); }
 
   ArrayRef<Operand> getAllOperands() const { return Operands.asArray(); }
   MutableArrayRef<Operand> getAllOperands() { return Operands.asArray(); }
 
-  SuccessorListTy getSuccessors() {
+  TermInst::SuccessorListTy getSuccessors() {
     return MutableArrayRef<SILSuccessor>{getSuccessorBuf(),
                            static_cast<size_t>(getNumCases() + hasDefault())};
   }
 
-  unsigned getNumCases() const {
-    return SILInstruction::Bits.SwitchEnumInstBase.NumCases;
-  }
+  unsigned getNumCases() const { return SILInstruction::Bits.SEIBase.NumCases; }
+
   std::pair<EnumElementDecl*, SILBasicBlock*>
   getCase(unsigned i) const {
     assert(i < getNumCases() && "case out of bounds");
     return {getCaseBuf()[i], getSuccessorBuf()[i].getBB()};
   }
+
   ProfileCounter getCaseCount(unsigned i) const {
     assert(i < getNumCases() && "case out of bounds");
     return getSuccessorBuf()[i].getCount();
   }
 
   // Swap the cases at indices \p i and \p j.
-  void swapCase(unsigned i, unsigned j);
+  void swapCase(unsigned i, unsigned j) {
+    assert(i < getNumCases() && "First index is out of bounds?!");
+    assert(j < getNumCases() && "Second index is out of bounds?!");
+
+    auto *succs = getSuccessorBuf();
+
+    // First grab our destination blocks.
+    SILBasicBlock *iBlock = succs[i].getBB();
+    SILBasicBlock *jBlock = succs[j].getBB();
+
+    // Then destroy the sil successors and reinitialize them with the new things
+    // that they are pointing at.
+    succs[i].~SILSuccessor();
+    ::new (succs + i) SILSuccessor(this, jBlock);
+    succs[j].~SILSuccessor();
+    ::new (succs + j) SILSuccessor(this, iBlock);
+
+    // Now swap our cases.
+    auto *cases = getCaseBuf();
+    std::swap(cases[i], cases[j]);
+  }
 
   /// Return the block that will be branched to on the specified enum
   /// case.
@@ -7893,22 +7957,69 @@ class SwitchEnumInstBase : public TermInst {
   }
 
   /// If the default refers to exactly one case decl, return it.
-  NullablePtr<EnumElementDecl> getUniqueCaseForDefault();
+  NullablePtr<EnumElementDecl> getUniqueCaseForDefault() {
+    auto enumValue = getOperand();
+    SILType enumType = enumValue->getType();
+
+    auto *f = SILInstruction::getFunction();
+    if (!enumType.isEffectivelyExhaustiveEnumType(f))
+      return nullptr;
+
+    EnumDecl *decl = enumType.getEnumOrBoundGenericEnum();
+    assert(decl && "switch_enum operand is not an enum");
+
+    SmallPtrSet<EnumElementDecl *, 4> unswitchedElts;
+    for (auto elt : decl->getAllElements())
+      unswitchedElts.insert(elt);
+
+    for (unsigned i = 0, e = getNumCases(); i != e; ++i) {
+      auto Entry = getCase(i);
+      unswitchedElts.erase(Entry.first);
+    }
+
+    if (unswitchedElts.size() == 1)
+      return *unswitchedElts.begin();
+
+    return nullptr;
+  }
 
   /// If the given block only has one enum element decl matched to it,
   /// return it.
-  NullablePtr<EnumElementDecl> getUniqueCaseForDestination(SILBasicBlock *BB);
-
-  bool hasDefault() const {
-    return SILInstruction::Bits.SwitchEnumInstBase.HasDefault;
+  NullablePtr<EnumElementDecl>
+  getUniqueCaseForDestination(SILBasicBlock *block) {
+    SILValue value = getOperand();
+    SILType enumType = value->getType();
+    EnumDecl *decl = enumType.getEnumOrBoundGenericEnum();
+    assert(decl && "switch_enum operand is not an enum");
+    (void)decl;
+
+    EnumElementDecl *eltDecl = nullptr;
+    for (unsigned i : range(getNumCases())) {
+      auto entry = getCase(i);
+      if (entry.second == block) {
+        if (eltDecl != nullptr)
+          return nullptr;
+        eltDecl = entry.first;
+      }
+    }
+    if (!eltDecl && hasDefault() && getDefaultBB() == block) {
+      return getUniqueCaseForDefault();
+    }
+    return eltDecl;
   }
 
+  bool hasDefault() const { return SILInstruction::Bits.SEIBase.HasDefault; }
+
   SILBasicBlock *getDefaultBB() const {
     assert(hasDefault() && "doesn't have a default");
     return getSuccessorBuf()[getNumCases()];
   }
 
-  NullablePtr<SILBasicBlock> getDefaultBBOrNull() const;
+  NullablePtr<SILBasicBlock> getDefaultBBOrNull() const {
+    if (!hasDefault())
+      return nullptr;
+    return getDefaultBB();
+  }
 
   ProfileCounter getDefaultCount() const {
     assert(hasDefault() && "doesn't have a default");
@@ -7925,7 +8036,7 @@ class SwitchEnumInstBase : public TermInst {
 /// passed into the corresponding destination block as an argument.
 class SwitchEnumInst
     : public InstructionBase<SILInstructionKind::SwitchEnumInst,
-                             SwitchEnumInstBase> {
+                             SwitchEnumInstBase<OwnershipForwardingTermInst>> {
   friend SILBuilder;
 
 private:
@@ -7937,7 +8048,7 @@ class SwitchEnumInst
       Optional<ArrayRef<ProfileCounter>> CaseCounts,
       ProfileCounter DefaultCount)
       : InstructionBase(DebugLoc, Operand, DefaultBB, CaseBBs, CaseCounts,
-                        DefaultCount) {}
+                        DefaultCount, Operand.getOwnershipKind()) {}
   static SwitchEnumInst *
   create(SILDebugLocation DebugLoc, SILValue Operand, SILBasicBlock *DefaultBB,
          ArrayRef<std::pair<EnumElementDecl *, SILBasicBlock *>> CaseBBs,
@@ -7948,7 +8059,7 @@ class SwitchEnumInst
 /// A switch on an enum's discriminator in memory.
 class SwitchEnumAddrInst
     : public InstructionBase<SILInstructionKind::SwitchEnumAddrInst,
-                             SwitchEnumInstBase> {
+                             SwitchEnumInstBase<TermInst>> {
   friend SILBuilder;
 
 private:

include/swift/SIL/SILNode.h

@@ -362,12 +362,25 @@ class alignas(8) SILNode {
   IBWTO_BITFIELD(SwitchValueInst, TermInst, 1,
     HasDefault : 1
   );
-  SWIFT_INLINE_BITFIELD_FULL(SwitchEnumInstBase, TermInst, 1+32,
+
+  // Special handling for SwitchEnumInstBase.
+  //
+  // We assume all subsequent SwitchEnumBit uses do not use any further bits.
+  SWIFT_INLINE_BITFIELD(SEIBase, TermInst, 32-NumTermInstBits,
+    // Does this switch enum inst have a default block.
     HasDefault : 1,
-    : NumPadBits,
-    NumCases : 32
+    // Number of cases 
+    NumCases : 31 - NumTermInstBits;
+    template <typename BaseTy>
+    friend class SwitchEnumInstBase;
   );
 
+#define SEIB_BITFIELD_EMPTY(T, U) \
+  IBWTO_BITFIELD_EMPTY(T, U)
+
+  SEIB_BITFIELD_EMPTY(SwitchEnumInst, SEIBase);
+  SEIB_BITFIELD_EMPTY(SwitchEnumAddrInst, SEIBase);
+
   SWIFT_INLINE_BITFIELD_EMPTY(MultipleValueInstruction, SILInstruction);
 
   SWIFT_INLINE_BITFIELD(BeginCOWMutationInst, MultipleValueInstruction, 1,

include/swift/SIL/SILType.h

@@ -429,6 +429,10 @@ class SILType {
   SILType getEnumElementType(EnumElementDecl *elt, SILModule &M,
                              TypeExpansionContext context) const;
 
+  /// Given that this is an enum type, return true if this type is effectively
+  /// exhausted.
+  bool isEffectivelyExhaustiveEnumType(SILFunction *f);
+
   /// Given that this is a tuple type, return the lowered type of the
   /// given tuple element.  The result will have the same value
   /// category as the base type.