Add AccessedStorageWithBase to conviently recover the base's VarDecl

Author: atrick

include/swift/SIL/MemAccessUtils.h

@@ -507,7 +507,7 @@ class AccessedStorage {
     llvm_unreachable("unhandled kind");
   }
 
-  /// Return trye if the given access is guaranteed to be within a heap object.
+  /// Return true if the given access is guaranteed to be within a heap object.
   bool isObjectAccess() const {
     return getKind() == Class || getKind() == Tail;
   }
@@ -696,6 +696,28 @@ template <> struct DenseMapInfo<swift::AccessedStorage> {
 
 namespace swift {
 
+/// For convenience, encapsulate and AccessedStorage value along with its
+/// accessed base address.
+struct AccessedStorageWithBase {
+  AccessedStorage storage;
+  // The base of the formal access. For class storage, it is the
+  // ref_element_addr. For global storage it is the global_addr or initializer
+  // apply. For other storage, it is the same as accessPath.getRoot().
+  //
+  // Base may be invalid for global_addr -> address_to_pointer -> phi patterns.
+  // FIXME: add a structural requirement to SIL so base is always valid in OSSA.
+  SILValue base;
+
+  AccessedStorageWithBase(AccessedStorage storage, SILValue base)
+      : storage(storage), base(base) {}
+
+  /// Identical to AccessedStorage::compute but preserves the access base.
+  static AccessedStorageWithBase compute(SILValue sourceAddress);
+
+  /// Identical to AccessedStorage::computeInScope but preserves the base.
+  static AccessedStorageWithBase computeInScope(SILValue sourceAddress);
+};
+
 /// Return an AccessedStorage value that identifies formally accessed storage
 /// for \p beginAccess, considering any outer access scope as having distinct
 /// storage from this access scope. This is useful for exclusivity checking

lib/SIL/Utils/MemAccessUtils.cpp

@@ -714,16 +714,26 @@ class FindAccessedStorageVisitor
 
 } // end anonymous namespace
 
-AccessedStorage AccessedStorage::compute(SILValue sourceAddress) {
+AccessedStorageWithBase
+AccessedStorageWithBase::compute(SILValue sourceAddress) {
   FindAccessedStorageVisitor visitor(NestedAccessType::IgnoreAccessBegin);
   visitor.findStorage(sourceAddress);
-  return visitor.getStorage();
+  return {visitor.getStorage(), visitor.getBase()};
 }
 
-AccessedStorage AccessedStorage::computeInScope(SILValue sourceAddress) {
+AccessedStorageWithBase
+AccessedStorageWithBase::computeInScope(SILValue sourceAddress) {
   FindAccessedStorageVisitor visitor(NestedAccessType::StopAtAccessBegin);
   visitor.findStorage(sourceAddress);
-  return visitor.getStorage();
+  return {visitor.getStorage(), visitor.getBase()};
+}
+
+AccessedStorage AccessedStorage::compute(SILValue sourceAddress) {
+  return AccessedStorageWithBase::compute(sourceAddress).storage;
+}
+
+AccessedStorage AccessedStorage::computeInScope(SILValue sourceAddress) {
+  return AccessedStorageWithBase::computeInScope(sourceAddress).storage;
 }
 
 //===----------------------------------------------------------------------===//

lib/SIL/Verifier/LoadBorrowInvalidationChecker.cpp

@@ -454,11 +454,11 @@ bool LoadBorrowNeverInvalidatedAnalysis::
 bool LoadBorrowNeverInvalidatedAnalysis::isNeverInvalidated(
     LoadBorrowInst *lbi) {
 
-  SILValue address = getAccessBegin(lbi->getOperand());
+  SILValue address = getAccessScope(lbi->getOperand());
   if (!address)
     return false;
 
-  auto storage = findAccessedStorage(address);
+  auto storage = AccessedStorage::compute(address);
   // If we couldn't find an access storage, return that we are assumed to write.
   if (!storage) {
     llvm::errs() << "Couldn't compute access storage?!\n";

lib/SILOptimizer/Transforms/AccessEnforcementWMO.cpp

@@ -67,22 +67,21 @@ using llvm::DenseMap;
 using llvm::SmallDenseSet;
 
 // Get the VarDecl that represents the DisjointAccessLocation for the given
-// AccessedStorage. Returns nullptr for any storage that can't be partitioned
-// into a disjoint location.
+// storage and access base. Returns nullptr for any storage that can't be
+// partitioned into a disjoint location.
 //
-// identifyFormalAccess may only return Unidentified storage for a global
-// variable access if the global is defined in a different module.
-//
-// WARNING: Retrieving VarDecl for Class access is not constant time.
-const VarDecl *getDisjointAccessLocation(const AccessedStorage &storage) {
+// Global storage is expected to be disjoint because identifyFormalAccess may
+// only return Unidentified storage for a global variable access if the global
+// is defined in a different module.
+const VarDecl *
+getDisjointAccessLocation(AccessedStorageWithBase storageAndBase) {
+  auto storage = storageAndBase.storage;
   switch (storage.getKind()) {
   case AccessedStorage::Global:
-    // A global variable may return a null decl. These variables are
-    // implementation details that aren't formally accessed.
-    return storage.getGlobal()->getDecl();
-  case AccessedStorage::Class: {
-    return cast<VarDecl>(storage.getDecl());
-  }
+  case AccessedStorage::Class:
+    // Class and Globals are always a VarDecl, but the global decl may have a
+    // null value for global_addr -> phi.
+    return cast_or_null<VarDecl>(storage.getDecl(storageAndBase.base));
   case AccessedStorage::Box:
   case AccessedStorage::Stack:
   case AccessedStorage::Tail:
@@ -169,15 +168,17 @@ void GlobalAccessRemoval::perform() {
 
 void GlobalAccessRemoval::visitInstruction(SILInstruction *I) {
   if (auto *BAI = dyn_cast<BeginAccessInst>(I)) {
-    auto storage = AccessedStorage::compute(BAI->getSource());
-    const VarDecl *decl = getDisjointAccessLocation(storage);
-    recordAccess(BAI, decl, storage.getKind(), BAI->hasNoNestedConflict());
+    auto storageAndBase = AccessedStorageWithBase::compute(BAI->getSource());
+    const VarDecl *decl = getDisjointAccessLocation(storageAndBase);
+    recordAccess(BAI, decl, storageAndBase.storage.getKind(),
+                 BAI->hasNoNestedConflict());
     return;
   }
   if (auto *BUAI = dyn_cast<BeginUnpairedAccessInst>(I)) {
-    auto storage = AccessedStorage::compute(BUAI->getSource());
-    const VarDecl *decl = getDisjointAccessLocation(storage);
-    recordAccess(BUAI, decl, storage.getKind(), BUAI->hasNoNestedConflict());
+    auto storageAndBase = AccessedStorageWithBase::compute(BUAI->getSource());
+    const VarDecl *decl = getDisjointAccessLocation(storageAndBase);
+    recordAccess(BUAI, decl, storageAndBase.storage.getKind(),
+                 BUAI->hasNoNestedConflict());
     return;
   }
   if (auto *KPI = dyn_cast<KeyPathInst>(I)) {