Merge pull request swiftlang#63556 from atrick/new-access-base-visitor

Add visitAccessPathBaseUses API

Author: gottesmm

include/swift/SIL/MemAccessUtils.h

@@ -1397,6 +1397,12 @@ bool visitAccessStorageUses(AccessUseVisitor &visitor, AccessStorage storage,
 bool visitAccessPathUses(AccessUseVisitor &visitor, AccessPath accessPath,
                          SILFunction *function);
 
+/// Similar to visitAccessPathUses, but the visitor is restricted to a specific
+/// access base, such as a particular ref_element_addr.
+bool visitAccessPathBaseUses(AccessUseVisitor &visitor,
+                             AccessPathWithBase accessPathWithBase,
+                             SILFunction *function);
+
 } // end namespace swift
 
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/MemAccessUtils.cpp

@@ -1534,6 +1534,16 @@ class AccessPathDefUseTraversal {
   // The origin of the def-use traversal.
   AccessStorage 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' is typically invalid, maning that all uses of 'storage' for the
+  // access path will be visited. When 'base' is set, the the visitor is
+  // restricted to a specific access base, such as a particular
+  // ref_element_addr.
+  SILValue base;
+
   // Indices of the path to match from inner to outer component.
   // A cursor is used to represent the most recently visited def.
   // During def-use traversal, the cursor starts at the end of pathIndices and
@@ -1581,6 +1591,20 @@ class AccessPathDefUseTraversal {
     });
   }
 
+  AccessPathDefUseTraversal(AccessUseVisitor &visitor,
+                            AccessPathWithBase accessPathWithBase,
+                            SILFunction *function)
+      : visitor(visitor), storage(accessPathWithBase.accessPath.getStorage()),
+        base(accessPathWithBase.getAccessBase().getBaseAddress()) {
+
+    auto accessPath = accessPathWithBase.accessPath;
+    assert(accessPath.isValid());
+
+    initializePathIndices(accessPath);
+
+    initializeDFS(base);
+  }
+
   // Return true is all uses have been visited.
   bool visitUses() {
     // Return false if initialization failed.
@@ -1601,8 +1625,8 @@ class AccessPathDefUseTraversal {
       if (phi->isPhi())
         visitedPhis.insert(phi);
     }
-    pushUsers(root,
-              DFSEntry(nullptr, storage.isReference(), pathIndices.size(), 0));
+    bool isRef = !base && storage.isReference();
+    pushUsers(root, DFSEntry(nullptr, isRef, pathIndices.size(), 0));
   }
 
   void pushUsers(SILValue def, const DFSEntry &dfs) {
@@ -1654,6 +1678,10 @@ void AccessPathDefUseTraversal::initializePathIndices(AccessPath accessPath) {
   if (int offset = accessPath.getOffset()) {
     pathIndices.push_back(AccessPath::Index::forOffset(offset));
   }
+  // If traversal starts at the base address, then class storage is irrelevant.
+  if (base)
+    return;
+
   // The search will start from the object root, not the formal access base,
   // so add the class index to the front.
   if (storage.getKind() == AccessStorage::Class) {
@@ -1966,6 +1994,13 @@ bool swift::visitAccessPathUses(AccessUseVisitor &visitor,
   return AccessPathDefUseTraversal(visitor, accessPath, function).visitUses();
 }
 
+bool swift::visitAccessPathBaseUses(AccessUseVisitor &visitor,
+                                    AccessPathWithBase accessPathWithBase,
+                                    SILFunction *function) {
+  return AccessPathDefUseTraversal(visitor, accessPathWithBase, function)
+      .visitUses();
+}
+
 bool swift::visitAccessStorageUses(AccessUseVisitor &visitor,
                                    AccessStorage storage,
                                    SILFunction *function) {

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -68,6 +68,7 @@
 
 #include "swift/AST/Type.h"
 #include "swift/Basic/TaggedUnion.h"
+#include "swift/SIL/MemAccessUtils.h"
 #include "swift/SIL/PrunedLiveness.h"
 #include "swift/SIL/SILArgumentArrayRef.h"
 #include "swift/SIL/SILBasicBlock.h"
@@ -621,6 +622,44 @@ struct SimplifyCFGTryJumpThreading : UnitTest {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// MARK: AccessPath Unit Tests
+//===----------------------------------------------------------------------===//
+
+struct AccessUseTestVisitor : public AccessUseVisitor {
+  AccessUseTestVisitor()
+    : AccessUseVisitor(AccessUseType::Overlapping,
+                       NestedAccessType::IgnoreAccessBegin) {}
+
+  bool visitUse(Operand *op, AccessUseType useTy) override {
+    switch (useTy) {
+    case AccessUseType::Exact:
+      llvm::errs() << "Exact Use: ";
+      break;
+    case AccessUseType::Inner:
+      llvm::errs() << "Inner Use: ";
+      break;
+    case AccessUseType::Overlapping:
+      llvm::errs() << "Overlapping Use ";
+      break;
+    }
+    llvm::errs() << *op->getUser();
+    return true;
+  }
+};
+
+struct AccessPathBaseTest : UnitTest {
+  AccessPathBaseTest(UnitTestRunner *pass) : UnitTest(pass) {}
+  void invoke(Arguments &arguments) override {
+    auto value = arguments.takeValue();
+    getFunction()->dump();
+    llvm::outs() << "Access path base: " << value;
+    auto accessPathWithBase = AccessPathWithBase::compute(value);
+    AccessUseTestVisitor visitor;
+    visitAccessPathBaseUses(visitor, accessPathWithBase, getFunction());
+  }
+};
+
 /// [new_tests] Add the new UnitTest subclass above this line. 
 
 //===----------------------------------------------------------------------===//
@@ -637,6 +676,7 @@ void UnitTestRunner::withTest(StringRef name, Doit doit) {
   }
 
     // Alphabetical mapping from string to unit test subclass.
+    ADD_UNIT_TEST_SUBCLASS("accesspath-base", AccessPathBaseTest)
     ADD_UNIT_TEST_SUBCLASS("canonicalize-ossa-lifetime", CanonicalizeOSSALifetimeTest)
     ADD_UNIT_TEST_SUBCLASS("canonicalize-borrow-scope",
                            CanonicalizeBorrowScopeTest)

test/SILOptimizer/accesspath_unit.sil

@@ -0,0 +1,50 @@
+// RUN: %target-sil-opt -unit-test-runner %s -o /dev/null 2>&1 | %FileCheck %s
+
+sil_stage raw
+
+import Builtin
+
+import Swift
+
+struct S {}
+
+class Klass {
+  var f: S
+}
+
+// CHECK-LABEL: begin running test 1 of 2 on testRefElement: accesspath-base with: @trace[0]
+// CHECK: [[P1:%.*]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK: [[A1:%.*]] = begin_access [read] [dynamic] [[P1]] : $*S
+// CHECK: [[P2:%.*]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK: [[A2:%.*]] = begin_access [read] [dynamic] [[P2]] : $*S
+// CHECK: Access path base:  [[P1]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK-NEXT: Exact Use:   %{{.*}} = load [trivial] [[A1]] : $*S
+// CHECK-NEXT: Exact Use:   end_access [[A1]] : $*S
+// CHECK-LABEL: end running test 1 of 2 on testRefElement: accesspath-base with: @trace[0]
+
+// CHECK-LABEL: begin running test 2 of 2 on testRefElement: accesspath-base with: @trace[1]
+// CHECK: [[P1:%.*]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK: [[A1:%.*]] = begin_access [read] [dynamic] [[P1]] : $*S
+// CHECK: [[P2:%.*]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK: [[A2:%.*]] = begin_access [read] [dynamic] [[P2]] : $*S
+// CHECK: Access path base:  [[P2]] = ref_element_addr %0 : $Klass, #Klass.f
+// CHECK-NEXT: Exact Use:   %{{.*}} = load [trivial] [[A2]] : $*S
+// CHECK-NEXT: Exact Use:   end_access [[A2]] : $*S
+// CHECK-LABEL: end running test 2 of 2 on testRefElement: accesspath-base with: @trace[1]
+sil hidden [ossa] @testRefElement : $@convention(thin) (@guaranteed Klass) -> () {
+bb0(%0 : @guaranteed $Klass):
+  test_specification "accesspath-base @trace[0]"
+  %p1 = ref_element_addr %0 : $Klass, #Klass.f
+  debug_value [trace] %p1 : $*S
+  %a1 = begin_access [read] [dynamic] %p1 : $*S
+  %l1 = load [trivial] %a1 : $*S
+  end_access %a1 : $*S
+  test_specification "accesspath-base @trace[1]"
+  %p2 = ref_element_addr %0 : $Klass, #Klass.f
+  debug_value [trace] %p2 : $*S
+  %a2 = begin_access [read] [dynamic] %p2 : $*S
+  %l2 = load [trivial] %a2 : $*S
+  end_access %a2 : $*S
+  %99 = tuple ()
+  return %99 : $()
+}