[AutoDiff] Fix mangling of '@noDerivative' in function types.

rxwei · rxwei · commit 82886bf77a7a · 2021-04-07T15:42:12.000-07:00
`@noDerivative` was not mangled in function types, and was resolved incorrectly when there's an ownership specifier. It is fixed by this patch with the following changes:

* Add `NoDerivative` demangle node represented by a `k` operator.
    ```
    list-type ::= type identifier? 'k'? 'z'? 'h'? 'n'? 'd'?  // type with optional label, '@noDerivative', inout convention, shared convention, owned convention, and variadic specifier
    ```
* Fix `NoDerivative`'s overflown offset in `ParameterTypeFlags` (`7` -&gt; `6`).
* In type decoder and type resolver where attributed type nodes are processed, add support for nested attributed nodes, e.g. `inout @noDerivative T`.
* Add `TypeResolverContext::InoutFunctionInput` so that when we resolve an `inout @noDerivative T` parameter, the `@noDerivative T` checking logic won't get a `TypeResolverContext::None` set by the caller.

Resolves rdar://75916833.
diff --git a/docs/ABI/Mangling.rst b/docs/ABI/Mangling.rst
@@ -582,7 +582,7 @@ Types
   type-list ::= empty-list
 
                                                   // FIXME: Consider replacing 'h' with a two-char code
-  list-type ::= type identifier? 'z'? 'h'? 'n'? 'd'?   // type with optional label, inout convention, shared convention, owned convention, and variadic specifier
+  list-type ::= type identifier? 'k'? 'z'? 'h'? 'n'? 'd'?  // type with optional label, '@noDerivative', inout convention, shared convention, owned convention, and variadic specifier
 
   METATYPE-REPR ::= 't'                      // Thin metatype representation
   METATYPE-REPR ::= 'T'                      // Thick metatype representation
@@ -666,7 +666,7 @@ mangled in to disambiguate.
   COROUTINE-KIND ::= 'A'                     // yield-once coroutine
   COROUTINE-KIND ::= 'G'                     // yield-many coroutine
 
-  SENDABLE ::= 'h'                         // @Sendable
+  SENDABLE ::= 'h'                           // @Sendable
   ASYNC ::= 'H'                              // @async
 
   PARAM-CONVENTION ::= 'i'                   // indirect in
diff --git a/include/swift/AST/Attr.h b/include/swift/AST/Attr.h
@@ -2324,8 +2324,6 @@ class TypeAttributes {
 
   Optional<Convention> ConventionArguments;
 
-  // Indicates whether the type's '@differentiable' attribute has a 'linear'
-  // argument.
   DifferentiabilityKind differentiabilityKind =
       DifferentiabilityKind::NonDifferentiable;
 
diff --git a/include/swift/AST/Types.h b/include/swift/AST/Types.h
@@ -1929,7 +1929,7 @@ class ParameterTypeFlags {
     NonEphemeral = 1 << 2,
     OwnershipShift = 3,
     Ownership    = 7 << OwnershipShift,
-    NoDerivative = 1 << 7,
+    NoDerivative = 1 << 6,
     NumBits = 7
   };
   OptionSet<ParameterFlags> value;
diff --git a/include/swift/Demangling/DemangleNodes.def b/include/swift/Demangling/DemangleNodes.def
@@ -312,6 +312,7 @@ NODE(AutoDiffSelfReorderingReabstractionThunk)
 NODE(AutoDiffSubsetParametersThunk)
 NODE(AutoDiffDerivativeVTableThunk)
 NODE(DifferentiabilityWitness)
+NODE(NoDerivative)
 NODE(IndexSubset)
 NODE(AsyncAwaitResumePartialFunction)
 NODE(AsyncSuspendResumePartialFunction)
diff --git a/include/swift/Demangling/TypeDecoder.h b/include/swift/Demangling/TypeDecoder.h
@@ -72,6 +72,7 @@ class FunctionParam {
   void setValueOwnership(ValueOwnership ownership) {
     Flags = Flags.withValueOwnership(ownership);
   }
+  void setNoDerivative() { Flags = Flags.withNoDerivative(true); }
   void setFlags(ParameterFlags flags) { Flags = flags; };
 
   FunctionParam withLabel(StringRef label) const {
@@ -1375,28 +1376,39 @@ class TypeDecoder {
         node = node->getFirstChild();
         hasParamFlags = true;
       };
-      switch (node->getKind()) {
-      case NodeKind::InOut:
-        setOwnership(ValueOwnership::InOut);
-        break;
 
-      case NodeKind::Shared:
-        setOwnership(ValueOwnership::Shared);
-        break;
+      bool recurse = true;
+      while (recurse) {
+        switch (node->getKind()) {
+        case NodeKind::InOut:
+          setOwnership(ValueOwnership::InOut);
+          break;
 
-      case NodeKind::Owned:
-        setOwnership(ValueOwnership::Owned);
-        break;
+        case NodeKind::Shared:
+          setOwnership(ValueOwnership::Shared);
+          break;
 
-      case NodeKind::AutoClosureType:
-      case NodeKind::EscapingAutoClosureType: {
-        param.setAutoClosure();
-        hasParamFlags = true;
-        break;
-      }
+        case NodeKind::Owned:
+          setOwnership(ValueOwnership::Owned);
+          break;
 
-      default:
-        break;
+        case NodeKind::NoDerivative:
+          param.setNoDerivative();
+          node = node->getFirstChild();
+          hasParamFlags = true;
+          break;
+
+        case NodeKind::AutoClosureType:
+        case NodeKind::EscapingAutoClosureType:
+          param.setAutoClosure();
+          hasParamFlags = true;
+          recurse = false;
+          break;
+
+        default:
+          recurse = false;
+          break;
+        }
       }
 
       auto paramType = decodeMangledType(node);
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -2541,6 +2541,9 @@ void ASTMangler::appendTypeListElement(Identifier name, Type elementType,
   else
     appendType(elementType, forDecl);
 
+  if (flags.isNoDerivative()) {
+    appendOperator("k");
+  }
   switch (flags.getValueOwnership()) {
   case ValueOwnership::Default:
     /* nothing */
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -793,6 +793,9 @@ NodePointer Demangler::demangleOperator() {
                                                 popTypeAndGetChild()));
     case 'i': return demangleSubscript();
     case 'j': return demangleDifferentiableFunctionType();
+    case 'k':
+      return createType(
+          createWithChild(Node::Kind::NoDerivative, popTypeAndGetChild()));
     case 'l': return demangleGenericSignature(/*hasParamCounts*/ false);
     case 'm': return createType(createWithChild(Node::Kind::Metatype,
                                                 popNode(Node::Kind::Type)));
diff --git a/lib/Demangling/NodePrinter.cpp b/lib/Demangling/NodePrinter.cpp
@@ -569,6 +569,7 @@ class NodePrinter {
     case Node::Kind::AutoDiffSubsetParametersThunk:
     case Node::Kind::AutoDiffFunctionKind:
     case Node::Kind::DifferentiabilityWitness:
+    case Node::Kind::NoDerivative:
     case Node::Kind::IndexSubset:
     case Node::Kind::AsyncAwaitResumePartialFunction:
     case Node::Kind::AsyncSuspendResumePartialFunction:
@@ -1421,6 +1422,10 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     Printer << "__owned ";
     print(Node->getChild(0));
     return nullptr;
+  case Node::Kind::NoDerivative:
+    Printer << "@noDerivative ";
+    print(Node->getChild(0));
+    return nullptr;
   case Node::Kind::NonObjCAttribute:
     Printer << "@nonobjc ";
     return nullptr;
diff --git a/lib/Demangling/OldDemangler.cpp b/lib/Demangling/OldDemangler.cpp
@@ -2063,6 +2063,14 @@ class OldDemangler {
       inout->addChild(type, Factory);
       return inout;
     }
+    if (c == 'k') {
+      auto noDerivative = Factory.createNode(Node::Kind::NoDerivative);
+      auto type = demangleTypeImpl();
+      if (!type)
+        return nullptr;
+      noDerivative->addChild(type, Factory);
+      return noDerivative;
+    }
     if (c == 'S') {
       return demangleSubstitutionIndex();
     }
diff --git a/lib/Demangling/OldRemangler.cpp b/lib/Demangling/OldRemangler.cpp
@@ -1485,6 +1485,11 @@ void Remangler::mangleInOut(Node *node) {
   mangleSingleChildNode(node); // type
 }
 
+void Remangler::mangleNoDerivative(Node *node) {
+  Buffer << 'k';
+  mangleSingleChildNode(node); // type
+}
+
 void Remangler::mangleTuple(Node *node) {
   size_t NumElems = node->getNumChildren();
   if (NumElems > 0 &&
diff --git a/lib/Demangling/Remangler.cpp b/lib/Demangling/Remangler.cpp
@@ -1676,6 +1676,11 @@ void Remangler::mangleOwned(Node *node) {
   Buffer << 'n';
 }
 
+void Remangler::mangleNoDerivative(Node *node) {
+  mangleSingleChildNode(node);
+  Buffer << 'k';
+}
+
 void Remangler::mangleInfixOperator(Node *node) {
   mangleIdentifierImpl(node, /*isOperator*/ true);
   Buffer << "oi";
diff --git a/lib/Sema/TypeCheckType.cpp b/lib/Sema/TypeCheckType.cpp
@@ -2524,10 +2524,12 @@ TypeResolver::resolveAttributedType(TypeAttributes &attrs, TypeRepr *repr,
   }
 
   if (attrs.has(TAK_noDerivative)) {
-    // @noDerivative is only valid on function parameters, or on function
-    // results in SIL.
+    // @noDerivative is valid on function parameters (AST and SIL) or on
+    // function results (SIL-only).
     bool isNoDerivativeAllowed =
-        isParam || (isResult && (options & TypeResolutionFlags::SILType));
+        isParam ||
+        options.is(TypeResolverContext::InoutFunctionInput) ||
+        (isResult && (options & TypeResolutionFlags::SILType));
     auto *SF = getDeclContext()->getParentSourceFile();
     if (SF && !isDifferentiableProgrammingEnabled(*SF)) {
       diagnose(
@@ -2645,7 +2647,7 @@ TypeResolver::resolveASTFunctionTypeParams(TupleTypeRepr *inputRepr,
     if (auto *ATR = dyn_cast<AttributedTypeRepr>(eltTypeRepr))
       autoclosure = ATR->getAttrs().has(TAK_autoclosure);
 
-    ValueOwnership ownership;
+    ValueOwnership ownership = ValueOwnership::Default;
 
     auto *nestedRepr = eltTypeRepr;
 
@@ -2657,31 +2659,35 @@ TypeResolver::resolveASTFunctionTypeParams(TupleTypeRepr *inputRepr,
       nestedRepr = tupleRepr->getElementType(0);
     }
 
-    switch (nestedRepr->getKind()) {
-    case TypeReprKind::Shared:
-      ownership = ValueOwnership::Shared;
-      break;
-    case TypeReprKind::InOut:
-      ownership = ValueOwnership::InOut;
-      break;
-    case TypeReprKind::Owned:
-      ownership = ValueOwnership::Owned;
-      break;
-    default:
-      ownership = ValueOwnership::Default;
-      break;
+    if (auto *specifierRepr = dyn_cast<SpecifierTypeRepr>(nestedRepr)) {
+      switch (specifierRepr->getKind()) {
+      case TypeReprKind::Shared:
+        ownership = ValueOwnership::Shared;
+        nestedRepr = specifierRepr->getBase();
+        break;
+      case TypeReprKind::InOut:
+        ownership = ValueOwnership::InOut;
+        nestedRepr = specifierRepr->getBase();
+        break;
+      case TypeReprKind::Owned:
+        ownership = ValueOwnership::Owned;
+        nestedRepr = specifierRepr->getBase();
+        break;
+      default:
+        break;
+      }
     }
 
     bool noDerivative = false;
-    if (auto *attrTypeRepr = dyn_cast<AttributedTypeRepr>(eltTypeRepr)) {
+    if (auto *attrTypeRepr = dyn_cast<AttributedTypeRepr>(nestedRepr)) {
       if (attrTypeRepr->getAttrs().has(TAK_noDerivative)) {
         if (diffKind == DifferentiabilityKind::NonDifferentiable &&
             isDifferentiableProgrammingEnabled(
                 *getDeclContext()->getParentSourceFile()))
-          diagnose(eltTypeRepr->getLoc(),
+          diagnose(nestedRepr->getLoc(),
                    diag::attr_only_on_parameters_of_differentiable,
                    "@noDerivative")
-              .highlight(eltTypeRepr->getSourceRange());
+              .highlight(nestedRepr->getSourceRange());
         else
           noDerivative = true;
       }
@@ -3476,7 +3482,7 @@ TypeResolver::resolveSpecifierTypeRepr(SpecifierTypeRepr *repr,
   if (isa<InOutTypeRepr>(repr)
       && !isa<ImplicitlyUnwrappedOptionalTypeRepr>(repr->getBase())) {
     // Anything within an inout isn't a parameter anymore.
-    options.setContext(None);
+    options.setContext(TypeResolverContext::InoutFunctionInput);
   }
 
   return resolveType(repr->getBase(), options);
@@ -3554,6 +3560,7 @@ NeverNullType TypeResolver::resolveImplicitlyUnwrappedOptionalType(
   bool doDiag = false;
   switch (options.getContext()) {
   case TypeResolverContext::None:
+  case TypeResolverContext::InoutFunctionInput:
     if (!isDirect || !(options & allowIUO))
       doDiag = true;
     break;
diff --git a/lib/Sema/TypeCheckType.h b/lib/Sema/TypeCheckType.h
@@ -85,6 +85,9 @@ enum class TypeResolverContext : uint8_t {
   /// Whether this is a variadic function input.
   VariadicFunctionInput,
 
+  /// Whether this is an 'inout' function input.
+  InoutFunctionInput,
+
   /// Whether we are in the result type of a function, including multi-level
   /// tuple return values. See also: TypeResolutionFlags::Direct
   FunctionResult,
@@ -198,6 +201,7 @@ class TypeResolutionOptions {
     case Context::None:
     case Context::FunctionInput:
     case Context::VariadicFunctionInput:
+    case Context::InoutFunctionInput:
     case Context::FunctionResult:
     case Context::ExtensionBinding:
     case Context::SubscriptDecl:
diff --git a/stdlib/public/Reflection/TypeRef.cpp b/stdlib/public/Reflection/TypeRef.cpp
@@ -582,6 +582,9 @@ class DemanglingForTypeRef
         parent->addChild(input, Dem);
         input = parent;
       };
+      if (flags.isNoDerivative()) {
+        wrapInput(Node::Kind::NoDerivative);
+      }
       switch (flags.getValueOwnership()) {
       case ValueOwnership::Default:
         /* nothing */
diff --git a/stdlib/public/runtime/Demangle.cpp b/stdlib/public/runtime/Demangle.cpp
@@ -482,6 +482,9 @@ swift::_swift_buildDemanglingForMetadata(const Metadata *type,
         parent->addChild(input, Dem);
         input = parent;
       };
+      if (flags.isNoDerivative()) {
+        wrapInput(Node::Kind::NoDerivative);
+      }
       switch (flags.getValueOwnership()) {
       case ValueOwnership::Default:
         /* nothing */
diff --git a/test/AutoDiff/validation-test/function_type_metadata.swift b/test/AutoDiff/validation-test/function_type_metadata.swift
@@ -26,14 +26,30 @@ if #available(macOS 9999, iOS 9999, watchOS 9999, tvOS 9999, *) {
       Swift.Optional<Swift.Float>
       """,
       String(reflecting: (@differentiable(reverse) (Float?) -> Float?).self))
-    // FIXME(rdar://75916833): Mangle '@noDerivative' in function types.
-    // expectEqual(
-    //   """
-    //   @differentiable(reverse) (Swift.Optional<Swift.Float>, \
-    //   @noDerivative Swift.Int) -> Swift.Optional<Swift.Float>
-    //   """,
-    //   String(reflecting: (
-    //       @differentiable(reverse) (Float?, @noDerivative Int) -> Float?).self))
+    expectEqual(
+      """
+      @differentiable(reverse) (Swift.Optional<Swift.Float>, \
+      @noDerivative Swift.Int) -> Swift.Optional<Swift.Float>
+      """,
+      String(reflecting: (
+          @differentiable(reverse)
+              (Float?, @noDerivative Int) -> Float?).self))
+    expectEqual(
+      """
+      @differentiable(reverse) (Swift.Optional<Swift.Float>, \
+      __owned @noDerivative Swift.Int) -> Swift.Optional<Swift.Float>
+      """,
+      String(reflecting: (
+          @differentiable(reverse)
+              (Float?, __owned @noDerivative Int) -> Float?).self))
+    expectEqual(
+      """
+      @differentiable(reverse) (Swift.Optional<Swift.Float>, \
+      inout @noDerivative Swift.Int) -> Swift.Optional<Swift.Float>
+      """,
+      String(reflecting: (
+          @differentiable(reverse)
+              (Float?, inout @noDerivative Int) -> Float?).self))
   }
 }
 
diff --git a/test/Demangle/Inputs/manglings.txt b/test/Demangle/Inputs/manglings.txt
@@ -399,3 +399,5 @@ $s4diff1hyyS2ijfXEF ---> diff.h(@differentiable(_forward) (Swift.Int) -> Swift.I
 $s4diff1hyyS2ijrXEF ---> diff.h(@differentiable(reverse) (Swift.Int) -> Swift.Int) -> ()
 $s4diff1hyyS2ijdXEF ---> diff.h(@differentiable (Swift.Int) -> Swift.Int) -> ()
 $s4diff1hyyS2ijlXEF ---> diff.h(@differentiable(_linear) (Swift.Int) -> Swift.Int) -> ()
+$s4test3fooyyS2f_SfkztjrXEF ---> test.foo(@differentiable(reverse) (Swift.Float, inout @noDerivative Swift.Float) -> Swift.Float) -> ()
+$s4test3fooyyS2f_SfktjrXEF ---> test.foo(@differentiable(reverse) (Swift.Float, @noDerivative Swift.Float) -> Swift.Float) -> ()
diff --git a/unittests/Reflection/TypeRef.cpp b/unittests/Reflection/TypeRef.cpp
@@ -250,6 +250,23 @@ TEST(TypeRefTest, UniqueFunctionTypeRef) {
       FunctionMetadataDifferentiabilityKind::Reverse);
   EXPECT_EQ(F17, F18);
   EXPECT_NE(F17, F19);
+
+  // Test differentiable with @noDerivative.
+  {
+    auto parameters = Parameters1;
+    parameters[1].setNoDerivative();
+    auto f1 = Builder.createFunctionType(
+        parameters, Result, FunctionTypeFlags().withDifferentiable(true),
+        FunctionMetadataDifferentiabilityKind::Reverse);
+    auto f2 = Builder.createFunctionType(
+        parameters, Result, FunctionTypeFlags().withDifferentiable(true),
+        FunctionMetadataDifferentiabilityKind::Reverse);
+    auto f3 = Builder.createFunctionType(
+        Parameters1, Result, FunctionTypeFlags().withDifferentiable(true),
+        FunctionMetadataDifferentiabilityKind::Reverse);
+    EXPECT_EQ(f1, f2);
+    EXPECT_NE(f1, f3);
+  }
 }
 
 TEST(TypeRefTest, UniqueProtocolTypeRef) {
