SILGen: Rewrite captured local archetypes into primary archetypes

Author: slavapestov

lib/SILGen/CMakeLists.txt

@@ -28,6 +28,7 @@ add_swift_host_library(swiftSILGen STATIC
   SILGenFunction.cpp
   SILGenGlobalVariable.cpp
   SILGenLazyConformance.cpp
+  SILGenLocalArchetype.cpp
   SILGenLValue.cpp
   SILGenPack.cpp
   SILGenPattern.cpp

lib/SILGen/SILGen.cpp

@@ -1260,6 +1260,9 @@ void SILGenModule::postEmitFunction(SILDeclRef constant,
                                     SILFunction *F) {
   emitLazyConformancesForFunction(F);
 
+  auto sig = Types.getGenericSignatureWithCapturedEnvironments(constant);
+  recontextualizeCapturedLocalArchetypes(F, sig);
+
   assert(!F->isExternalDeclaration() && "did not emit any function body?!");
   LLVM_DEBUG(llvm::dbgs() << "lowered sil:\n";
              F->print(llvm::dbgs()));

lib/SILGen/SILGen.h

@@ -625,6 +625,11 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Emit a property descriptor for the given storage decl if it needs one.
   void tryEmitPropertyDescriptor(AbstractStorageDecl *decl);
 
+  /// Replace local archetypes captured from outer AST contexts with primary
+  /// archetypes.
+  void recontextualizeCapturedLocalArchetypes(
+      SILFunction *F, GenericSignatureWithCapturedEnvironments sig);
+
 private:
   /// The most recent declaration we considered for emission.
   SILDeclRef lastEmittedFunction;

lib/SILGen/SILGenLazyConformance.cpp

@@ -9,6 +9,11 @@
 // See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
 //
 //===----------------------------------------------------------------------===//
+//
+//  This file forces emission of lazily-generated ClangImporter-synthesized
+//  conformances.
+//
+//===----------------------------------------------------------------------===//
 
 #include "SILGen.h"
 #include "swift/AST/Decl.h"

lib/SILGen/SILGenLocalArchetype.cpp

@@ -0,0 +1,149 @@
+//===--- SILGenLocalArchetype.cpp - Local archetype transform -------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 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
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file implements the transformation which rewrites captured local
+//  archetypes into primary archetypes in the enclosing function's generic
+//  signature.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SILGen.h"
+#include "swift/AST/LocalArchetypeRequirementCollector.h"
+#include "swift/SIL/SILInstruction.h"
+#include "swift/SIL/SILCloner.h"
+
+using namespace swift;
+using namespace swift::Lowering;
+
+namespace {
+
+class LocalArchetypeTransform : public SILCloner<LocalArchetypeTransform> {
+  friend class SILCloner<LocalArchetypeTransform>;
+  friend class SILInstructionVisitor<LocalArchetypeTransform>;
+
+  GenericSignatureWithCapturedEnvironments sig;
+  GenericEnvironment *env;
+
+public:
+  LocalArchetypeTransform(SILFunction *F,
+                          GenericSignatureWithCapturedEnvironments sig)
+      : SILCloner(*F), env(sig.genericSig.getGenericEnvironment()) {
+
+    assert(!sig.capturedEnvs.empty() && "Why are we doing this?");
+
+    // The primary archetypes of the old generic environment map to
+    // primary archetypes of the new generic environment at the same
+    // index and depth.
+    Functor.SubsMap = env->getForwardingSubstitutionMap();
+
+    // Local archetypes map to generic parameters at higher depths.
+    MapLocalArchetypesOutOfContext mapOutOfContext(sig.baseGenericSig,
+                                                   sig.capturedEnvs);
+
+    // For each captured environment...
+    for (auto *capturedEnv : sig.capturedEnvs) {
+      // For each introduced generic parameter...
+      auto localParams = capturedEnv->getGenericSignature()
+          .getInnermostGenericParams();
+      for (auto *gp : localParams) {
+        // Get the local archetype from the captured environment.
+        auto origArchetypeTy = capturedEnv->mapTypeIntoContext(gp)
+            ->castTo<LocalArchetypeType>();
+
+        // Map the local archetype to an interface type in the new generic
+        // signature.
+        auto substInterfaceTy = mapOutOfContext(origArchetypeTy);
+
+        // Map this interface type into the new generic environment to get
+        // a primary archetype.
+        auto substArchetypeTy = env->mapTypeIntoContext(substInterfaceTy)
+            ->castTo<PrimaryArchetypeType>();
+
+        // Remember this correspondence.
+        registerLocalArchetypeRemapping(origArchetypeTy, substArchetypeTy);
+      }
+    }
+  }
+
+  void doIt() {
+    auto &F = getBuilder().getFunction();
+
+    // Collect the old basic blocks that we're going to delete.
+    llvm::SmallVector<SILBasicBlock *, 4> bbs;
+    for (auto &bb : F)
+      bbs.push_back(&bb);
+
+    // Make F.mapTypeIntoContext() use the new environment.
+    F.setGenericEnvironment(env);
+
+    // Start by cloning the entry block.
+    auto *origEntryBlock = F.getEntryBlock();
+    auto *clonedEntryBlock = F.createBasicBlock();
+
+    // Clone arguments.
+    SmallVector<SILValue, 4> entryArgs;
+    entryArgs.reserve(origEntryBlock->getArguments().size());
+    for (auto &origArg : origEntryBlock->getArguments()) {
+
+      // Remap the argument type into the new generic environment.
+      SILType mappedType = remapType(origArg->getType());
+      auto *NewArg = clonedEntryBlock->createFunctionArgument(
+          mappedType, origArg->getDecl(), true);
+      NewArg->copyFlags(cast<SILFunctionArgument>(origArg));
+      entryArgs.push_back(NewArg);
+    }
+
+    // Clone the remaining body.
+    getBuilder().setInsertionPoint(clonedEntryBlock);
+    cloneFunctionBody(&F, clonedEntryBlock, entryArgs,
+                      true /*replaceOriginalFunctionInPlace*/);
+
+    // Insert the new entry block at the beginning.
+    F.moveBlockBefore(clonedEntryBlock, F.begin());
+
+    // FIXME: This should be a common utility.
+
+    // Erase the old basic blocks.
+    for (auto *bb : bbs) {
+      for (SILArgument *arg : bb->getArguments()) {
+        arg->replaceAllUsesWithUndef();
+        // To appease the ownership verifier, just set to None.
+        arg->setOwnershipKind(OwnershipKind::None);
+      }
+
+      // Instructions in the dead block may be used by other dead blocks. Replace
+      // any uses of them with undef values.
+      while (!bb->empty()) {
+        // Grab the last instruction in the bb.
+        auto *inst = &bb->back();
+
+        // Replace any still-remaining uses with undef values and erase.
+        inst->replaceAllUsesOfAllResultsWithUndef();
+        inst->eraseFromParent();
+      }
+
+      // Finally, erase the basic block itself.
+      bb->eraseFromParent();
+    }
+  }
+};
+
+} // end anonymous namespace
+
+void SILGenModule::recontextualizeCapturedLocalArchetypes(
+    SILFunction *F, GenericSignatureWithCapturedEnvironments sig) {
+  if (sig.capturedEnvs.empty())
+    return;
+
+  LocalArchetypeTransform(F, sig).doIt();
+  M.reclaimUnresolvedLocalArchetypeDefinitions();
+}

test/Interpreter/element_archetype_captures.swift

@@ -0,0 +1,78 @@
+// RUN: %target-run-simple-swift
+
+func callee<X: P1, T: P2, U: P3, V: P4>(x: X, t: T, u: U, v: V)
+    -> (any P1, any P2, any P3, any P4) {
+  return (x, t, u, v)
+}
+
+protocol P1 { var f1: Bool { get } }
+protocol P2 { var f2: Bool { get } }
+protocol P3 { var f3: Bool { get } }
+protocol P4 { var f4: Bool { get } }
+
+func packFunction<X, each T, each U, each V>(x: X,
+                                             ts: repeat each T,
+                                             us: repeat each U,
+                                             vs: repeat each V)
+    where X: P1,
+          (repeat (each T, each U)): Any,
+          repeat each T: P2,
+          repeat each U: P3,
+          repeat each V: P4 {
+
+  var result: [(any P1, any P2, any P3, any P4)] = []
+  var bools: [Bool] = []
+
+  for t in repeat each ts {
+    for u in repeat each us {
+      for v in repeat each vs {
+
+        // Local function
+        do {
+          func localFunc() {
+            result.append(callee(x: x, t: t, u: u, v: v))
+            bools.append(true && x.f1 && t.f2 && u.f3 && v.f4)
+          }
+
+          localFunc()
+          let fn = localFunc
+          fn()
+
+          let closure = { localFunc() }
+          closure()
+        }
+
+        // Generic local function
+        do {
+          func localFunc<Y: P1>(_ y: Y) {
+            result.append(callee(x: y, t: t, u: u, v: v))
+            bools.append(true && y.f1 && t.f2 && u.f3 && v.f4)
+          }
+
+          localFunc(x)
+          let fn: (X) -> () = localFunc
+          fn(x)
+
+          let closure = { localFunc(x) }
+          closure()
+        }
+      }
+    }
+  }
+
+  assert(result.count == bools.count)
+
+  for (r, b) in zip(result, bools) {
+    assert(b == (r.0.f1 && r.1.f2 && r.2.f3 && r.3.f4))
+  }
+}
+
+struct S1: P1 { var f1: Bool }
+struct S2: P2 { var f2: Bool }
+struct S3: P3 { var f3: Bool }
+struct S4: P4 { var f4: Bool }
+
+packFunction(x: S1(f1: true),
+             ts: S2(f2: true), S2(f2: true),
+             us: S3(f3: true), S3(f3: false),
+             vs: S4(f4: false), S4(f4: false), S4(f4: true), S4(f4: true))

test/SILGen/element_archetype_captures.swift

@@ -0,0 +1,32 @@
+// RUN: %target-swift-emit-silgen %s | %FileCheck %s
+
+// CHECK-LABEL: sil hidden [ossa] @$s26element_archetype_captures6calleeyyx_q_q0_tSTRzSTR_7ElementQy_ACRtzr1_lF : $@convention(thin) <T, U, V where T : Sequence, U : Sequence, T.Element == U.Element> (@in_guaranteed T, @in_guaranteed U, @in_guaranteed V) -> () {
+
+func callee<T, U, V>(_: T, _: U, _: V) where T: Sequence, U: Sequence, T.Element == U.Element {}
+
+// CHECK-LABEL: sil hidden [ossa] @$s26element_archetype_captures12packFunction2ts2us2vsyxxQp_q_xQpq0_q0_QptRvzRv_Rv0_STRzSTR_7ElementQy_AFRtzr1_lF : $@convention(thin) <each T, each U, each V where repeat each T : Sequence, repeat each U : Sequence, repeat (each T).Element == (each U).Element> (@pack_guaranteed Pack{repeat each T}, @pack_guaranteed Pack{repeat each U}, @pack_guaranteed Pack{repeat each V}) -> () {
+func packFunction<each T, each U, each V>(ts: repeat each T, us: repeat each U, vs: repeat each V)
+    where repeat each T: Sequence,
+          repeat each U: Sequence,
+          repeat (each T).Element == (each U).Element {
+  for (t, u) in repeat (each ts, each us) {
+
+    // // CHECK-LABEL: sil private [ossa] @$s26element_archetype_captures12packFunction2ts2us2vsyxxQp_q_xQpq0_q0_QptRvzRv_Rv0_STRzSTR_7ElementQy_AFRtzr1_lF10middleFuncL_yyRvzRv_Rv0_STRzSTR_AgHRSr1_lF : $@convention(thin) <each T, each U, each V where repeat each T : Sequence, repeat each U : Sequence, repeat (each T).Element == (each U).Element><τ_1_0, τ_1_1 where τ_1_0 : Sequence, τ_1_1 : Sequence, τ_1_0.Element == τ_1_1.Element> (@in_guaranteed (repeat each V), @in_guaranteed τ_1_0, @in_guaranteed τ_1_1) -> () {
+    func middleFunc() {
+      for v in repeat each vs {
+
+        // CHECK-LABEL: sil private [ossa] @$s26element_archetype_captures12packFunction2ts2us2vsyxxQp_q_xQpq0_q0_QptRvzRv_Rv0_STRzSTR_7ElementQy_AFRtzr1_lF10middleFuncL_yyRvzRv_Rv0_STRzSTR_AgHRSr1_lF05innerK0L_yyRvzRv_Rv0_STRzSTR_AgHRSr1_lF : $@convention(thin) <each T, each U, each V where repeat each T : Sequence, repeat each U : Sequence, repeat (each T).Element == (each U).Element><τ_1_0, τ_1_1 where τ_1_0 : Sequence, τ_1_1 : Sequence, τ_1_0.Element == τ_1_1.Element><τ_2_0> (@in_guaranteed τ_1_0, @in_guaranteed τ_1_1, @in_guaranteed τ_2_0) -> () {
+        func innerFunc() {
+
+          // CHECK: [[FN:%.*]] = function_ref @$s26element_archetype_captures6calleeyyx_q_q0_tSTRzSTR_7ElementQy_ACRtzr1_lF : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2 where τ_0_0 : Sequence, τ_0_1 : Sequence, τ_0_0.Element == τ_0_1.Element> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_1, @in_guaranteed τ_0_2) -> ()
+          apply [[FN]]<τ_1_0, τ_1_1, τ_2_0>(%0, %1, %2) : $@convention(thin) <τ_0_0, τ_0_1, τ_0_2 where τ_0_0 : Sequence, τ_0_1 : Sequence, τ_0_0.Element == τ_0_1.Element> (@in_guaranteed τ_0_0, @in_guaranteed τ_0_1, @in_guaranteed τ_0_2) -> ()
+          callee(t, u, v)
+        }
+
+        innerFunc()
+      }
+    }
+
+    middleFunc()
+  }
+}