Swift Optimizer: add the SSAUpdater utility

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -145,17 +145,15 @@ extension MutatingContext {
     SubstitutionMap(_bridged.getContextSubstitutionMap(type.bridged))
   }
 
-  // Private utilities
-
-  fileprivate func notifyInstructionsChanged() {
+  func notifyInstructionsChanged() {
     _bridged.asNotificationHandler().notifyChanges(.instructionsChanged)
   }
 
-  fileprivate func notifyCallsChanged() {
+  func notifyCallsChanged() {
     _bridged.asNotificationHandler().notifyChanges(.callsChanged)
   }
 
-  fileprivate func notifyBranchesChanged() {
+  func notifyBranchesChanged() {
     _bridged.asNotificationHandler().notifyChanges(.branchesChanged)
   }
 }

SwiftCompilerSources/Sources/Optimizer/Utilities/CMakeLists.txt

@@ -11,5 +11,6 @@ swift_compiler_sources(Optimizer
   OptUtils.swift
   WalkUtils.swift
   AccessUtils.swift
+  SSAUpdater.swift
   StaticInitCloner.swift
 )

SwiftCompilerSources/Sources/Optimizer/Utilities/SSAUpdater.swift

@@ -0,0 +1,51 @@
+//===--- SSAUpdater.swift -------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+import OptimizerBridging
+
+/// Utility for updating SSA for a set of SIL instructions defined in multiple blocks.
+struct SSAUpdater<Context: MutatingContext> {
+  let context: Context
+
+  init(type: Type, ownership: Ownership, _ context: Context) {
+    self.context = context
+    context._bridged.SSAUpdater_initialize(type.bridged, ownership._bridged)
+  }
+
+  mutating func addAvailableValue(_ value: Value, in block: BasicBlock) {
+    context._bridged.SSAUpdater_addAvailableValue(block.bridged, value.bridged)
+  }
+
+  /// Construct SSA for a value that is live at the *end* of a basic block.
+  mutating func getValue(atEndOf block: BasicBlock) -> Value {
+    context.notifyInstructionsChanged()
+    return context._bridged.SSAUpdater_getValueAtEndOfBlock(block.bridged).value
+  }
+
+  /// Construct SSA for a value that is live in the *middle* of a block.
+  /// This handles the case where the use is before a definition of the value in the same block.
+  ///
+  ///   bb1:
+  ///     %1 = def
+  ///     br bb2
+  ///   bb2:
+  ///       = use(?)
+  ///    %2 = def
+  ///    cond_br bb2, bb3
+  ///
+  /// In this case we need to insert a phi argument in bb2, merging %1 and %2.
+  mutating func getValue(inMiddleOf block: BasicBlock) -> Value {
+    context.notifyInstructionsChanged()
+    return context._bridged.SSAUpdater_getValueInMiddleOfBlock(block.bridged).value
+  }
+}

include/swift/SILOptimizer/OptimizerBridging.h

@@ -452,6 +452,26 @@ struct BridgedPassContext {
                                           invocation->getTransform());
   }
 
+  // SSAUpdater
+
+  void SSAUpdater_initialize(swift::SILType type, BridgedValue::Ownership ownership) const {
+    invocation->initializeSSAUpdater(type, castToOwnership(ownership));
+  }
+
+  void SSAUpdater_addAvailableValue(BridgedBasicBlock block, BridgedValue value) const {
+    invocation->getSSAUpdater()->addAvailableValue(block.getBlock(), value.getSILValue());
+  }
+
+  SWIFT_IMPORT_UNSAFE
+  BridgedValue SSAUpdater_getValueAtEndOfBlock(BridgedBasicBlock block) const {
+    return {invocation->getSSAUpdater()->getValueAtEndOfBlock(block.getBlock())};
+  }
+
+  SWIFT_IMPORT_UNSAFE
+  BridgedValue SSAUpdater_getValueInMiddleOfBlock(BridgedBasicBlock block) const {
+    return {invocation->getSSAUpdater()->getValueInMiddleOfBlock(block.getBlock())};
+  }
+
   // Options
 
   bool enableStackProtection() const {

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -17,6 +17,7 @@
 #include "swift/SILOptimizer/Analysis/Analysis.h"
 #include "swift/SILOptimizer/PassManager/PassPipeline.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
+#include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
@@ -68,6 +69,8 @@ class SwiftPassInvocation {
   /// All slabs, allocated by the pass.
   SILModule::SlabList allocatedSlabs;
 
+  SILSSAUpdater *ssaUpdater = nullptr;
+
   static constexpr int BlockSetCapacity = 8;
   char blockSetStorage[sizeof(BasicBlockSet) * BlockSetCapacity];
   bool aliveBlockSets[BlockSetCapacity];
@@ -82,7 +85,7 @@ class SwiftPassInvocation {
 
   bool needFixStackNesting = false;
 
-  void endPassRunChecks();
+  void endPass();
 
 public:
   SwiftPassInvocation(SILPassManager *passManager, SILFunction *function,
@@ -143,6 +146,17 @@ public:
 
   void setNeedFixStackNesting(bool newValue) { needFixStackNesting = newValue; }
   bool getNeedFixStackNesting() const { return needFixStackNesting; }
+
+  void initializeSSAUpdater(SILType type, ValueOwnershipKind ownership) {
+    if (!ssaUpdater)
+      ssaUpdater = new SILSSAUpdater;
+    ssaUpdater->initialize(type, ownership);
+  }
+
+  SILSSAUpdater *getSSAUpdater() const {
+    assert(ssaUpdater && "SSAUpdater not initialized");
+    return ssaUpdater;
+  }
 };
 
 /// The SIL pass manager.

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1360,7 +1360,7 @@ void SwiftPassInvocation::startInstructionPassRun(SILInstruction *inst) {
 }
 
 void SwiftPassInvocation::finishedModulePassRun() {
-  endPassRunChecks();
+  endPass();
   assert(!function && transform && "not running a pass");
   assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
          && "unhandled change notifications at end of module pass");
@@ -1368,22 +1368,26 @@ void SwiftPassInvocation::finishedModulePassRun() {
 }
 
 void SwiftPassInvocation::finishedFunctionPassRun() {
-  endPassRunChecks();
+  endPass();
   endTransformFunction();
   assert(allocatedSlabs.empty() && "StackList is leaking slabs");
   transform = nullptr;
 }
 
 void SwiftPassInvocation::finishedInstructionPassRun() {
-  endPassRunChecks();
+  endPass();
 }
 
-void SwiftPassInvocation::endPassRunChecks() {
+void SwiftPassInvocation::endPass() {
   assert(allocatedSlabs.empty() && "StackList is leaking slabs");
   assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
   assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
   assert(numClonersAllocated == 0 && "Not all cloners deallocated");
   assert(!needFixStackNesting && "Stack nesting not fixed");
+  if (ssaUpdater) {
+    delete ssaUpdater;
+    ssaUpdater = nullptr;
+  }
 }
 
 void SwiftPassInvocation::beginTransformFunction(SILFunction *function) {