From 12a4d471c08c7160cc74fa413e289832cf13f696 Mon Sep 17 00:00:00 2001
From: John McCall <rjmccall@apple.com>
Date: Fri, 7 May 2021 18:44:15 -0400
Subject: [PATCH] In DI, cache whether a memory object is a box.

Boxes tend to have a small number of uses, so frequently finding the
unique projection isn't too bad.  Non-boxes, however, can have a large
number of uses: for example, a class instance has expected uses
proportionate to the number of stored properties.  So if we do a linear
scan of the uses on a non-box instruction, we'll scale quadratically.

Fixes SR-14532.
---
 .../Mandatory/DIMemoryUseCollector.cpp            | 15 ++++++++-------
 lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h |  9 ++++++++-
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp b/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp
index 3e113b40d1b04..66654a01e23c0 100644
--- a/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp
+++ b/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp
@@ -93,13 +93,8 @@ static unsigned getElementCountRec(TypeExpansionContext context,
 }
 
 static std::pair<SILType, bool>
-computeMemorySILType(MarkUninitializedInst *MemoryInst) {
+computeMemorySILType(MarkUninitializedInst *MUI, SILValue Address) {
   // Compute the type of the memory object.
-  auto *MUI = MemoryInst;
-  SILValue Address = MUI;
-  if (auto *PBI = Address->getSingleUserOfType<ProjectBoxInst>()) {
-    Address = PBI;
-  }
   SILType MemorySILType = Address->getType().getObjectType();
 
   // If this is a let variable we're initializing, remember this so we don't
@@ -118,7 +113,13 @@ DIMemoryObjectInfo::DIMemoryObjectInfo(MarkUninitializedInst *MI)
     : MemoryInst(MI) {
   auto &Module = MI->getModule();
 
-  std::tie(MemorySILType, IsLet) = computeMemorySILType(MemoryInst);
+  SILValue Address = MemoryInst;
+  if (auto PBI = MemoryInst->getSingleUserOfType<ProjectBoxInst>()) {
+    IsBox = true;
+    Address = PBI;
+  }
+
+  std::tie(MemorySILType, IsLet) = computeMemorySILType(MI, Address);
 
   // Compute the number of elements to track in this memory object.
   // If this is a 'self' in a delegating initializer, we only track one bit:
diff --git a/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h b/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h
index 88ebcd5f0b537..f8821ae19b654 100644
--- a/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h
+++ b/lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h
@@ -72,6 +72,9 @@ class DIMemoryObjectInfo {
   /// non-empty.
   bool HasDummyElement = false;
 
+  /// True if this object has a single user of type ProjectBoxInst.
+  bool IsBox = false;
+
 public:
   DIMemoryObjectInfo(MarkUninitializedInst *MemoryInst);
 
@@ -98,8 +101,12 @@ class DIMemoryObjectInfo {
   /// instruction. For alloc_box though it returns the project_box associated
   /// with the memory info.
   SingleValueInstruction *getUninitializedValue() const {
-    if (auto *pbi = MemoryInst->getSingleUserOfType<ProjectBoxInst>())
+    if (IsBox) {
+      // TODO: consider just storing the ProjectBoxInst in this case.
+      auto *pbi = MemoryInst->getSingleUserOfType<ProjectBoxInst>();
+      assert(pbi);
       return pbi;
+    }
     return MemoryInst;
   }