[OSLogOptimization] Fix a bug in the replaceAllUsesAndFixLifetimes

function of the OSLogOptimization pass that happens when folding a
guaranteed value with a constant (owned) value. The fix inserts the
constant value at the beginning of the borrowed scope of the guaranteed
value rather than at the definition of the guaranteed value.

Update the SIL tests for the new folding pattern and add a test that
catches this bug.

Also, re-enable the OSLogPrototypeExecTest.swift that was disabled
due to this bug.

Author: ravikandhadai

lib/SILOptimizer/Mandatory/OSLogOptimization.cpp

@@ -82,6 +82,7 @@
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/CFG.h"
 #include "swift/SIL/InstructionUtils.h"
+#include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SIL/SILConstants.h"
@@ -510,7 +511,6 @@ static SILValue emitCodeForConstantArray(ArrayRef<SILValue> elements,
   assert(astContext.getArrayDecl() ==
          arrayType->getNominalOrBoundGenericNominal());
   SILModule &module = builder.getModule();
-  SILFunction &fun = builder.getFunction();
 
   // Create a SILValue for the number of elements.
   unsigned numElements = elements.size();
@@ -538,19 +538,10 @@ static SILValue emitCodeForConstantArray(ArrayRef<SILValue> elements,
       loc, arrayAllocateRef, subMap, ArrayRef<SILValue>(numElementsSIL), false);
 
   // Extract the elements of the tuple returned by the call to the allocator.
-  SILValue arraySIL;
-  SILValue storagePointerSIL;
-  if (fun.hasOwnership()) {
-    DestructureTupleInst *destructureInst =
-        builder.createDestructureTuple(loc, applyInst);
-    arraySIL = destructureInst->getResults()[0];
-    storagePointerSIL = destructureInst->getResults()[1];
-  } else {
-    SILType arraySILType = SILType::getPrimitiveObjectType(arrayType);
-    arraySIL = builder.createTupleExtract(loc, applyInst, 0, arraySILType);
-    storagePointerSIL = builder.createTupleExtract(
-        loc, applyInst, 1, SILType::getRawPointerType(astContext));
-  }
+  DestructureTupleInst *destructureInst =
+      builder.createDestructureTuple(loc, applyInst);
+  SILValue arraySIL = destructureInst->getResults()[0];
+  SILValue storagePointerSIL = destructureInst->getResults()[1];
 
   if (elements.empty()) {
     // Nothing more to be done if we are creating an empty array.
@@ -846,20 +837,21 @@ getEndPointsOfDataDependentChain(SILValue value, SILFunction *fun,
   }
 }
 
-/// Given an instruction \p inst, invoke the given clean-up function \p cleanup
-/// on its lifetime frontier, which are instructions that follow the last use of
-/// the results of \c inst. E.g. the clean-up function could destory/release
-/// the function result.
-static void
-cleanupAtEndOfLifetime(SILInstruction *inst,
-                       llvm::function_ref<void(SILInstruction *)> cleanup) {
-  ValueLifetimeAnalysis lifetimeAnalysis = ValueLifetimeAnalysis(inst);
-  ValueLifetimeAnalysis::Frontier frontier;
-  (void)lifetimeAnalysis.computeFrontier(
-      frontier, ValueLifetimeAnalysis::AllowToModifyCFG);
-  for (SILInstruction *lifetimeEndInst : frontier) {
-    cleanup(lifetimeEndInst);
-  }
+/// Given a guaranteed SILValue \p value, return a borrow-scope introducing
+/// value, if there is exactly one such introducing value. Otherwise, return
+/// None. There can be multiple borrow scopes for a SILValue iff it is derived
+/// from a guaranteed basic block parameter representing a phi node.
+static Optional<BorrowScopeIntroducingValue>
+getUniqueBorrowScopeIntroducingValue(SILValue value) {
+  assert(value.getOwnershipKind() == ValueOwnershipKind::Guaranteed &&
+         "parameter must be a guarenteed value");
+  SmallVector<BorrowScopeIntroducingValue, 4> borrowIntroducers;
+  getUnderlyingBorrowIntroducingValues(value, borrowIntroducers);
+  assert(borrowIntroducers.size() > 0 &&
+         "folding guaranteed value with no borrow introducer");
+  if (borrowIntroducers.size() > 1)
+    return None;
+  return borrowIntroducers[0];
 }
 
 /// Replace all uses of \c originalVal by \c foldedVal and adjust lifetimes of
@@ -888,45 +880,7 @@ static void replaceAllUsesAndFixLifetimes(SILValue foldedVal,
     return;
   }
   assert(!foldedVal->getType().isTrivial(*fun));
-
-  if (!fun->hasOwnership()) {
-    // In non-ownership SIL, handle only folding of struct_extract instruction,
-    // which is the only important instruction that should be folded by this
-    // pass. The logic used here is not correct in general and overfits a
-    // specific form of SIL. This code should be removed once OSSA is enabled
-    // for this pass.
-    // TODO: this code can be safely removed once ownership SIL becomes the
-    // default SIL this pass works on.
-    assert(isa<StructExtractInst>(originalInst) &&
-           !originalVal->getType().isAddress());
-
-    // First, replace all uses of originalVal by foldedVal, and then adjust
-    // their lifetimes if necessary.
-    originalVal->replaceAllUsesWith(foldedVal);
-
-    unsigned retainCount = 0;
-    unsigned consumeCount = 0;
-    for (Operand *use : foldedVal->getUses()) {
-      SILInstruction *user = use->getUser();
-      if (isa<ReleaseValueInst>(user) || isa<StoreInst>(user))
-        consumeCount++;
-      if (isa<RetainValueInst>(user))
-        retainCount++;
-      // Note that there could other consuming operations but they are not
-      // handled here as this code should be phased out soon.
-    }
-    if (consumeCount > retainCount) {
-      // The original value was at +1 and therefore consumed at the end. Since
-      // the foldedVal is also at +1 there is nothing to be done.
-      return;
-    }
-    cleanupAtEndOfLifetime(foldedInst, [&](SILInstruction *lifetimeEndInst) {
-      SILBuilderWithScope builder(lifetimeEndInst);
-      builder.emitReleaseValue(lifetimeEndInst->getLoc(), foldedVal);
-    });
-    return;
-  }
-
+  assert(fun->hasOwnership());
   assert(foldedVal.getOwnershipKind() == ValueOwnershipKind::Owned &&
          "constant value must have owned ownership kind");
 
@@ -942,22 +896,36 @@ static void replaceAllUsesAndFixLifetimes(SILValue foldedVal,
     return;
   }
 
-  // Here, originalVal is not owned. Hence, borrow form foldedVal and use the
-  // borrow in place of originalVal. Also, destroy foldedVal at the end of its
-  // lifetime.
+  // Here, originalVal is guaranteed. It must belong to a borrow scope that
+  // begins at a scope introducing instruction e.g. begin_borrow or load_borrow.
+  // The foldedVal should also have been inserted at the beginning of the scope.
+  // Therefore, create a borrow of foldedVal at the beginning of the scope and
+  // use the borrow in place of the originalVal. Also, end the borrow and
+  // destroy foldedVal at the end of the borrow scope.
   assert(originalVal.getOwnershipKind() == ValueOwnershipKind::Guaranteed);
 
-  SILBuilderWithScope builder(&*std::next(foldedInst->getIterator()));
-  BeginBorrowInst *borrow =
-      builder.createBeginBorrow(foldedInst->getLoc(), foldedVal);
+  Optional<BorrowScopeIntroducingValue> originalScopeBegin =
+      getUniqueBorrowScopeIntroducingValue(originalVal);
+  assert(originalScopeBegin &&
+         "value without a unique borrow scope should not have been folded");
+  SILInstruction *scopeBeginInst =
+      originalScopeBegin->value->getDefiningInstruction();
+  assert(scopeBeginInst);
+
+  SILBuilderWithScope builder(scopeBeginInst);
+  SILValue borrow =
+      builder.emitBeginBorrowOperation(scopeBeginInst->getLoc(), foldedVal);
 
   originalVal->replaceAllUsesWith(borrow);
 
-  cleanupAtEndOfLifetime(borrow, [&](SILInstruction *lifetimeEndInst) {
-    SILBuilderWithScope builder(lifetimeEndInst);
-    builder.createEndBorrow(lifetimeEndInst->getLoc(), borrow);
-    builder.emitDestroyValueOperation(lifetimeEndInst->getLoc(), foldedVal);
-  });
+  SmallVector<SILInstruction *, 4> scopeEndingInsts;
+  originalScopeBegin->getLocalScopeEndingInstructions(scopeEndingInsts);
+
+  for (SILInstruction *scopeEndingInst : scopeEndingInsts) {
+    SILBuilderWithScope builder(scopeEndingInst);
+    builder.emitEndBorrowOperation(scopeEndingInst->getLoc(), borrow);
+    builder.emitDestroyValueOperation(scopeEndingInst->getLoc(), foldedVal);
+  }
   return;
 }
 
@@ -977,13 +945,31 @@ static void substituteConstants(FoldState &foldState) {
     assert(definingInst);
     SILFunction *fun = definingInst->getFunction();
 
-    // Do not attempt to fold anything but struct_extract in non-OSSA.
-    // TODO: this condition should be removed once migration OSSA is complete.
-    if (!fun->hasOwnership() && !isa<StructExtractInst>(definingInst))
-      continue;
+    // Find an insertion point for inserting the new constant value. If we are
+    // folding a value like struct_extract within a borrow scope, we need to
+    // insert the constant value at the beginning of the borrow scope. This
+    // is because the borrowed value is expected to be alive during its entire
+    // borrow scope and could be stored into memory and accessed indirectly
+    // without a copy e.g. using store_borrow within the borrow scope. On the
+    // other hand, if we are folding an owned value, we can insert the constant
+    // value at the point where the owned value is defined.
+    SILInstruction *insertionPoint = definingInst;
+    if (constantSILValue.getOwnershipKind() == ValueOwnershipKind::Guaranteed) {
+      Optional<BorrowScopeIntroducingValue> borrowIntroducer =
+          getUniqueBorrowScopeIntroducingValue(constantSILValue);
+      if (!borrowIntroducer) {
+        // This case happens only if constantSILValue is derived from a
+        // guaranteed basic block parameter. This is unlikley because the values
+        // that have to be folded should just be a struct-extract of an owned
+        // instance of OSLogMessage.
+        continue;
+      }
+      insertionPoint = borrowIntroducer->value->getDefiningInstruction();
+      assert(insertionPoint && "borrow scope begnning is a parameter");
+    }
 
-    SILBuilderWithScope builder(definingInst);
-    SILLocation loc = definingInst->getLoc();
+    SILBuilderWithScope builder(insertionPoint);
+    SILLocation loc = insertionPoint->getLoc();
     CanType instType = constantSILValue->getType().getASTType();
     SILValue foldedSILVal = emitCodeForSymbolicValue(
         constantSymbolicVal, instType, builder, loc, foldState.stringInfo);
@@ -1067,6 +1053,7 @@ static bool constantFold(SILInstruction *start,
                          SingleValueInstruction *oslogMessage,
                          unsigned assertConfig) {
   SILFunction *fun = start->getFunction();
+  assert(fun->hasOwnership() && "function not in ownership SIL");
 
   // Initialize fold state.
   SmallVector<SILInstruction *, 2> endUsersOfOSLogMessage;