SILOptimizer: Replace [].append(contentsOf:) with [].append(element:)

Author: Ben Ng

include/swift/AST/ASTContext.h

@@ -450,7 +450,10 @@ class ASTContext {
 
   /// Retrieve the declaration of Swift.==(Int, Int) -> Bool.
   FuncDecl *getEqualIntDecl() const;
-  
+
+  /// Retrieve the declaration of Array.append(element:)
+  FuncDecl *getArrayAppendElementDecl() const;
+
   /// Retrieve the declaration of Swift._unimplementedInitializer.
   FuncDecl *getUnimplementedInitializerDecl(LazyResolver *resolver) const;
 

include/swift/SILOptimizer/Analysis/ArraySemantic.h

@@ -34,8 +34,12 @@ enum class ArrayCallKind {
   kMakeMutable,
   kMutateUnknown,
   kWithUnsafeMutableBufferPointer,
+  kAppendContentsOf,
+  kAppendElement,
   // The following two semantic function kinds return the result @owned
-  // instead of operating on self passed as parameter.
+  // instead of operating on self passed as parameter. If you are adding
+  // a function, and it has a self parameter, make sure that it is defined
+  // before this comment.
   kArrayInit,
   kArrayUninitialized
 };
@@ -131,6 +135,12 @@ class ArraySemanticsCall {
   /// Returns true on success, false otherwise.
   bool replaceByValue(SILValue V);
 
+  /// Replace a call to append(contentsOf: ) with a series of
+  /// append(element: ) calls.
+  bool replaceByAppendingValues(SILModule &M, SILFunction *AppendFn,
+                                llvm::SmallVectorImpl<SILValue> &Vals,
+                                ArrayRef<Substitution> Subs);
+
   /// Hoist the call to the insert point.
   void hoist(SILInstruction *InsertBefore, DominanceInfo *DT) {
     hoistOrCopy(InsertBefore, DT, false);

lib/AST/ASTContext.cpp

@@ -162,6 +162,9 @@ struct ASTContext::Implementation {
   /// func ==(Int, Int) -> Bool
   FuncDecl *EqualIntDecl = nullptr;
 
+  /// func append(Element) -> void
+  FuncDecl *ArrayAppendElementDecl = nullptr;
+
   /// func _unimplementedInitializer(className: StaticString).
   FuncDecl *UnimplementedInitializerDecl = nullptr;
 
@@ -883,6 +886,55 @@ FuncDecl *ASTContext::getEqualIntDecl() const {
   return nullptr;
 }
 
+FuncDecl *ASTContext::getArrayAppendElementDecl() const {
+  if (Impl.ArrayAppendElementDecl)
+    return Impl.ArrayAppendElementDecl;
+
+  auto AppendFunctions = getArrayDecl()->lookupDirect(getIdentifier("append"));
+
+  for (auto CandidateFn : AppendFunctions) {
+    auto FnDecl = dyn_cast<FuncDecl>(CandidateFn);
+    auto Attrs = FnDecl->getAttrs();
+    for (auto *A : Attrs.getAttributes<SemanticsAttr, false>()) {
+      if (A->Value != "array.append_element")
+        continue;
+
+#ifndef NDEBUG
+      auto ParamLists = FnDecl->getParameterLists();
+      assert(ParamLists.size() == 2 && "Should have two parameter lists");
+      assert(ParamLists[0]->size() == 1 &&
+             "First parameter list should have one parameter");
+      auto SelfInOutTy = ParamLists[0]->get(0)->getInterfaceType()->getAs<InOutType>();
+      assert(SelfInOutTy && "Self parameter should be an inout Type");
+      auto SelfGenericStructTy =
+          SelfInOutTy->getObjectType()->getAs<BoundGenericStructType>();
+      assert(SelfGenericStructTy &&
+             "Self parameter should be a BoundGenericStructType Type");
+      assert(SelfGenericStructTy->getDecl() == getArrayDecl() &&
+             "Self parameter should be an Array");
+
+      assert(ParamLists[1]->size() == 1 &&
+             "Second parameter list should have one parameter");
+      auto ElementType = ParamLists[1]
+                             ->get(0)
+                             ->getInterfaceType()
+                             ->getAs<GenericTypeParamType>();
+      assert(ElementType &&
+             "First parameter replacement should be a GenericTypeParamType");
+      assert(ElementType->getName() == getIdentifier("Element") &&
+             "The GenericTypeParamType's name should be \"Element\"");
+
+      assert(FnDecl->getResultInterfaceType()->isVoid() &&
+             "The return type should be void");
+#endif
+      Impl.ArrayAppendElementDecl = FnDecl;
+      return FnDecl;
+    }
+  }
+
+  return NULL;
+}
+
 FuncDecl *
 ASTContext::getUnimplementedInitializerDecl(LazyResolver *resolver) const {
   if (Impl.UnimplementedInitializerDecl)

lib/SILOptimizer/Analysis/ArraySemantic.cpp

@@ -162,7 +162,10 @@ ArrayCallKind swift::ArraySemanticsCall::getKind() const {
             .Case("array.get_element_address",
                   ArrayCallKind::kGetElementAddress)
             .Case("array.mutate_unknown", ArrayCallKind::kMutateUnknown)
-            .Case("array.withUnsafeMutableBufferPointer", ArrayCallKind::kWithUnsafeMutableBufferPointer)
+            .Case("array.withUnsafeMutableBufferPointer",
+                  ArrayCallKind::kWithUnsafeMutableBufferPointer)
+            .Case("array.append_contentsOf", ArrayCallKind::kAppendContentsOf)
+            .Case("array.append_element", ArrayCallKind::kAppendElement)
             .Default(ArrayCallKind::kNone);
     if (Tmp != ArrayCallKind::kNone) {
       assert(Kind == ArrayCallKind::kNone && "Multiple array semantic "
@@ -681,3 +684,42 @@ bool swift::ArraySemanticsCall::replaceByValue(SILValue V) {
   removeCall();
   return true;
 }
+
+bool swift::ArraySemanticsCall::replaceByAppendingValues(
+    SILModule &M, SILFunction *AppendFn, SmallVectorImpl<SILValue> &Vals,
+    ArrayRef<Substitution> Subs) {
+  assert(getKind() == ArrayCallKind::kAppendContentsOf &&
+         "Must be an append_contentsOf call");
+  assert(AppendFn && "Must provide an append SILFunction");
+
+  // We only handle loadable types.
+  if (any_of(Vals, [&M](SILValue V) -> bool {
+        return !V->getType().isLoadable(M);
+      }))
+    return false;
+
+  auto ArrRef = SemanticsCall->getArgument(1);
+  SILBuilderWithScope Builder(SemanticsCall);
+  auto Loc = SemanticsCall->getLoc();
+  auto *FnRef = Builder.createFunctionRef(Loc, AppendFn);
+  auto FnTy = FnRef->getType();
+
+  for (auto &V : Vals) {
+    auto SubTy = V->getType();
+    auto &ValLowering = Builder.getModule().getTypeLowering(SubTy);
+    auto CopiedVal = ValLowering.emitCopyValue(Builder, Loc, V);
+    auto *AllocStackInst = Builder.createAllocStack(Loc, SubTy);
+    ValLowering.emitStoreOfCopy(Builder, Loc, CopiedVal, AllocStackInst,
+                                IsInitialization_t::IsInitialization);
+    SILValue Args[] = {AllocStackInst, ArrRef};
+    Builder.createApply(Loc, FnRef, FnTy.substGenericArgs(M, Subs),
+                        FnTy.castTo<SILFunctionType>()->getAllResultsType(), Subs,
+                        Args, false);
+    Builder.createDeallocStack(Loc, AllocStackInst);
+    ValLowering.emitDestroyValue(Builder, Loc, CopiedVal);
+  }
+
+  removeCall();
+
+  return true;
+}

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -489,6 +489,8 @@ static bool isNonMutatingArraySemanticCall(SILInstruction *Inst) {
   case ArrayCallKind::kWithUnsafeMutableBufferPointer:
   case ArrayCallKind::kArrayInit:
   case ArrayCallKind::kArrayUninitialized:
+  case ArrayCallKind::kAppendContentsOf:
+  case ArrayCallKind::kAppendElement:
     return false;
   }
 
@@ -825,6 +827,8 @@ static bool mayChangeArrayValueToNonUniqueState(ArraySemanticsCall &Call) {
   case ArrayCallKind::kWithUnsafeMutableBufferPointer:
   case ArrayCallKind::kArrayInit:
   case ArrayCallKind::kArrayUninitialized:
+  case ArrayCallKind::kAppendContentsOf:
+  case ArrayCallKind::kAppendElement:
     return true;
   }