IRGen: Generalize the nullable optimization for single-payload enums.

jckarter · jckarter · commit 950ee3b0de9b · 2020-03-03T15:00:27.000-08:00
Augment the `isSingleRetainablePointer` check that allows IRGen to avoid adding branching around retain/release operations on enums that use the null pointer extra inhabitant with a more general "can value witness extra inhabitants" method on TypeInfo, which says whether a type's retain/release operations are safe to invoke on some or all of its extra inhabitants. This lets us generalize the optimization to include things like `String?` or `ClassProtocol?` which are common types with a nullable pointer in them.
diff --git a/lib/IRGen/GenEnum.cpp b/lib/IRGen/GenEnum.cpp
@@ -812,6 +812,14 @@ namespace {
         return false;
       return singleton->isSingleRetainablePointer(expansion, rc);
     }
+    
+    bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                             unsigned index) const override {
+      auto singleton = getSingleton();
+      if (!singleton)
+        return false;
+      return singleton->canValueWitnessExtraInhabitantsUpTo(IGM, index);
+    }
   };
 
   /// Implementation strategy for no-payload enums, in other words, 'C-like'
@@ -1674,6 +1682,9 @@ namespace {
       /// copy and destroy can pass through to retain and release entry
       /// points.
       NullableRefcounted,
+      /// The payload's value witnesses can handle the extra inhabitants we use
+      /// for no-payload tags, so we can forward all our calls to them.
+      ForwardToPayload,
     };
 
     CopyDestroyStrategy CopyDestroyKind;
@@ -1802,6 +1813,26 @@ namespace {
           && cast<FixedTypeInfo>(payloadTI)
             .getFixedExtraInhabitantCount(IGM) > 0) {
         CopyDestroyKind = NullableRefcounted;
+      // If the payload's value witnesses can accept the extra inhabitants we
+      // use, then we can forward to them instead of checking for empty tags.
+      // TODO: Do this for all types, not just loadable types.
+      } else if (tik >= TypeInfoKind::Loadable) {
+        ReferenceCounting refCounting;
+        (void)refCounting;
+        // Ensure that asking `canValueWitnessExtraInhabitantsUpTo` doesn't
+        // regress any places we were previously able to ask
+        // `isSingleRetainablePointer`.
+        assert(
+          (!payloadTI.isSingleRetainablePointer(ResilienceExpansion::Maximal,
+                                                &refCounting)
+           || payloadTI.canValueWitnessExtraInhabitantsUpTo(IGM, 0))
+          && "single-refcounted thing should be able to value-witness "
+             "extra inhabitant zero");
+        
+        unsigned numTags = ElementsWithNoPayload.size();
+        if (payloadTI.canValueWitnessExtraInhabitantsUpTo(IGM, numTags-1)) {
+          CopyDestroyKind = ForwardToPayload;
+        }
       }
     }
 
@@ -1811,6 +1842,12 @@ namespace {
       assert(NumExtraInhabitantTagValues != ~0U);
       return NumExtraInhabitantTagValues;
     }
+    
+    bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                             unsigned index) const override {
+      return getPayloadTypeInfo().canValueWitnessExtraInhabitantsUpTo(IGM,
+                                           index + NumExtraInhabitantTagValues);
+    }
 
     /// Emit a call into the runtime to get the current enum payload tag.
     /// This returns a tag index in the range [0..NumElements-1].
@@ -2412,6 +2449,7 @@ namespace {
       switch (CopyDestroyKind) {
       case NullableRefcounted:
         return IGM.getReferenceType(Refcounting);
+      case ForwardToPayload:
       case POD:
       case Normal:
       case ABIInaccessible:
@@ -2427,6 +2465,7 @@ namespace {
       case NullableRefcounted:
         IGF.emitStrongRetain(ptr, Refcounting, IGF.getDefaultAtomicity());
         return;
+      case ForwardToPayload:
       case POD:
       case Normal:
       case ABIInaccessible:
@@ -2440,6 +2479,7 @@ namespace {
       case NullableRefcounted:
         IGF.emitFixLifetime(ptr);
         return;
+      case ForwardToPayload:
       case POD:
       case Normal:
       case ABIInaccessible:
@@ -2453,6 +2493,7 @@ namespace {
       case NullableRefcounted:
         IGF.emitStrongRelease(ptr, Refcounting, IGF.getDefaultAtomicity());
         return;
+      case ForwardToPayload:
       case POD:
       case Normal:
       case ABIInaccessible:
@@ -2471,6 +2512,25 @@ namespace {
       if (extraTag)
         out.add(extraTag);
     }
+    
+    void unpackIntoPayloadExplosion(IRGenFunction &IGF,
+                                    Explosion &asEnumIn,
+                                    Explosion &asPayloadOut) const {
+      auto &payloadTI = getLoadablePayloadTypeInfo();
+      // Unpack as an instance of the payload type and use its copy operation.
+      auto srcBits = EnumPayload::fromExplosion(IGF.IGM, asEnumIn,
+                                                PayloadSchema);
+      payloadTI.unpackFromEnumPayload(IGF, srcBits, asPayloadOut, 0);
+    }
+    
+    void packFromPayloadExplosion(IRGenFunction &IGF,
+                                  Explosion &asPayloadIn,
+                                  Explosion &asEnumOut) const {
+      auto &payloadTI = getLoadablePayloadTypeInfo();
+      auto payload = EnumPayload::zero(IGF.IGM, PayloadSchema);
+      payloadTI.packIntoEnumPayload(IGF, payload, asPayloadIn, 0);
+      payload.explode(IGF.IGM, asEnumOut);
+    }
 
   public:
     void copy(IRGenFunction &IGF, Explosion &src, Explosion &dest,
@@ -2530,6 +2590,15 @@ namespace {
         dest.add(val);
         return;
       }
+
+      case ForwardToPayload: {
+        auto &payloadTI = getLoadablePayloadTypeInfo();
+        Explosion srcAsPayload, destAsPayload;
+        unpackIntoPayloadExplosion(IGF, src, srcAsPayload);
+        payloadTI.copy(IGF, srcAsPayload, destAsPayload, atomicity);
+        packFromPayloadExplosion(IGF, destAsPayload, dest);
+        return;
+      }
       }
     }
 
@@ -2585,6 +2654,16 @@ namespace {
         releaseRefcountedPayload(IGF, ptr);
         return;
       }
+
+      case ForwardToPayload: {
+        auto &payloadTI = getLoadablePayloadTypeInfo();
+        // Unpack as an instance of the payload type and use its consume
+        // operation.
+        Explosion srcAsPayload;
+        unpackIntoPayloadExplosion(IGF, src, srcAsPayload);
+        payloadTI.consume(IGF, srcAsPayload, atomicity);
+        return;
+      }
       }
     }
 
@@ -2630,6 +2709,16 @@ namespace {
         fixLifetimeOfRefcountedPayload(IGF, ptr);
         return;
       }
+
+      case ForwardToPayload: {
+        auto &payloadTI = getLoadablePayloadTypeInfo();
+        // Unpack as an instance of the payload type and use its fixLifetime
+        // operation.
+        Explosion srcAsPayload;
+        unpackIntoPayloadExplosion(IGF, srcAsPayload, srcAsPayload);
+        payloadTI.fixLifetime(IGF, src);
+        return;
+      }
       }
 
     }
@@ -2667,13 +2756,22 @@ namespace {
         }
 
         case NullableRefcounted: {
-          // Load the value as swift.refcounted, then hand to swift_release.
+          // Apply the payload's operation.
           addr = IGF.Builder.CreateBitCast(
               addr, getRefcountedPtrType(IGM)->getPointerTo());
           llvm::Value *ptr = IGF.Builder.CreateLoad(addr);
           releaseRefcountedPayload(IGF, ptr);
           return;
         }
+
+        case ForwardToPayload: {
+          auto &payloadTI = getPayloadTypeInfo();
+          // Apply the payload's operation.
+          addr = IGF.Builder.CreateBitCast(
+              addr, payloadTI.getStorageType()->getPointerTo());
+          payloadTI.destroy(IGF, addr, getPayloadType(IGF.IGM, T), isOutlined);
+          return;
+        }
         }
       } else {
         if (!IGF.IGM.getOptions().UseTypeLayoutValueHandling) {
@@ -2702,8 +2800,9 @@ namespace {
                                    Address addr) const override {
       // There is no need to bitcast from the enum address. Loading from the
       // reference type emits a bitcast to the proper reference type first.
-      return cast<LoadableTypeInfo>(getPayloadTypeInfo()).loadRefcountedPtr(
-        IGF, loc, addr).getValue();
+      return getLoadablePayloadTypeInfo()
+        .loadRefcountedPtr(IGF, loc, addr)
+        .getValue();
     }
   private:
     llvm::ConstantInt *getZeroExtraTagConstant(IRGenModule &IGM) const {
@@ -2843,6 +2942,18 @@ namespace {
         releaseRefcountedPayload(IGF, oldPtr);
         return;
       }
+
+      case ForwardToPayload: {
+        auto &payloadTI = getPayloadTypeInfo();
+        // Apply the payload's operation.
+        dest = IGF.Builder.CreateBitCast(dest,
+                                  payloadTI.getStorageType()->getPointerTo());
+        src = IGF.Builder.CreateBitCast(src,
+                                  payloadTI.getStorageType()->getPointerTo());
+        payloadTI.assign(IGF, dest, src, isTake,
+                         getPayloadType(IGF.IGM, T), isOutlined);
+        return;
+      }
       }
 
     }
@@ -2911,6 +3022,18 @@ namespace {
         IGF.Builder.CreateStore(srcPtr, destAddr);
         return;
       }
+
+      case ForwardToPayload: {
+        auto &payloadTI = getPayloadTypeInfo();
+        // Apply the payload's operation.
+        dest = IGF.Builder.CreateBitCast(dest,
+                                  payloadTI.getStorageType()->getPointerTo());
+        src = IGF.Builder.CreateBitCast(src,
+                                  payloadTI.getStorageType()->getPointerTo());
+        payloadTI.initialize(IGF, dest, src, isTake,
+                             getPayloadType(IGF.IGM, T), isOutlined);
+        return;
+      }
       }
     }
 
@@ -6093,6 +6216,10 @@ namespace {
                                           SILType ty) const override {
       return Strategy.buildTypeLayoutEntry(IGM, ty);
     }
+    bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                             unsigned index) const override {
+      return Strategy.canValueWitnessExtraInhabitantsUpTo(IGM, index);
+    }
   };
 
   template <class Base>
diff --git a/lib/IRGen/GenEnum.h b/lib/IRGen/GenEnum.h
@@ -463,8 +463,10 @@ class EnumImplStrategy {
 
   void emitResilientTagIndices(IRGenModule &IGM) const;
 
-protected:
-
+  virtual bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                                   unsigned index) const {
+    return false;
+  }
 
 private:
   EnumImplStrategy(const EnumImplStrategy &) = delete;
diff --git a/lib/IRGen/GenExistential.cpp b/lib/IRGen/GenExistential.cpp
@@ -1049,6 +1049,11 @@ class ClassExistentialTypeInfo final
     if (refcounting) *refcounting = Refcounting;
     return getNumStoredProtocols() == 0;
   }
+  
+  bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                           unsigned index) const override {
+    return index == 0;
+  }
 
   const LoadableTypeInfo &
   getValueTypeInfoForExtraInhabitants(IRGenModule &IGM) const {
diff --git a/lib/IRGen/GenRecord.h b/lib/IRGen/GenRecord.h
@@ -595,6 +595,28 @@ class RecordTypeInfo<Impl, Base, FieldImpl,
     return 0;
   }
 
+  bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                           unsigned index) const override {
+    if (auto field = asImpl().getFixedExtraInhabitantProvidingField(IGM)) {
+      // The non-extra-inhabitant-providing fields of the type must be
+      // trivial, because an enum may contain garbage values in those fields'
+      // storage which the value witness operation won't handle.
+      for (auto &otherField : asImpl().getFields()) {
+        if (field == &otherField)
+          continue;
+        auto &ti = otherField.getTypeInfo();
+        if (!ti.isPOD(ResilienceExpansion::Maximal)) {
+          return false;
+        }
+      }
+
+      return field->getTypeInfo()
+        .canValueWitnessExtraInhabitantsUpTo(IGM, index);
+    }
+    
+    return false;
+  }
+
   APInt getFixedExtraInhabitantValue(IRGenModule &IGM,
                                      unsigned bits,
                                      unsigned index) const override {
diff --git a/lib/IRGen/GenStruct.cpp b/lib/IRGen/GenStruct.cpp
@@ -219,7 +219,7 @@ namespace {
         return false;
       return fields[0].getTypeInfo().isSingleRetainablePointer(expansion, rc);
     }
-    
+       
     void verify(IRGenTypeVerifierFunction &IGF,
                 llvm::Value *metadata,
                 SILType structType) const override {
diff --git a/lib/IRGen/GenValueWitness.cpp b/lib/IRGen/GenValueWitness.cpp
@@ -1245,6 +1245,18 @@ FixedPacking TypeInfo::getFixedPacking(IRGenModule &IGM) const {
   return FixedPacking::Allocate;
 }
 
+bool TypeInfo::canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                                   unsigned index) const {
+  // If this type is POD, then its value witnesses are trivial, so can handle
+  // any bit pattern.
+  if (isPOD(ResilienceExpansion::Maximal)) {
+    return true;
+  }
+  
+  // By default, assume that extra inhabitants must be branched out on.
+  return false;
+}
+
 Address TypeInfo::indexArray(IRGenFunction &IGF, Address base,
                              llvm::Value *index, SILType T) const {
   // The stride of a Swift type may not match its LLVM size. If we know we have
diff --git a/lib/IRGen/HeapTypeInfo.h b/lib/IRGen/HeapTypeInfo.h
@@ -75,6 +75,13 @@ class HeapTypeInfo
     return true;
   }
   
+  bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                           unsigned index) const override {
+    // Refcounting functions are no-ops when passed a null pointer, which is the
+    // first extra inhabitant.
+    return index == 0;
+  }
+  
   IsaEncoding getIsaEncoding(ResilienceExpansion expansion) const {
     switch (asDerived().getReferenceCounting()) {
     // We can access the isa of pure Swift heap objects directly.
diff --git a/lib/IRGen/TypeInfo.h b/lib/IRGen/TypeInfo.h
@@ -408,6 +408,17 @@ class TypeInfo {
   /// have extra inhabitants based on type arguments?
   virtual bool mayHaveExtraInhabitants(IRGenModule &IGM) const = 0;
 
+  /// Returns true if the value witness operations on this type work correctly
+  /// with extra inhabitants up to the given index.
+  ///
+  /// An example of this is retainable pointers. The first extra inhabitant for
+  /// these types is the null pointer, on which swift_retain is a harmless
+  /// no-op. If this predicate returns true, then a single-payload enum with
+  /// this type as its payload (like Optional<T>) can avoid additional branching
+  /// on the enum tag for value witness operations.
+  virtual bool canValueWitnessExtraInhabitantsUpTo(IRGenModule &IGM,
+                                                   unsigned index) const;
+  
   /// Get the tag of a single payload enum with a payload of this type (\p T) e.g
   /// Optional<T>.
   virtual llvm::Value *getEnumTagSinglePayload(IRGenFunction &IGF,
diff --git a/test/IRGen/enum_single_payload_forward_to_payload.swift b/test/IRGen/enum_single_payload_forward_to_payload.swift
diff --git a/test/IRGen/weak.sil b/test/IRGen/weak.sil

Original file line number	Diff line number	Diff line change
`@@ -219,7 +219,7 @@ namespace {`
`219`	`219`	`return false;`
`220`	`220`	`return fields[0].getTypeInfo().isSingleRetainablePointer(expansion, rc);`
`221`	`221`	`}`
`222`		`-`
	`222`	`+`
`223`	`223`	`void verify(IRGenTypeVerifierFunction &IGF,`
`224`	`224`	`llvm::Value *metadata,`
`225`	`225`	`SILType structType) const override {`