From 653a5aae2a265d09e6dc6e0c71043ff764d27e81 Mon Sep 17 00:00:00 2001
From: Andrew Trick <atrick@apple.com>
Date: Tue, 3 Mar 2020 22:40:41 -0800
Subject: [PATCH] Prevent begin_borrow on addresses and cleanup address access
 utils.

---
 include/swift/SIL/MemAccessUtils.h |  8 ++++++
 include/swift/SIL/SILBuilder.h     |  1 +
 lib/SIL/MemAccessUtils.cpp         | 44 ++++++++++++++++--------------
 3 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/include/swift/SIL/MemAccessUtils.h b/include/swift/SIL/MemAccessUtils.h
index 1d2ce42b3f6..f54966525a0 100644
--- a/include/swift/SIL/MemAccessUtils.h
+++ b/include/swift/SIL/MemAccessUtils.h
@@ -53,6 +53,14 @@ namespace swift {
 /// (pointer-to-address is not stripped).
 SILValue stripAccessMarkers(SILValue v);
 
+/// Return a non-null address-type SingleValueInstruction if \p v is the result
+/// of an address projection that may be inside of a formal access, such as
+/// (begin_borrow, struct_element_addr, tuple_element_addr).
+///
+/// The resulting projection must have an address-type operand at index zero
+/// representing the projected address.
+SingleValueInstruction *isAccessProjection(SILValue v);
+
 /// Attempt to return the address corresponding to a variable's formal access
 /// by stripping indexing and address projections.
 ///
diff --git a/include/swift/SIL/SILBuilder.h b/include/swift/SIL/SILBuilder.h
index ed1e3b7a033..6401020d377 100644
--- a/include/swift/SIL/SILBuilder.h
+++ b/include/swift/SIL/SILBuilder.h
@@ -728,6 +728,7 @@ public:
   }
 
   BeginBorrowInst *createBeginBorrow(SILLocation Loc, SILValue LV) {
+    assert(!LV->getType().isAddress());
     return insert(new (getModule())
                       BeginBorrowInst(getSILDebugLocation(Loc), LV));
   }
diff --git a/lib/SIL/MemAccessUtils.cpp b/lib/SIL/MemAccessUtils.cpp
index 92cb96801d1..f5937e12683 100644
--- a/lib/SIL/MemAccessUtils.cpp
+++ b/lib/SIL/MemAccessUtils.cpp
@@ -22,37 +22,39 @@
 using namespace swift;
 
 SILValue swift::stripAccessMarkers(SILValue v) {
-  while (true) {
-    switch (v->getKind()) {
-    default:
-      return v;
-
-    case ValueKind::BeginBorrowInst:
-    case ValueKind::BeginAccessInst:
-      v = cast<SingleValueInstruction>(v)->getOperand(0);
-    }
+  while (auto *bai = dyn_cast<BeginAccessInst>(v)) {
+    v = bai->getOperand();
   }
+  return v;
+}
+
+// The resulting projection must have an address-type operand at index zero
+// representing the projected address.
+SingleValueInstruction *swift::isAccessProjection(SILValue v) {
+  switch (v->getKind()) {
+  default:
+    return nullptr;
+
+  case ValueKind::StructElementAddrInst:
+  case ValueKind::TupleElementAddrInst:
+  case ValueKind::UncheckedTakeEnumDataAddrInst:
+  case ValueKind::TailAddrInst:
+  case ValueKind::IndexAddrInst:
+    return cast<SingleValueInstruction>(v);
+  };
 }
 
 // TODO: When the optimizer stops stripping begin_access markers, then we should
 // be able to assert that the result is a BeginAccessInst and the default case
 // is unreachable.
 SILValue swift::getAccessedAddress(SILValue v) {
-  assert(v->getType().isAddress());
   while (true) {
-    switch (v->getKind()) {
-    default:
+    assert(v->getType().isAddress());
+    auto *projection = isAccessProjection(v);
+    if (!projection)
       return v;
 
-    case ValueKind::BeginBorrowInst:
-    case ValueKind::StructElementAddrInst:
-    case ValueKind::TupleElementAddrInst:
-    case ValueKind::UncheckedTakeEnumDataAddrInst:
-    case ValueKind::TailAddrInst:
-    case ValueKind::IndexAddrInst:
-      v = cast<SingleValueInstruction>(v)->getOperand(0);
-      continue;
-    };
+    v = projection->getOperand(0);
   }
 }