[WIP] [NO-MERGE] Make "unique" flag more robust and update access enforcement opts to make use of it.

docs/SIL.rst

@@ -2238,6 +2238,7 @@ alloc_ref
 ::
 
   sil-instruction ::= 'alloc_ref'
+                        ('[' 'unique' ']')?
                         ('[' 'objc' ']')?
                         ('[' 'stack' ']')?
                         ('[' 'tail_elems' sil-type '*' sil-operand ']')*
@@ -2273,6 +2274,11 @@ The instructions ``ref_tail_addr`` and ``tail_addr`` can be used to project
 the tail elements.
 The ``objc`` attribute cannot be used together with ``tail_elems``.
 
+The optional ``unique`` attribute indicates that this is the only reference to
+the allocated object. In other words, the reference will not escape or be
+written to dynamically. This allows access of the object to be static instead
+of dynamic.
+
 alloc_ref_dynamic
 `````````````````
 ::

include/swift/SIL/InstructionUtils.h

@@ -173,6 +173,14 @@ struct PolymorphicBuiltinSpecializedOverloadInfo {
 /// return SILValue().
 SILValue getStaticOverloadForSpecializedPolymorphicBuiltin(BuiltinInst *bi);
 
+SILFunction *getDestructor(AllocRefInstBase *ari);
+
+/// \returns None if \p ref is unique. Otherwise, returns an error message in
+/// the form of a string. If \p verbose is true, will print out diagnostics,
+/// such as the reason for failure. This is helpful when the funciton is being
+/// used by the SILVerifier.
+Optional<StringRef> isReferenceUnique(AllocRefInstBase *ref, bool verbose);
+
 } // end namespace swift
 
 #endif

include/swift/SIL/MemAccessUtils.h

@@ -362,6 +362,9 @@ class AccessedStorage {
   bool isLetAccess(SILFunction *F) const;
 
   bool isUniquelyIdentified() const {
+    if (auto *ref = dyn_cast<AllocRefInst>(value))
+      return ref->isUniqueReference();
+
     switch (getKind()) {
     case Box:
     case Stack:

include/swift/SIL/SILBuilder.h

@@ -421,17 +421,17 @@ class SILBuilder {
                                          Var, hasDynamicLifetime));
   }
 
-  AllocRefInst *createAllocRef(SILLocation Loc, SILType ObjectType,
-                               bool objc, bool canAllocOnStack,
+  AllocRefInst *createAllocRef(SILLocation Loc, SILType ObjectType, bool objc,
+                               bool canAllocOnStack, bool isUnique,
                                ArrayRef<SILType> ElementTypes,
                                ArrayRef<SILValue> ElementCountOperands) {
     // AllocRefInsts expand to function calls and can therefore not be
     // counted towards the function prologue.
     assert(!Loc.isInPrologue());
-    return insert(AllocRefInst::create(getSILDebugLocation(Loc), getFunction(),
-                                       ObjectType, objc, canAllocOnStack,
-                                       ElementTypes, ElementCountOperands,
-                                       C.OpenedArchetypes));
+    return insert(AllocRefInst::create(
+        getSILDebugLocation(Loc), getFunction(), ObjectType, objc,
+        canAllocOnStack, isUnique, ElementTypes, ElementCountOperands,
+        C.OpenedArchetypes));
   }
 
   AllocRefDynamicInst *createAllocRefDynamic(SILLocation Loc, SILValue operand,

include/swift/SIL/SILCloner.h

@@ -817,10 +817,9 @@ SILCloner<ImplClass>::visitAllocRefInst(AllocRefInst *Inst) {
   for (SILType OrigElemType : Inst->getTailAllocatedTypes()) {
     ElemTypes.push_back(getOpType(OrigElemType));
   }
-  auto *NewInst = getBuilder().createAllocRef(getOpLocation(Inst->getLoc()),
-                                      getOpType(Inst->getType()),
-                                      Inst->isObjC(), Inst->canAllocOnStack(),
-                                      ElemTypes, CountArgs);
+  auto *NewInst = getBuilder().createAllocRef(
+      getOpLocation(Inst->getLoc()), getOpType(Inst->getType()), Inst->isObjC(),
+      Inst->canAllocOnStack(), Inst->isUniqueReference(), ElemTypes, CountArgs);
   recordClonedInstruction(Inst, NewInst);
 }
 

include/swift/SIL/SILInstruction.h

@@ -1510,12 +1510,9 @@ class AllocStackInst final
 /// elements, the remaining operands are opened archetype operands.
 class AllocRefInstBase : public AllocationInst {
 protected:
-
-  AllocRefInstBase(SILInstructionKind Kind,
-                   SILDebugLocation DebugLoc,
-                   SILType ObjectType,
-                   bool objc, bool canBeOnStack,
-                   ArrayRef<SILType> ElementTypes);
+  AllocRefInstBase(SILInstructionKind Kind, SILDebugLocation DebugLoc,
+                   SILType ObjectType, bool objc, bool canBeOnStack,
+                   bool isUnique, ArrayRef<SILType> ElementTypes);
 
   SILType *getTypeStorage();
   const SILType *getTypeStorage() const {
@@ -1535,6 +1532,14 @@ class AllocRefInstBase : public AllocationInst {
     SILInstruction::Bits.AllocRefInstBase.OnStack = OnStack;
   }
 
+  bool isUniqueReference() const {
+    return SILInstruction::Bits.AllocRefInstBase.uniqueReference;
+  }
+
+  void setUniqueReference(bool isUnique = true) {
+    SILInstruction::Bits.AllocRefInstBase.uniqueReference = isUnique;
+  }
+
   ArrayRef<SILType> getTailAllocatedTypes() const {
     return {getTypeStorage(), getNumTailTypes()};
   }
@@ -1573,22 +1578,20 @@ class AllocRefInst final
   friend AllocRefInstBase;
   friend SILBuilder;
 
-  AllocRefInst(SILDebugLocation DebugLoc, SILFunction &F,
-               SILType ObjectType,
-               bool objc, bool canBeOnStack,
-               ArrayRef<SILType> ElementTypes,
-               ArrayRef<SILValue> AllOperands)
+  AllocRefInst(SILDebugLocation DebugLoc, SILFunction &F, SILType ObjectType,
+               bool objc, bool canBeOnStack, bool isUnique,
+               ArrayRef<SILType> ElementTypes, ArrayRef<SILValue> AllOperands)
       : InstructionBaseWithTrailingOperands(AllOperands, DebugLoc, ObjectType,
-                        objc, canBeOnStack, ElementTypes) {
+                                            objc, canBeOnStack, isUnique,
+                                            ElementTypes) {
     assert(AllOperands.size() >= ElementTypes.size());
     std::uninitialized_copy(ElementTypes.begin(), ElementTypes.end(),
                             getTrailingObjects<SILType>());
   }
 
   static AllocRefInst *create(SILDebugLocation DebugLoc, SILFunction &F,
-                              SILType ObjectType,
-                              bool objc, bool canBeOnStack,
-                              ArrayRef<SILType> ElementTypes,
+                              SILType ObjectType, bool objc, bool canBeOnStack,
+                              bool isUnique, ArrayRef<SILType> ElementTypes,
                               ArrayRef<SILValue> ElementCountOperands,
                               SILOpenedArchetypesState &OpenedArchetypes);
 
@@ -1616,13 +1619,11 @@ class AllocRefDynamicInst final
   friend AllocRefInstBase;
   friend SILBuilder;
 
-  AllocRefDynamicInst(SILDebugLocation DebugLoc,
-                      SILType ty,
-                      bool objc,
+  AllocRefDynamicInst(SILDebugLocation DebugLoc, SILType ty, bool objc,
                       ArrayRef<SILType> ElementTypes,
                       ArrayRef<SILValue> AllOperands)
       : InstructionBaseWithTrailingOperands(AllOperands, DebugLoc, ty, objc,
-                                            false, ElementTypes) {
+                                            false, false, ElementTypes) {
     assert(AllOperands.size() >= ElementTypes.size() + 1);
     std::uninitialized_copy(ElementTypes.begin(), ElementTypes.end(),
                             getTrailingObjects<SILType>());

include/swift/SIL/SILNode.h

@@ -201,12 +201,11 @@ class alignas(8) SILNode {
     NumOperands : 32-NumAllocationInstBits,
     VarInfo : 32
   );
-  IBWTO_BITFIELD(AllocRefInstBase, AllocationInst, 32-NumAllocationInstBits,
-    ObjC : 1,
-    OnStack : 1,
-    NumTailTypes : 32-1-1-NumAllocationInstBits
-  );
-  static_assert(32-1-1-NumAllocationInstBits >= 16, "Reconsider bitfield use?");
+  IBWTO_BITFIELD(AllocRefInstBase, AllocationInst, 32 - NumAllocationInstBits,
+                 ObjC : 1, OnStack : 1, uniqueReference : 1,
+                 NumTailTypes : 32 - 1 - 1 - 1 - NumAllocationInstBits);
+  static_assert(32 - 1 - 1 - 1 - NumAllocationInstBits >= 16,
+                "Reconsider bitfield use?");
 
   UIWTDOB_BITFIELD_EMPTY(AllocValueBufferInst, AllocationInst);
 

include/swift/SILOptimizer/PassManager/Passes.def

@@ -279,6 +279,8 @@ IRGEN_PASS(LoadableByAddress, "loadable-address",
      "SIL Large Loadable type by-address lowering.")
 PASS(MandatorySILLinker, "mandatory-linker",
      "Deserialize all referenced SIL functions that are shared or transparent")
+PASS(MarkReferenceUnique, "mark-reference-unique",
+     "If possible, mark alloc_ref instruction with 'unique' flag")
 PASS(PerformanceSILLinker, "performance-linker",
      "Deserialize all referenced SIL functions")
 PASS(RawSILInstLowering, "raw-sil-inst-lowering",

lib/SIL/IR/SILInstructions.cpp

@@ -200,22 +200,22 @@ DeallocStackInst *AllocStackInst::getSingleDeallocStack() const {
 }
 
 AllocRefInstBase::AllocRefInstBase(SILInstructionKind Kind,
-                                   SILDebugLocation Loc,
-                                   SILType ObjectType,
-                                   bool objc, bool canBeOnStack,
+                                   SILDebugLocation Loc, SILType ObjectType,
+                                   bool objc, bool canBeOnStack, bool isUnique,
                                    ArrayRef<SILType> ElementTypes)
     : AllocationInst(Kind, Loc, ObjectType) {
   SILInstruction::Bits.AllocRefInstBase.ObjC = objc;
   SILInstruction::Bits.AllocRefInstBase.OnStack = canBeOnStack;
+  SILInstruction::Bits.AllocRefInstBase.uniqueReference = isUnique;
   SILInstruction::Bits.AllocRefInstBase.NumTailTypes = ElementTypes.size();
   assert(SILInstruction::Bits.AllocRefInstBase.NumTailTypes ==
          ElementTypes.size() && "Truncation");
   assert(!objc || ElementTypes.empty());
 }
 
 AllocRefInst *AllocRefInst::create(SILDebugLocation Loc, SILFunction &F,
-                                   SILType ObjectType,
-                                   bool objc, bool canBeOnStack,
+                                   SILType ObjectType, bool objc,
+                                   bool canBeOnStack, bool isUnique,
                                    ArrayRef<SILType> ElementTypes,
                                    ArrayRef<SILValue> ElementCountOperands,
                                    SILOpenedArchetypesState &OpenedArchetypes) {
@@ -233,7 +233,7 @@ AllocRefInst *AllocRefInst::create(SILDebugLocation Loc, SILFunction &F,
                                                         ElementTypes.size());
   auto Buffer = F.getModule().allocateInst(Size, alignof(AllocRefInst));
   return ::new (Buffer) AllocRefInst(Loc, F, ObjectType, objc, canBeOnStack,
-                                     ElementTypes, AllOperands);
+                                     isUnique, ElementTypes, AllOperands);
 }
 
 AllocRefDynamicInst *

lib/SIL/IR/SILPrinter.cpp

@@ -1153,6 +1153,8 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
       *this << "[objc] ";
     if (ARI->canAllocOnStack())
       *this << "[stack] ";
+    if (ARI->isUniqueReference())
+      *this << "[unique] ";
     auto Types = ARI->getTailAllocatedTypes();
     auto Counts = ARI->getTailAllocatedCounts();
     for (unsigned Idx = 0, NumTypes = Types.size(); Idx < NumTypes; ++Idx) {

lib/SIL/Parser/ParseSIL.cpp

@@ -3862,6 +3862,7 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
   case SILInstructionKind::AllocRefDynamicInst: {
     bool IsObjC = false;
     bool OnStack = false;
+    bool isUnique = false;
     SmallVector<SILType, 2> ElementTypes;
     SmallVector<SILValue, 2> ElementCounts;
     while (P.consumeIf(tok::l_square)) {
@@ -3872,6 +3873,8 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
         IsObjC = true;
       } else if (Optional == "stack") {
         OnStack = true;
+      } else if (Optional == "unique") {
+        isUnique = true;
       } else if (Optional == "tail_elems") {
         SILType ElemTy;
         if (parseSILType(ElemTy) || !P.Tok.isAnyOperator() ||
@@ -3916,7 +3919,7 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
                                           ElementTypes, ElementCounts);
     } else {
       ResultVal = B.createAllocRef(InstLoc, ObjectType, IsObjC, OnStack,
-                                   ElementTypes, ElementCounts);
+                                   isUnique, ElementTypes, ElementCounts);
     }
     break;
   }

lib/SIL/Utils/InstructionUtils.cpp

@@ -653,3 +653,138 @@ swift::getStaticOverloadForSpecializedPolymorphicBuiltin(BuiltinInst *bi) {
 
   return newBI;
 }
+
+// Analyzing the body of this class destructor is valid because the object is
+// dead. This means that the object is never passed to objc_setAssociatedObject,
+// so its destructor cannot be extended at runtime.
+SILFunction *swift::getDestructor(AllocRefInstBase *ARI) {
+  // We only support classes.
+  ClassDecl *ClsDecl = ARI->getType().getClassOrBoundGenericClass();
+  if (!ClsDecl)
+    return nullptr;
+
+  // Look up the destructor of ClsDecl.
+  DestructorDecl *Destructor = ClsDecl->getDestructor();
+  assert(Destructor && "getDestructor() should never return a nullptr.");
+
+  // Find the destructor name via SILDeclRef.
+  // FIXME: When destructors get moved into vtables, update this to use the
+  // vtable for the class.
+  SILDeclRef Ref(Destructor);
+  SILFunction *Fn = ARI->getModule().lookUpFunction(Ref);
+  if (!Fn || Fn->empty()) {
+    LLVM_DEBUG(llvm::dbgs() << "    Could not find destructor.\n");
+    return nullptr;
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "    Found destructor!\n");
+
+  // If the destructor has an objc_method calling convention, we cannot
+  // analyze it since it could be swapped out from under us at runtime.
+  if (Fn->getRepresentation() == SILFunctionTypeRepresentation::ObjCMethod) {
+    LLVM_DEBUG(llvm::dbgs() << "        Found Objective-C destructor. Can't "
+                               "analyze!\n");
+    return nullptr;
+  }
+
+  return Fn;
+}
+
+Optional<StringRef> swift::isReferenceUnique(AllocRefInstBase *ref,
+                                             bool verbose) {
+  SmallPtrSet<Operand *, 8> seenUses;
+  SmallVector<Operand *, 8> usesToCheck;
+  usesToCheck.append(ref->use_begin(), ref->use_end());
+
+  while (!usesToCheck.empty()) {
+    auto *use = usesToCheck.pop_back_val();
+    auto *user = use->getUser();
+
+    // Keep track of the operands that we've already seen to prevent recursion.
+    // If we've already seen this operand, continue.
+    if (!seenUses.insert(use).second)
+      continue;
+
+    if (auto br = dyn_cast<BranchInst>(user)) {
+      unsigned i = 0;
+      for (auto arg : br->getArgs()) {
+        if (arg == use->get())
+          break;
+        i++;
+      }
+      auto arg = br->getDestBB()->getArgument(i);
+      usesToCheck.append(arg->use_begin(), arg->use_end());
+      continue;
+    }
+
+    if (auto apply = FullApplySite::isa(user)) {
+      if (!apply.getReferencedFunctionOrNull()) {
+        if (verbose) {
+          llvm::dbgs() << "Could not find the referenced function of the "
+                          "following full apply site: \n";
+          apply.getInstruction()->print(llvm::dbgs());
+        }
+        return {"Found unkown use of unique reference."};
+      }
+
+      auto referencedFn = apply.getReferencedFunctionOrNull();
+      if (referencedFn->empty()) {
+        if (verbose) {
+          llvm::dbgs() << "Cannot see the body of the referenced function: \n";
+          referencedFn->print(llvm::dbgs());
+        }
+        return {"Found unkown use of unique reference."};
+      }
+
+      unsigned i = 0;
+      for (auto arg : apply.getArguments()) {
+        if (arg == use->get())
+          break;
+        i++;
+      }
+      auto arg = referencedFn->getArgument(i);
+      usesToCheck.append(arg->use_begin(), arg->use_end());
+      continue;
+    }
+
+    if (isa<StrongReleaseInst>(user)) {
+      auto *destructor = getDestructor(ref);
+      if (!destructor)
+        return {"Cannot find destructor for reference type while checking "
+                "strong_release does not escape the unique reference."};
+      auto *selfInDestructor = destructor->begin()->getArgument(0);
+      usesToCheck.append(selfInDestructor->use_begin(),
+                         selfInDestructor->use_end());
+      continue;
+    }
+
+    // These instructions are OK.
+    if (isa<LoadInst>(user) || isa<DeallocRefInst>(user) ||
+        isa<EndAccessInst>(user) || isa<SetDeallocatingInst>(user) ||
+        isa<DebugValueInst>(user) || isa<DebugValueAddrInst>(user) ||
+        isa<ClassMethodInst>(user) || isa<StrongRetainInst>(user))
+      continue;
+
+    if (auto *store = dyn_cast<StoreInst>(user)) {
+      if (store->getDest() != use->get())
+        return {"Store escapes address. Store instruction must store into "
+                "the unique reference."};
+      continue;
+    }
+
+    if (isa<BeginAccessInst>(user) || isa<RefElementAddrInst>(user) ||
+        isa<StructElementAddrInst>(user) || isa<TupleElementAddrInst>(user)) {
+      auto *userVal = cast<SingleValueInstruction>(user);
+      usesToCheck.append(userVal->use_begin(), userVal->use_end());
+      continue;
+    }
+
+    if (verbose) {
+      llvm::dbgs() << "The following instruction may escape the reference: \n";
+      user->print(llvm::dbgs());
+    }
+    return {"Found unkown use of unique reference."};
+  }
+
+  return None;
+}