Merge pull request #32560 from slavapestov/clean-up-closure-discriminators

slavapestov · web-flow · commit 7f8c6796ef78 · 2020-06-25T22:15:09.000-04:00
AutoDiff: Clean up closure discriminator assignment
diff --git a/lib/Sema/CodeSynthesis.cpp b/lib/Sema/CodeSynthesis.cpp
@@ -1414,22 +1414,3 @@ void swift::addFixedLayoutAttr(NominalTypeDecl *nominal) {
   // Add `@_fixed_layout` to the nominal.
   nominal->getAttrs().add(new (C) FixedLayoutAttr(/*Implicit*/ true));
 }
-
-Expr *DiscriminatorFinder::walkToExprPost(Expr *E) {
-  auto *ACE = dyn_cast<AbstractClosureExpr>(E);
-  if (!ACE)
-    return E;
-
-  unsigned Discriminator = ACE->getDiscriminator();
-  assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
-         "Existing closures should have valid discriminators");
-  if (Discriminator >= NextDiscriminator)
-    NextDiscriminator = Discriminator + 1;
-  return E;
-}
-
-unsigned DiscriminatorFinder::getNextDiscriminator() {
-  if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
-    llvm::report_fatal_error("Out of valid closure discriminators");
-  return NextDiscriminator++;
-}
diff --git a/lib/Sema/CodeSynthesis.h b/lib/Sema/CodeSynthesis.h
@@ -76,20 +76,6 @@ bool hasLetStoredPropertyWithInitialValue(NominalTypeDecl *nominal);
 /// Add `@_fixed_layout` attribute to the nominal type, if possible.
 void addFixedLayoutAttr(NominalTypeDecl *nominal);
 
-/// Find available closure discriminators.
-///
-/// The parser typically takes care of assigning unique discriminators to
-/// closures, but the parser is unavailable during semantic analysis.
-class DiscriminatorFinder : public ASTWalker {
-  unsigned NextDiscriminator = 0;
-
-public:
-  Expr *walkToExprPost(Expr *E) override;
-
-  // Get the next available closure discriminator.
-  unsigned getNextDiscriminator();
-};
-
 } // end namespace swift
 
 #endif
diff --git a/lib/Sema/DebuggerTestingTransform.cpp b/lib/Sema/DebuggerTestingTransform.cpp
@@ -34,6 +34,35 @@ using namespace swift;
 
 namespace {
 
+/// Find available closure discriminators.
+///
+/// The parser typically takes care of assigning unique discriminators to
+/// closures, but the parser is unavailable during semantic analysis.
+class DiscriminatorFinder : public ASTWalker {
+  unsigned NextDiscriminator = 0;
+
+public:
+  Expr *walkToExprPost(Expr *E) override {
+    auto *ACE = dyn_cast<AbstractClosureExpr>(E);
+    if (!ACE)
+      return E;
+
+    unsigned Discriminator = ACE->getDiscriminator();
+    assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
+           "Existing closures should have valid discriminators");
+    if (Discriminator >= NextDiscriminator)
+      NextDiscriminator = Discriminator + 1;
+    return E;
+  }
+
+  // Get the next available closure discriminator.
+  unsigned getNextDiscriminator() {
+    if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
+      llvm::report_fatal_error("Out of valid closure discriminators");
+    return NextDiscriminator++;
+  }
+};
+
 /// Instrument decls with sanity-checks which the debugger can evaluate.
 class DebuggerTestingTransform : public ASTWalker {
   ASTContext &Ctx;
diff --git a/lib/Sema/DerivedConformanceDifferentiable.cpp b/lib/Sema/DerivedConformanceDifferentiable.cpp
@@ -388,10 +388,7 @@ deriveBodyDifferentiable_zeroTangentVectorInitializer(
                               DeclNameLoc(), /*Implicit*/ true);
 
     // Create closure expression.
-    DiscriminatorFinder DF;
-    for (Decl *D : parentDC->getParentSourceFile()->getTopLevelDecls())
-      D->walk(DF);
-    auto discriminator = DF.getNextDiscriminator();
+    unsigned discriminator = 0;
     auto resultTy = funcDecl->getMethodInterfaceType()
                         ->castTo<AnyFunctionType>()
                         ->getResult();
@@ -502,10 +499,7 @@ deriveBodyDifferentiable_zeroTangentVectorInitializer(
 
   // Create closure expression:
   // `{ TangentVector(x: x_zeroTangentVectorInitializer(), ...) }`.
-  DiscriminatorFinder DF;
-  for (Decl *D : parentDC->getParentSourceFile()->getTopLevelDecls())
-    D->walk(DF);
-  auto discriminator = DF.getNextDiscriminator();
+  unsigned discriminator = 0;
   auto resultTy = funcDecl->getMethodInterfaceType()
                       ->castTo<AnyFunctionType>()
                       ->getResult();
diff --git a/test/AutoDiff/Sema/DerivedConformances/derived_zero_tangent_vector_initializer.swift b/test/AutoDiff/Sema/DerivedConformances/derived_zero_tangent_vector_initializer.swift
@@ -137,8 +137,8 @@ enum CustomTangentVectorEnum<T: Differentiable>: Differentiable {
 // CHECK:   [[Y_PROP:%.*]] = struct_extract [[SELF]] : $SelfTangentVectorStruct, #SelfTangentVectorStruct.y
 // CHECK:   [[Y_ZERO_INIT_FN:%.*]] = function_ref @$sSd{{.*}}E28zeroTangentVectorInitializerSdycvg : $@convention(method) (Double) -> @owned @callee_guaranteed () -> Double
 // CHECK:   [[Y_ZERO_INIT:%.*]] = apply [[Y_ZERO_INIT_FN]]([[Y_PROP]])
-// CHECK:   // function_ref closure #2 in SelfTangentVectorStruct.zeroTangentVectorInitializer.getter
-// CHECK:   [[CLOSURE_FN:%.*]] = function_ref @${{.*}}SelfTangentVectorStructV0bgH11InitializerACycvgACycfU0_
+// CHECK:   // function_ref closure #1 in SelfTangentVectorStruct.zeroTangentVectorInitializer.getter
+// CHECK:   [[CLOSURE_FN:%.*]] = function_ref @${{.*}}SelfTangentVectorStructV0bgH11InitializerACycvgACycfU_
 // CHECK:   [[X_ZERO_INIT_COPY:%.*]] = copy_value [[X_ZERO_INIT]]
 // CHECK:   [[Y_ZERO_INIT_COPY:%.*]] = copy_value [[Y_ZERO_INIT]]
 // CHECK:   [[ZERO_INIT:%.*]] = partial_apply [callee_guaranteed] [[CLOSURE_FN]]([[X_ZERO_INIT_COPY]], [[Y_ZERO_INIT_COPY]])
@@ -148,8 +148,8 @@ enum CustomTangentVectorEnum<T: Differentiable>: Differentiable {
 // CHECK-LABEL: // CustomTangentVectorStruct.zeroTangentVectorInitializer.getter
 // CHECK-NEXT: sil hidden [_semantics "autodiff.nonvarying"] [ossa] @${{.*}}CustomTangentVectorStructV0bgH11Initializer0gH0Qzycvg : $@convention(method) <T, U where T : Differentiable, U : Differentiable> (@in_guaranteed CustomTangentVectorStruct<T, U>) -> @owned @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <T.TangentVector> {
 // CHECK: bb0([[SELF:%.*]] : $*CustomTangentVectorStruct<T, U>):
-// CHECK:   // function_ref closure #3 in CustomTangentVectorStruct.zeroTangentVectorInitializer.getter
-// CHECK:   function_ref @${{.*}}CustomTangentVectorStructV0bgH11Initializer0gH0QzycvgAFycfU1_
+// CHECK:   // function_ref closure #1 in CustomTangentVectorStruct.zeroTangentVectorInitializer.getter
+// CHECK:   function_ref @${{.*}}CustomTangentVectorStructV0bgH11Initializer0gH0QzycvgAFycfU_
 // CHECK: }
 
 // CHECK-LABEL: // MemberwiseTangentVectorClass.zeroTangentVectorInitializer.getter
@@ -163,8 +163,8 @@ enum CustomTangentVectorEnum<T: Differentiable>: Differentiable {
 // CHECK:   [[Y_PROP:%.*]] = apply [[Y_PROP_METHOD]]([[SELF]])
 // CHECK:   [[Y_ZERO_INIT_FN:%.*]] = function_ref @$sSd{{.*}}E28zeroTangentVectorInitializerSdycvg : $@convention(method) (Double) -> @owned @callee_guaranteed () -> Double
 // CHECK:   [[Y_ZERO_INIT:%.*]] = apply [[Y_ZERO_INIT_FN]]([[Y_PROP]])
-// CHECK:   // function_ref closure #4 in MemberwiseTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK:   [[CLOSURE_FN:%.*]] = function_ref @${{.*}}MemberwiseTangentVectorClassC0bgH11InitializerAC0gH0VycvgAFycfU2_
+// CHECK:   // function_ref closure #1 in MemberwiseTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK:   [[CLOSURE_FN:%.*]] = function_ref @${{.*}}MemberwiseTangentVectorClassC0bgH11InitializerAC0gH0VycvgAFycfU_
 // CHECK:   [[X_ZERO_INIT_COPY:%.*]] = copy_value [[X_ZERO_INIT]]
 // CHECK:   [[Y_ZERO_INIT_COPY:%.*]] = copy_value [[Y_ZERO_INIT]]
 // CHECK:   [[ZERO_INIT:%.*]] = partial_apply [callee_guaranteed] [[CLOSURE_FN]]([[X_ZERO_INIT_COPY]], [[Y_ZERO_INIT_COPY]])
@@ -174,29 +174,29 @@ enum CustomTangentVectorEnum<T: Differentiable>: Differentiable {
 // CHECK-LABEL: // SelfTangentVectorClass.zeroTangentVectorInitializer.getter
 // CHECK-NEXT: sil hidden [_semantics "autodiff.nonvarying"] [ossa] @${{.*}}SelfTangentVectorClassC0bgH11InitializerACycvg : $@convention(method) (@guaranteed SelfTangentVectorClass) -> @owned @callee_guaranteed () -> @owned SelfTangentVectorClass {
 // CHECK: bb0([[SELF:%.*]] : @guaranteed $SelfTangentVectorClass):
-// CHECK:   // function_ref closure #5 in SelfTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK:   function_ref @${{.*}}SelfTangentVectorClassC0bgH11InitializerACycvgACycfU3_
+// CHECK:   // function_ref closure #1 in SelfTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK:   function_ref @${{.*}}SelfTangentVectorClassC0bgH11InitializerACycvgACycfU_
 // CHECK: }
 
 // CHECK-LABEL: // CustomTangentVectorClass.zeroTangentVectorInitializer.getter
 // CHECK-NEXT: sil hidden [_semantics "autodiff.nonvarying"] [ossa] @${{.*}}CustomTangentVectorClassC0bgH11Initializer0gH0Qzycvg : $@convention(method) <T, U where T : Differentiable, U : Differentiable> (@guaranteed CustomTangentVectorClass<T, U>) -> @owned @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <T.TangentVector> {
 // CHECK: bb0(%0 : @guaranteed $CustomTangentVectorClass<T, U>):
-// CHECK:   // function_ref closure #6 in CustomTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK:   function_ref @${{.*}}CustomTangentVectorClassC0bgH11Initializer0gH0QzycvgAFycfU4_
+// CHECK:   // function_ref closure #1 in CustomTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK:   function_ref @${{.*}}CustomTangentVectorClassC0bgH11Initializer0gH0QzycvgAFycfU_
 // CHECK: }
 
 // CHECK-LABEL: // SelfTangentVectorEnum.zeroTangentVectorInitializer.getter
 // CHECK-NEXT: sil hidden [_semantics "autodiff.nonvarying"] [ossa] @${{.*}}SelfTangentVectorEnumO0bgH11InitializerACycvg : $@convention(method) (@guaranteed SelfTangentVectorEnum) -> @owned @callee_guaranteed () -> @owned SelfTangentVectorEnum {
 // CHECK: bb0([[SELF:%.*]] : @guaranteed $SelfTangentVectorEnum):
-// CHECK:   // function_ref closure #7 in SelfTangentVectorEnum.zeroTangentVectorInitializer.getter
-// CHECK:   function_ref @${{.*}}SelfTangentVectorEnumO0bgH11InitializerACycvgACycfU5_
+// CHECK:   // function_ref closure #1 in SelfTangentVectorEnum.zeroTangentVectorInitializer.getter
+// CHECK:   function_ref @${{.*}}SelfTangentVectorEnumO0bgH11InitializerACycvgACycfU_
 // CHECK: }
 
 // CHECK-LABEL: // CustomTangentVectorEnum.zeroTangentVectorInitializer.getter
 // CHECK-NEXT: sil hidden [_semantics "autodiff.nonvarying"] [ossa] @${{.*}}CustomTangentVectorEnumO0bgH11Initializer0gH0Qzycvg : $@convention(method) <T where T : Differentiable> (@in_guaranteed CustomTangentVectorEnum<T>) -> @owned @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <T.TangentVector> {
 // CHECK: bb0([[SELF:%.*]] : $*CustomTangentVectorEnum<T>):
-// CHECK:   // function_ref closure #8 in CustomTangentVectorEnum.zeroTangentVectorInitializer.getter
-// CHECK:   function_ref @${{.*}}CustomTangentVectorEnumO0bgH11Initializer0gH0QzycvgAFycfU6_
+// CHECK:   // function_ref closure #1 in CustomTangentVectorEnum.zeroTangentVectorInitializer.getter
+// CHECK:   function_ref @${{.*}}CustomTangentVectorEnumO0bgH11Initializer0gH0QzycvgAFycfU_
 // CHECK: }
 
 // CHECK-LABEL: // closure #1 in MemberwiseTangentVectorStruct.zeroTangentVectorInitializer.getter
@@ -206,44 +206,44 @@ enum CustomTangentVectorEnum<T: Differentiable>: Differentiable {
 // CHECK-NOT:   witness_method {{.*}}, #AdditiveArithmetic.zero!getter
 // CHECK: }
 
-// CHECK-LABEL: // closure #2 in SelfTangentVectorStruct.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorStructV0bgH11InitializerACycvgACycfU0_ : $@convention(thin) (@guaranteed @callee_guaranteed () -> Float, @guaranteed @callee_guaranteed () -> Double) -> SelfTangentVectorStruct {
+// CHECK-LABEL: // closure #1 in SelfTangentVectorStruct.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorStructV0bgH11InitializerACycvgACycfU_ : $@convention(thin) (@guaranteed @callee_guaranteed () -> Float, @guaranteed @callee_guaranteed () -> Double) -> SelfTangentVectorStruct {
 // CHECK:   // function_ref SelfTangentVectorStruct.init(x:y:)
 // CHECK-NOT:   // function_ref static {{.*}}.zero.getter
 // CHECK-NOT:   witness_method {{.*}}, #AdditiveArithmetic.zero!getter
 // CHECK: }
 
-// CHECK-LABEL: // closure #3 in CustomTangentVectorStruct.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}CustomTangentVectorStructV0bgH11Initializer0gH0QzycvgAFycfU1_ : $@convention(thin) <T, U where T : Differentiable, U : Differentiable> () -> @out T.TangentVector {
+// CHECK-LABEL: // closure #1 in CustomTangentVectorStruct.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}CustomTangentVectorStructV0bgH11Initializer0gH0QzycvgAFycfU_ : $@convention(thin) <T, U where T : Differentiable, U : Differentiable> () -> @out T.TangentVector {
 // CHECK:   witness_method $T.TangentVector, #AdditiveArithmetic.zero!getter
 // CHECK: }
 
-// CHECK-LABEL: // closure #4 in MemberwiseTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}MemberwiseTangentVectorClassC0bgH11InitializerAC0gH0VycvgAFycfU2_ : $@convention(thin) (@guaranteed @callee_guaranteed () -> Float, @guaranteed @callee_guaranteed () -> Double) -> MemberwiseTangentVectorClass.TangentVector {
+// CHECK-LABEL: // closure #1 in MemberwiseTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}MemberwiseTangentVectorClassC0bgH11InitializerAC0gH0VycvgAFycfU_ : $@convention(thin) (@guaranteed @callee_guaranteed () -> Float, @guaranteed @callee_guaranteed () -> Double) -> MemberwiseTangentVectorClass.TangentVector {
 // CHECK:   // function_ref MemberwiseTangentVectorClass.TangentVector.init(x:y:)
 // CHECK-NOT:   // function_ref static {{.*}}.zero.getter
 // CHECK-NOT:   witness_method {{.*}}, #AdditiveArithmetic.zero!getter
 // CHECK: }
 
-// CHECK-LABEL: // closure #5 in SelfTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorClassC0bgH11InitializerACycvgACycfU3_ : $@convention(thin) () -> @owned SelfTangentVectorClass {
+// CHECK-LABEL: // closure #1 in SelfTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorClassC0bgH11InitializerACycvgACycfU_ : $@convention(thin) () -> @owned SelfTangentVectorClass {
 // CHECK:   // function_ref static SelfTangentVectorClass.zero.getter
 // CHECK:   function_ref @${{.*}}SelfTangentVectorClassC0B0ACXDvgZ : $@convention(method) (@thick SelfTangentVectorClass.Type) -> @owned SelfTangentVectorClass
 // CHECK: }
 
-// CHECK-LABEL: // closure #6 in CustomTangentVectorClass.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}CustomTangentVectorClassC0bgH11Initializer0gH0QzycvgAFycfU4_ : $@convention(thin) <T, U where T : Differentiable, U : Differentiable> () -> @out T.TangentVector {
+// CHECK-LABEL: // closure #1 in CustomTangentVectorClass.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}CustomTangentVectorClassC0bgH11Initializer0gH0QzycvgAFycfU_ : $@convention(thin) <T, U where T : Differentiable, U : Differentiable> () -> @out T.TangentVector {
 // CHECK:   witness_method $T.TangentVector, #AdditiveArithmetic.zero!getter
 // CHECK: }
 
 // TODO(TF-1012): Implement memberwise `zeroTangentVectorInitializer` synthesis for enums.
-// CHECK-LABEL: // closure #7 in SelfTangentVectorEnum.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorEnumO0bgH11InitializerACycvgACycfU5_ : $@convention(thin) () -> @owned SelfTangentVectorEnum {
+// CHECK-LABEL: // closure #1 in SelfTangentVectorEnum.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @${{.*}}SelfTangentVectorEnumO0bgH11InitializerACycvgACycfU_ : $@convention(thin) () -> @owned SelfTangentVectorEnum {
 // CHECK:   // function_ref static SelfTangentVectorEnum.zero.getter
 // CHECK:   function_ref @${{.*}}SelfTangentVectorEnumO0B0ACvgZ : $@convention(method) (@thin SelfTangentVectorEnum.Type) -> @owned SelfTangentVectorEnum
 // CHECK: }
 
-// CHECK-LABEL: // closure #8 in CustomTangentVectorEnum.zeroTangentVectorInitializer.getter
-// CHECK-NEXT: sil private [ossa] @$s39derived_zero_tangent_vector_initializer23CustomTangentVectorEnumO0bgH11Initializer0gH0QzycvgAFycfU6_ : $@convention(thin) <T where T : Differentiable> () -> @out T.TangentVector {
+// CHECK-LABEL: // closure #1 in CustomTangentVectorEnum.zeroTangentVectorInitializer.getter
+// CHECK-NEXT: sil private [ossa] @$s39derived_zero_tangent_vector_initializer23CustomTangentVectorEnumO0bgH11Initializer0gH0QzycvgAFycfU_ : $@convention(thin) <T where T : Differentiable> () -> @out T.TangentVector {
 // CHECK:   witness_method $T.TangentVector, #AdditiveArithmetic.zero!getter
 // CHECK: }
