[WIP] Visitor for AccessedStorage use-def chain walk

include/swift/SIL/MemAccessUtils.h

@@ -36,9 +36,11 @@
 #ifndef SWIFT_SIL_MEMACCESSUTILS_H
 #define SWIFT_SIL_MEMACCESSUTILS_H
 
+#include "swift/SIL/ApplySite.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILGlobalVariable.h"
 #include "swift/SIL/SILInstruction.h"
 #include "llvm/ADT/DenseMap.h"
 
@@ -117,10 +119,6 @@ class AccessedStorage {
 
   static const char *getKindName(Kind k);
 
-  /// If the given address source is an identified access base, return the kind
-  /// of access base. Otherwise, return Unidentified.
-  static AccessedStorage::Kind classify(SILValue base);
-
   /// Directly create an AccessedStorage for class property access.
   static AccessedStorage forClass(SILValue object, unsigned propertyIndex) {
     AccessedStorage storage;
@@ -406,6 +404,58 @@ template <> struct DenseMapInfo<swift::AccessedStorage> {
 
 namespace swift {
 
+/// Abstract CRTP class for a visitor passed to \c visitAccessUseDefChain.
+template<typename Impl, typename Result = void>
+class AccessUseDefChainVisitor {
+protected:
+  Impl &asImpl() {
+    return static_cast<Impl&>(*this);
+  }
+public:
+  // Subclasses can provide a method for any identified access base:
+  // Result visitBase(SILValue base, AccessedStorage::Kind kind);
+
+  // Visitors for specific identified access kinds. These default to calling out
+  // to visitIdentified.
+
+  Result visitClassAccess(RefElementAddrInst *field) {
+    return asImpl().visitBase(field, AccessedStorage::Class);
+  }
+  Result visitArgumentAccess(SILFunctionArgument *arg) {
+    return asImpl().visitBase(arg, AccessedStorage::Argument);
+  }
+  Result visitBoxAccess(AllocBoxInst *box) {
+    return asImpl().visitBase(box, AccessedStorage::Box);
+  }
+  /// The argument may be either a GlobalAddrInst or the ApplyInst for a global accessor function.
+  Result visitGlobalAccess(SILValue global) {
+    return asImpl().visitBase(global, AccessedStorage::Global);
+  }
+  Result visitYieldAccess(BeginApplyResult *yield) {
+    return asImpl().visitBase(yield, AccessedStorage::Yield);
+  }
+  Result visitStackAccess(AllocStackInst *stack) {
+    return asImpl().visitBase(stack, AccessedStorage::Stack);
+  }
+  Result visitNestedAccess(BeginAccessInst *access) {
+    return asImpl().visitBase(access, AccessedStorage::Nested);
+  }
+
+  // Visitors for unidentified base values.
+
+  Result visitUnidentified(SILValue base) {
+    return asImpl().visitBase(base, AccessedStorage::Unidentified);
+  }
+
+  // Subclasses must provide implementations to visit non-access bases
+  // and phi arguments, and for incomplete projections from the access:
+  // void visitNonAccess(SILValue base);
+  // void visitPhi(SILPhiArgument *phi);
+  // void visitIncomplete(SILValue projectedAddr, SILValue parentAddr);
+
+  Result visit(SILValue sourceAddr);
+};
+
 /// Given an address accessed by an instruction that reads or modifies
 /// memory, return an AccessedStorage object that identifies the formal access.
 ///
@@ -468,6 +518,192 @@ void visitAccessedAddress(SILInstruction *I,
 /// instruction following this begin_access was not also erased.
 SILBasicBlock::iterator removeBeginAccess(BeginAccessInst *beginAccess);
 
+/// Return true if the given address value is produced by a special address
+/// producer that is only used for local initialization, not formal access.
+bool isAddressForLocalInitOnly(SILValue sourceAddr);
+/// Return true if the given apply invokes a global addressor defined in another
+/// module.
+bool isExternalGlobalAddressor(ApplyInst *AI);
+/// Return true if the given StructExtractInst extracts the RawPointer from
+/// Unsafe[Mutable]Pointer.
+bool isUnsafePointerExtraction(StructExtractInst *SEI);
+/// Given a block argument address base, check if it is actually a box projected
+/// from a switch_enum. This is a valid pattern at any SIL stage resulting in a
+/// block-type phi. In later SIL stages, the optimizer may form address-type
+/// phis, causing this assert if called on those cases.
+void checkSwitchEnumBlockArg(SILPhiArgument *arg);
+
+template<typename Impl, typename Result>
+Result AccessUseDefChainVisitor<Impl, Result>::visit(SILValue sourceAddr) {
+  // Handle immediately-identifiable instructions.
+  while (true) {
+    switch (sourceAddr->getKind()) {
+    // An AllocBox is a fully identified memory location.
+    case ValueKind::AllocBoxInst:
+      return asImpl().visitBoxAccess(cast<AllocBoxInst>(sourceAddr));
+    // An AllocStack is a fully identified memory location, which may occur
+    // after inlining code already subjected to stack promotion.
+    case ValueKind::AllocStackInst:
+      return asImpl().visitStackAccess(cast<AllocStackInst>(sourceAddr));
+    case ValueKind::GlobalAddrInst:
+      return asImpl().visitGlobalAccess(sourceAddr);
+    case ValueKind::ApplyInst: {
+      FullApplySite apply(cast<ApplyInst>(sourceAddr));
+      if (auto *funcRef = apply.getReferencedFunctionOrNull()) {
+        if (getVariableOfGlobalInit(funcRef)) {
+          return asImpl().visitGlobalAccess(sourceAddr);
+        }
+      }
+      // Try to classify further below.
+      break;
+    }
+    case ValueKind::RefElementAddrInst:
+      return asImpl().visitClassAccess(cast<RefElementAddrInst>(sourceAddr));
+    // A yield is effectively a nested access, enforced independently in
+    // the caller and callee.
+    case ValueKind::BeginApplyResult:
+      return asImpl().visitYieldAccess(cast<BeginApplyResult>(sourceAddr));
+    // A function argument is effectively a nested access, enforced
+    // independently in the caller and callee.
+    case ValueKind::SILFunctionArgument:
+      return asImpl().visitArgumentAccess(cast<SILFunctionArgument>(sourceAddr));
+    // View the outer begin_access as a separate location because nested
+    // accesses do not conflict with each other.
+    case ValueKind::BeginAccessInst:
+      return asImpl().visitNestedAccess(cast<BeginAccessInst>(sourceAddr));
+    default:
+      // Try to classify further below.
+      break;
+    }
+
+    // If the sourceAddr producer cannot immediately be classified, follow the
+    // use-def chain of sourceAddr, box, or RawPointer producers.
+    assert(sourceAddr->getType().isAddress()
+           || isa<SILBoxType>(sourceAddr->getType().getASTType())
+           || isa<BuiltinRawPointerType>(sourceAddr->getType().getASTType()));
+
+    // Handle other unidentified address sources.
+    switch (sourceAddr->getKind()) {
+    default:
+      if (isAddressForLocalInitOnly(sourceAddr))
+        return asImpl().visitUnidentified(sourceAddr);
+      return asImpl().visitNonAccess(sourceAddr);
+
+    case ValueKind::SILUndef:
+      return asImpl().visitUnidentified(sourceAddr);
+
+    case ValueKind::ApplyInst:
+      if (isExternalGlobalAddressor(cast<ApplyInst>(sourceAddr)))
+        return asImpl().visitUnidentified(sourceAddr);
+
+      // Don't currently allow any other calls to return an accessed address.
+      return asImpl().visitNonAccess(sourceAddr);
+
+    case ValueKind::StructExtractInst:
+      // Handle nested access to a KeyPath projection. The projection itself
+      // uses a Builtin. However, the returned UnsafeMutablePointer may be
+      // converted to an address and accessed via an inout argument.
+      if (isUnsafePointerExtraction(cast<StructExtractInst>(sourceAddr)))
+        return asImpl().visitUnidentified(sourceAddr);
+      return asImpl().visitNonAccess(sourceAddr);
+
+    case ValueKind::SILPhiArgument: {
+      auto *phiArg = cast<SILPhiArgument>(sourceAddr);
+      if (phiArg->isPhiArgument()) {
+        return asImpl().visitPhi(phiArg);
+      }
+
+      // A non-phi block argument may be a box value projected out of
+      // switch_enum. Address-type block arguments are not allowed.
+      if (sourceAddr->getType().isAddress())
+        return asImpl().visitNonAccess(sourceAddr);
+
+      checkSwitchEnumBlockArg(cast<SILPhiArgument>(sourceAddr));
+      return asImpl().visitUnidentified(sourceAddr);
+    }
+    // Load a box from an indirect payload of an opaque enum.
+    // We must have peeked past the project_box earlier in this loop.
+    // (the indirectness makes it a box, the load is for address-only).
+    //
+    // %payload_adr = unchecked_take_enum_data_addr %enum : $*Enum, #Enum.case
+    // %box = load [take] %payload_adr : $*{ var Enum }
+    //
+    // FIXME: this case should go away with opaque values.
+    //
+    // Otherwise return invalid AccessedStorage.
+    case ValueKind::LoadInst:
+      if (sourceAddr->getType().is<SILBoxType>()) {
+        SILValue operAddr = cast<LoadInst>(sourceAddr)->getOperand();
+        assert(isa<UncheckedTakeEnumDataAddrInst>(operAddr));
+        return asImpl().visitIncomplete(sourceAddr, operAddr);
+      }
+      return asImpl().visitNonAccess(sourceAddr);
+
+    // ref_tail_addr project an address from a reference.
+    // This is a valid address producer for nested @inout argument
+    // access, but it is never used for formal access of identified objects.
+    case ValueKind::RefTailAddrInst:
+      return asImpl().visitUnidentified(sourceAddr);
+
+    // Inductive single-operand cases:
+    // Look through address casts to find the source address.
+    case ValueKind::MarkUninitializedInst:
+    case ValueKind::OpenExistentialAddrInst:
+    case ValueKind::UncheckedAddrCastInst:
+    // Inductive cases that apply to any type.
+    case ValueKind::CopyValueInst:
+    case ValueKind::MarkDependenceInst:
+    // Look through a project_box to identify the underlying alloc_box as the
+    // accesed object. It must be possible to reach either the alloc_box or the
+    // containing enum in this loop, only looking through simple value
+    // propagation such as copy_value.
+    case ValueKind::ProjectBoxInst:
+    // Handle project_block_storage just like project_box.
+    case ValueKind::ProjectBlockStorageInst:
+    // Look through begin_borrow in case a local box is borrowed.
+    case ValueKind::BeginBorrowInst:
+      return asImpl().visitIncomplete(sourceAddr,
+        cast<SingleValueInstruction>(sourceAddr)->getOperand(0));
+
+    // Access to a Builtin.RawPointer. Treat this like the inductive cases
+    // above because some RawPointers originate from identified locations. See
+    // the special case for global addressors, which return RawPointer,
+    // above. AddressToPointer is also handled because it results from inlining a
+    // global addressor without folding the AddressToPointer->PointerToAddress.
+    //
+    // If the inductive search does not find a valid addressor, it will
+    // eventually reach the default case that returns in invalid location. This
+    // is correct for RawPointer because, although accessing a RawPointer is
+    // legal SIL, there is no way to guarantee that it doesn't access class or
+    // global storage, so returning a valid unidentified storage object would be
+    // incorrect. It is the caller's responsibility to know that formal access
+    // to such a location can be safely ignored.
+    //
+    // For example:
+    //
+    // - KeyPath Builtins access RawPointer. However, the caller can check
+    // that the access `isFromBuilin` and ignore the storage.
+    //
+    // - lldb generates RawPointer access for debugger variables, but SILGen
+    // marks debug VarDecl access as 'Unsafe' and SIL passes don't need the
+    // AccessedStorage for 'Unsafe' access.
+    case ValueKind::PointerToAddressInst:
+    case ValueKind::AddressToPointerInst:
+      return asImpl().visitIncomplete(sourceAddr,
+        cast<SingleValueInstruction>(sourceAddr)->getOperand(0));
+
+    // Address-to-address subobject projections.
+    case ValueKind::StructElementAddrInst:
+    case ValueKind::TupleElementAddrInst:
+    case ValueKind::UncheckedTakeEnumDataAddrInst:
+    case ValueKind::TailAddrInst:
+    case ValueKind::IndexAddrInst:
+      return asImpl().visitIncomplete(sourceAddr,
+        cast<SingleValueInstruction>(sourceAddr)->getOperand(0));
+    }
+  };
+}
+
 } // end namespace swift
 
 #endif