SIL: cache case indices in enum instructions.

As we do with field indices for struct instructions.
This avoids quadratic behavior in case of enums with lots of cases.
Also: cache field and enum case indices in the SILModule.

Author: eeckstein

include/swift/SIL/SILInstruction.h

@@ -812,6 +812,10 @@ class SILInstruction : public llvm::ilist_node<SILInstruction> {
 
   /// This is supportable but usually suggests a logic mistake.
   static bool classof(const ValueBase *) = delete;
+  
+protected:
+  unsigned getCachedFieldIndex(NominalTypeDecl *decl, VarDecl *property);
+  unsigned getCachedCaseIndex(EnumElementDecl *enumElement);
 };
 
 inline SILNodePointer::SILNodePointer(const SILInstruction *inst) :
@@ -6010,6 +6014,7 @@ class EnumInst
     : public InstructionBase<SILInstructionKind::EnumInst,
                              FirstArgOwnershipForwardingSingleValueInst> {
   friend SILBuilder;
+  enum : unsigned { InvalidCaseIndex = ~unsigned(0) };
 
   Optional<FixedOperandList<1>> OptionalOperand;
   EnumElementDecl *Element;
@@ -6019,6 +6024,8 @@ class EnumInst
            ValueOwnershipKind forwardingOwnershipKind)
       : InstructionBase(DebugLoc, ResultTy, forwardingOwnershipKind),
         Element(Element) {
+    SILNode::Bits.EnumInst.CaseIndex = InvalidCaseIndex;
+
     if (Operand) {
       OptionalOperand.emplace(this, Operand);
     }
@@ -6027,6 +6034,16 @@ class EnumInst
 public:
   EnumElementDecl *getElement() const { return Element; }
 
+  unsigned getCaseIndex() {
+    unsigned idx = SILNode::Bits.EnumInst.CaseIndex;
+    if (idx != InvalidCaseIndex)
+      return idx;
+
+    unsigned index = getCachedCaseIndex(getElement());
+    SILNode::Bits.EnumInst.CaseIndex = index;
+    return index;
+  }
+
   bool hasOperand() const { return OptionalOperand.hasValue(); }
   SILValue getOperand() const { return OptionalOperand->asValueArray()[0]; }
 
@@ -6048,6 +6065,7 @@ class UncheckedEnumDataInst
     : public UnaryInstructionBase<SILInstructionKind::UncheckedEnumDataInst,
                                   FirstArgOwnershipForwardingSingleValueInst> {
   friend SILBuilder;
+  enum : unsigned { InvalidCaseIndex = ~unsigned(0) };
 
   EnumElementDecl *Element;
 
@@ -6056,11 +6074,23 @@ class UncheckedEnumDataInst
                         ValueOwnershipKind forwardingOwnershipKind)
       : UnaryInstructionBase(DebugLoc, Operand, ResultTy,
                              forwardingOwnershipKind),
-        Element(Element) {}
+        Element(Element) {
+    SILNode::Bits.UncheckedEnumDataInst.CaseIndex = InvalidCaseIndex;
+  }
 
 public:
   EnumElementDecl *getElement() const { return Element; }
 
+  unsigned getCaseIndex() {
+    unsigned idx = SILNode::Bits.UncheckedEnumDataInst.CaseIndex;
+    if (idx != InvalidCaseIndex)
+      return idx;
+
+    unsigned index = getCachedCaseIndex(getElement());
+    SILNode::Bits.UncheckedEnumDataInst.CaseIndex = index;
+    return index;
+  }
+
   EnumDecl *getEnumDecl() const {
     auto *E = getOperand()->getType().getEnumOrBoundGenericEnum();
     assert(E && "Operand of unchecked_enum_data must be of enum type");
@@ -6086,15 +6116,28 @@ class InitEnumDataAddrInst
                                 SingleValueInstruction>
 {
   friend SILBuilder;
+  enum : unsigned { InvalidCaseIndex = ~unsigned(0) };
 
   EnumElementDecl *Element;
 
   InitEnumDataAddrInst(SILDebugLocation DebugLoc, SILValue Operand,
                        EnumElementDecl *Element, SILType ResultTy)
-      : UnaryInstructionBase(DebugLoc, Operand, ResultTy), Element(Element) {}
+      : UnaryInstructionBase(DebugLoc, Operand, ResultTy), Element(Element) {
+    SILNode::Bits.InitEnumDataAddrInst.CaseIndex = InvalidCaseIndex;
+  }
 
 public:
   EnumElementDecl *getElement() const { return Element; }
+
+  unsigned getCaseIndex() {
+    unsigned idx = SILNode::Bits.InitEnumDataAddrInst.CaseIndex;
+    if (idx != InvalidCaseIndex)
+      return idx;
+
+    unsigned index = getCachedCaseIndex(getElement());
+    SILNode::Bits.InitEnumDataAddrInst.CaseIndex = index;
+    return index;
+  }
 };
 
 /// InjectEnumAddrInst - Tags an enum as containing a case. The data for
@@ -6104,15 +6147,28 @@ class InjectEnumAddrInst
                                 NonValueInstruction>
 {
   friend SILBuilder;
+  enum : unsigned { InvalidCaseIndex = ~unsigned(0) };
 
   EnumElementDecl *Element;
 
   InjectEnumAddrInst(SILDebugLocation DebugLoc, SILValue Operand,
                      EnumElementDecl *Element)
-      : UnaryInstructionBase(DebugLoc, Operand), Element(Element) {}
+      : UnaryInstructionBase(DebugLoc, Operand), Element(Element) {
+    SILNode::Bits.InjectEnumAddrInst.CaseIndex = InvalidCaseIndex;
+  }
 
 public:
   EnumElementDecl *getElement() const { return Element; }
+
+  unsigned getCaseIndex() {
+    unsigned idx = SILNode::Bits.InjectEnumAddrInst.CaseIndex;
+    if (idx != InvalidCaseIndex)
+      return idx;
+
+    unsigned index = getCachedCaseIndex(getElement());
+    SILNode::Bits.InjectEnumAddrInst.CaseIndex = index;
+    return index;
+  }
 };
 
 /// Invalidate an enum value and take ownership of its payload data
@@ -6122,34 +6178,35 @@ class UncheckedTakeEnumDataAddrInst
                                 SingleValueInstruction>
 {
   friend SILBuilder;
+  enum : unsigned { InvalidCaseIndex = ~unsigned(0) };
 
   EnumElementDecl *Element;
 
   UncheckedTakeEnumDataAddrInst(SILDebugLocation DebugLoc, SILValue Operand,
                                 EnumElementDecl *Element, SILType ResultTy)
-      : UnaryInstructionBase(DebugLoc, Operand, ResultTy), Element(Element) {}
+      : UnaryInstructionBase(DebugLoc, Operand, ResultTy), Element(Element) {
+    SILNode::Bits.UncheckedTakeEnumDataAddrInst.CaseIndex = InvalidCaseIndex;
+  }
 
 public:
   EnumElementDecl *getElement() const { return Element; }
 
+  unsigned getCaseIndex() {
+    unsigned idx = SILNode::Bits.UncheckedTakeEnumDataAddrInst.CaseIndex;
+    if (idx != InvalidCaseIndex)
+      return idx;
+
+    unsigned index = getCachedCaseIndex(getElement());
+    SILNode::Bits.UncheckedTakeEnumDataAddrInst.CaseIndex = index;
+    return index;
+  }
+
   EnumDecl *getEnumDecl() const {
     auto *E = getOperand()->getType().getEnumOrBoundGenericEnum();
     assert(E && "Operand of unchecked_take_enum_data_addr must be of enum"
                 " type");
     return E;
   }
-
-  unsigned getElementNo() const {
-    unsigned i = 0;
-    for (EnumElementDecl *E : getEnumDecl()->getAllElements()) {
-      if (E == Element)
-        return i;
-      ++i;
-    }
-    llvm_unreachable(
-        "An unchecked_enum_data_addr's enumdecl should have at least "
-        "on element, the element that is being extracted");
-  }
 };
 
 // Abstract base class of all select instructions like select_enum,
@@ -6520,19 +6577,6 @@ class TupleElementAddrInst
   }
 };
 
-/// Get a unique index for a struct or class field in layout order.
-///
-/// Precondition: \p decl must be a non-resilient struct or class.
-///
-/// Precondition: \p field must be a stored property declared in \p decl,
-///               not in a superclass.
-///
-/// Postcondition: The returned index is unique across all properties in the
-///                object, including properties declared in a superclass.
-unsigned getFieldIndex(NominalTypeDecl *decl, VarDecl *property);
-
-unsigned getCaseIndex(EnumElementDecl *enumElement);
-
 unsigned getNumFieldsInNominal(NominalTypeDecl *decl);
 
 /// Get the property for a struct or class by its unique index, or nullptr if
@@ -6579,12 +6623,14 @@ class FieldIndexCacheBase : public ParentTy {
 
   VarDecl *getField() const { return field; }
 
-  unsigned getFieldIndex() const {
+  unsigned getFieldIndex() {
     unsigned idx = SILNode::Bits.FieldIndexCacheBase.FieldIndex;
     if (idx != InvalidFieldIndex)
       return idx;
-
-    return const_cast<FieldIndexCacheBase *>(this)->cacheFieldIndex();
+      
+    idx = ParentTy::getCachedFieldIndex(getParentDecl(), getField());
+    SILNode::Bits.FieldIndexCacheBase.FieldIndex = idx;
+    return idx;
   }
 
   NominalTypeDecl *getParentDecl() const {
@@ -6600,13 +6646,6 @@ class FieldIndexCacheBase : public ParentTy {
            kind == SILNodeKind::StructElementAddrInst ||
            kind == SILNodeKind::RefElementAddrInst;
   }
-
-private:
-  unsigned cacheFieldIndex() {
-    unsigned index = swift::getFieldIndex(getParentDecl(), getField());
-    SILNode::Bits.FieldIndexCacheBase.FieldIndex = index;
-    return index;
-  }
 };
 
 /// Extract a physical, fragile field out of a value of struct type.

include/swift/SIL/SILModule.h

@@ -303,6 +303,9 @@ class SILModule {
   /// This is the set of undef values we've created, for uniquing purposes.
   llvm::DenseMap<SILType, SILUndef *> UndefValues;
 
+  llvm::DenseMap<std::pair<Decl *, VarDecl *>, unsigned> fieldIndices;
+  llvm::DenseMap<EnumElementDecl *, unsigned> enumCaseIndices;
+
   /// The stage of processing this module is at.
   SILStage Stage;
 
@@ -448,6 +451,19 @@ class SILModule {
   /// This should only be the case during parsing or deserialization.
   bool hasUnresolvedOpenedArchetypeDefinitions();
 
+  /// Get a unique index for a struct or class field in layout order.
+  ///
+  /// Precondition: \p decl must be a non-resilient struct or class.
+  ///
+  /// Precondition: \p field must be a stored property declared in \p decl,
+  ///               not in a superclass.
+  ///
+  /// Postcondition: The returned index is unique across all properties in the
+  ///                object, including properties declared in a superclass.
+  unsigned getFieldIndex(NominalTypeDecl *decl, VarDecl *property);
+
+  unsigned getCaseIndex(EnumElementDecl *enumElement);
+
   /// Called by SILBuilder whenever a new instruction is created and inserted.
   void notifyAddedInstruction(SILInstruction *inst);
 

include/swift/SIL/SILNode.h

@@ -356,6 +356,36 @@ class alignas(8) SILNode :
     FieldIndex : 32
   );
 
+  SWIFT_INLINE_BITFIELD_FULL(EnumInst,
+                                      SingleValueInstruction, 32,
+    : NumPadBits,
+    CaseIndex : 32
+  );
+
+  SWIFT_INLINE_BITFIELD_FULL(UncheckedEnumDataInst,
+                                      SingleValueInstruction, 32,
+    : NumPadBits,
+    CaseIndex : 32
+  );
+
+  SWIFT_INLINE_BITFIELD_FULL(InjectEnumAddrInst,
+                                      SILInstruction, 32,
+    : NumPadBits,
+    CaseIndex : 32
+  );
+
+  SWIFT_INLINE_BITFIELD_FULL(InitEnumDataAddrInst,
+                                      SingleValueInstruction, 32,
+    : NumPadBits,
+    CaseIndex : 32
+  );
+
+  SWIFT_INLINE_BITFIELD_FULL(UncheckedTakeEnumDataAddrInst,
+                                      SingleValueInstruction, 32,
+    : NumPadBits,
+    CaseIndex : 32
+  );
+
   SWIFT_INLINE_BITFIELD_EMPTY(MethodInst, SingleValueInstruction);
   // Ensure that WitnessMethodInst bitfield does not overflow.
   IBWTO_BITFIELD_EMPTY(WitnessMethodInst, MethodInst);

lib/SIL/IR/SILInstruction.cpp

@@ -1387,6 +1387,15 @@ bool SILInstruction::isMetaInstruction() const {
   llvm_unreachable("Instruction not handled in isMetaInstruction()!");
 }
 
+unsigned SILInstruction::getCachedFieldIndex(NominalTypeDecl *decl,
+                                             VarDecl *property) {
+  return getModule().getFieldIndex(decl, property);
+}
+
+unsigned SILInstruction::getCachedCaseIndex(EnumElementDecl *enumElement) {
+  return getModule().getCaseIndex(enumElement);
+}
+
 //===----------------------------------------------------------------------===//
 //                                 Utilities
 //===----------------------------------------------------------------------===//

lib/SIL/IR/SILInstructions.cpp

@@ -1378,26 +1378,6 @@ bool TupleExtractInst::isEltOnlyNonTrivialElt() const {
   return true;
 }
 
-/// Get a unique index for a struct or class field in layout order.
-unsigned swift::getFieldIndex(NominalTypeDecl *decl, VarDecl *field) {
-  unsigned index = 0;
-  if (auto *classDecl = dyn_cast<ClassDecl>(decl)) {
-    for (auto *superDecl = classDecl->getSuperclassDecl(); superDecl != nullptr;
-         superDecl = superDecl->getSuperclassDecl()) {
-      index += superDecl->getStoredProperties().size();
-    }
-  }
-  for (VarDecl *property : decl->getStoredProperties()) {
-    if (field == property) {
-      return index;
-    }
-    ++index;
-  }
-  llvm_unreachable("The field decl for a struct_extract, struct_element_addr, "
-                   "or ref_element_addr must be an accessible stored "
-                   "property of the operand type");
-}
-
 unsigned swift::getNumFieldsInNominal(NominalTypeDecl *decl) {
   unsigned count = 0;
   if (auto *classDecl = dyn_cast<ClassDecl>(decl)) {
@@ -1409,16 +1389,6 @@ unsigned swift::getNumFieldsInNominal(NominalTypeDecl *decl) {
   return count + decl->getStoredProperties().size();
 }
 
-unsigned swift::getCaseIndex(EnumElementDecl *enumElement) {
-  unsigned idx = 0;
-  for (EnumElementDecl *e : enumElement->getParentEnum()->getAllElements()) {
-    if (e == enumElement)
-      return idx;
-    ++idx;
-  }
-  llvm_unreachable("enum element not found in enum decl");
-}
-
 /// Get the property for a struct or class by its unique index.
 VarDecl *swift::getIndexedField(NominalTypeDecl *decl, unsigned index) {
   if (auto *structDecl = dyn_cast<StructDecl>(decl)) {