MandatoryAllocBoxToStack: also handle new specialized functions

Add new created specializations to the worklist so that those are optimized as well.

rdar://154686063, rdar://154713388

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/DataStructures/Worklist.swift

@@ -125,7 +125,7 @@ struct FunctionWorklist {
     }
   }
 
-  mutating func pushIfNotVisited(contentsOf functions: [Function]) {
+  mutating func pushIfNotVisited<S: Sequence>(contentsOf functions: S) where S.Element == Function {
     for f in functions {
       pushIfNotVisited(f)
     }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -65,12 +65,16 @@ let allocBoxToStack = FunctionPass(name: "allocbox-to-stack") {
 let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
   (moduleContext: ModulePassContext) in
 
+  var worklist = FunctionWorklist()
+  worklist.pushIfNotVisited(contentsOf: moduleContext.functions)
+
   var originalsOfSpecializedFunctions = FunctionWorklist()
 
-  for function in moduleContext.functions {
+  while let function = worklist.pop() {
     moduleContext.transform(function: function) { context in
       let specFns = tryConvertBoxesToStack(in: function, context)
-      originalsOfSpecializedFunctions.pushIfNotVisited(contentsOf: specFns)
+      worklist.pushIfNotVisited(contentsOf: specFns.specializedFunctions)
+      originalsOfSpecializedFunctions.pushIfNotVisited(contentsOf: specFns.originalFunctions)
     }
   }
 
@@ -82,7 +86,7 @@ let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
 /// Converts all non-escaping `alloc_box` to `alloc_stack` and specializes called functions if a
 /// box is passed to a function.
 /// Returns the list of original functions for which a specialization has been created.
-private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPassContext) -> [Function] {
+private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPassContext) -> FunctionSpecializations {
   var promotableBoxes = Array<(AllocBoxInst, Flags)>()
   var functionsToSpecialize = FunctionSpecializations()
 
@@ -101,7 +105,7 @@ private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPa
     function.fixStackNesting(context)
   }
 
-  return functionsToSpecialize.originalFunctions
+  return functionsToSpecialize
 }
 
 private func findPromotableBoxes(in function: Function,
@@ -187,6 +191,7 @@ private struct FunctionSpecializations {
   private var originalToSpecialized = Dictionary<Function, Function>()
 
   var originalFunctions: [Function] { originals.functions }
+  var specializedFunctions: [Function] { originals.functions.lazy.map { originalToSpecialized[$0]! } }
 
   mutating func add(promotableArguments: [FunctionArgument]) {
     for arg in promotableArguments {

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -1198,3 +1198,36 @@ bb0(%0 : $*T, %1 : $*T, %2 : $*T):
   %15 = tuple ()
   return %15 : $()
 }
+
+// CHECK-LABEL: sil [ossa] @alloc_box_in_specialized_callee :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function 'alloc_box_in_specialized_callee'
+sil [ossa] @alloc_box_in_specialized_callee : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %2 = project_box %1, 0
+  store %0 to [trivial] %2
+  %4 = function_ref @callee_with_allocbox : $@convention(thin) (@guaranteed { var Int }) -> ()
+  %5 = apply %4(%1) : $@convention(thin) (@guaranteed { var Int }) -> ()
+  destroy_value %1
+  %r = tuple ()
+  return %r
+}
+
+// CHECK-LABEL: sil shared [ossa] @$s20callee_with_allocboxTf0s_n :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function '$s20callee_with_allocboxTf0s_n'
+
+// CHECK-LABEL: sil [ossa] @callee_with_allocbox :
+// CHECK-NOT:     alloc_box
+// CHECK-LABEL: } // end sil function 'callee_with_allocbox'
+sil [ossa] @callee_with_allocbox : $@convention(thin) (@guaranteed { var Int }) -> () {
+bb0(%0 : @guaranteed ${ var Int }):
+  %1 = alloc_box ${ var Int }
+  %2 = project_box %1 : ${ var Int }, 0
+  %3 = project_box %0 : ${ var Int }, 0
+  copy_addr %3 to %2
+  destroy_value %1
+  %r = tuple ()
+  return %r
+}