[AutoDiff] Fix reflection and metadata IRGen for differentiable function types

Author: DougGregor

lib/IRGen/MetadataRequest.cpp

@@ -1453,9 +1453,8 @@ namespace {
       llvm::Value *diffKindVal = nullptr;
       if (type->isDifferentiable()) {
         assert(metadataDifferentiabilityKind.isDifferentiable());
-        // FIXME: Shouldn't this use metadataDifferentiabilityKind?
-        diffKindVal = llvm::ConstantInt::get(IGF.IGM.SizeTy,
-                                             flags.getIntValue());
+        diffKindVal = llvm::ConstantInt::get(
+            IGF.IGM.SizeTy, metadataDifferentiabilityKind.getIntValue());
       } else if (type->getGlobalActor()) {
         diffKindVal = llvm::ConstantInt::get(
             IGF.IGM.SizeTy,

stdlib/public/Reflection/TypeRef.cpp

@@ -137,6 +137,27 @@ class PrintTypeRef : public TypeRefVisitor<PrintTypeRef, void> {
       break;
     }
 
+    switch (F->getDifferentiabilityKind().Value) {
+    case FunctionMetadataDifferentiabilityKind::NonDifferentiable:
+      break;
+
+    case FunctionMetadataDifferentiabilityKind::Forward:
+      printField("differentiable", "forward");
+      break;
+
+    case FunctionMetadataDifferentiabilityKind::Reverse:
+      printField("differentiable", "reverse");
+      break;
+
+    case FunctionMetadataDifferentiabilityKind::Normal:
+      printField("differentiable", "normal");
+      break;
+
+    case FunctionMetadataDifferentiabilityKind::Linear:
+      printField("differentiable", "linear");
+      break;
+    }
+
     if (auto globalActor = F->getGlobalActor()) {
       fprintf(file, "\n");
       Indent += 2;
@@ -683,7 +704,27 @@ class DemanglingForTypeRef
       funcNode->addChild(node, Dem);
     }
 
-    // FIXME: Differentiability is missing
+    if (F->getFlags().isDifferentiable()) {
+      MangledDifferentiabilityKind mangledKind;
+      switch (F->getDifferentiabilityKind().Value) {
+#define CASE(X) case FunctionMetadataDifferentiabilityKind::X: \
+        mangledKind = MangledDifferentiabilityKind::X; break;
+
+      CASE(NonDifferentiable)
+      CASE(Forward)
+      CASE(Reverse)
+      CASE(Normal)
+      CASE(Linear)
+#undef CASE
+      }
+
+      funcNode->addChild(
+          Dem.createNode(
+            Node::Kind::DifferentiableFunctionType,
+            (Node::IndexType)mangledKind),
+          Dem);
+    }
+
     if (F->getFlags().isThrowing())
       funcNode->addChild(Dem.createNode(Node::Kind::ThrowsAnnotation), Dem);
     if (F->getFlags().isSendable()) {

test/AutoDiff/validation-test/Inputs/AutoDiffTypes.swift

@@ -0,0 +1,5 @@
+import _Differentiation
+
+public struct HasAutoDiffTypes {
+  public var aFunction: @differentiable(reverse) (Float) -> Float
+}

test/AutoDiff/validation-test/Inputs/main.swift

@@ -0,0 +1 @@
+// this file intentionally left blank

test/AutoDiff/validation-test/reflection.swift

@@ -0,0 +1,16 @@
+// REQUIRES: no_asan
+// RUN: %empty-directory(%t)
+import _Differentiation
+
+// RUN: %target-build-swift -Xfrontend -enable-anonymous-context-mangled-names %S/Inputs/AutoDiffTypes.swift -parse-as-library -emit-module -emit-library -module-name TypesToReflect -o %t/%target-library-name(TypesToReflect)
+// RUN: %target-build-swift -Xfrontend -enable-anonymous-context-mangled-names %S/Inputs/AutoDiffTypes.swift %S/Inputs/main.swift -emit-module -emit-executable -module-name TypesToReflect -o %t/TypesToReflect
+
+// RUN: %target-swift-reflection-dump -binary-filename %t/%target-library-name(TypesToReflect) | %FileCheck %s
+// RUN: %target-swift-reflection-dump -binary-filename %t/TypesToReflect | %FileCheck %s
+
+// CHECK: FIELDS:
+// CHECK: =======
+// CHECK: TypesToReflect.HasAutoDiffTypes
+// CHECK: aFunction: @differentiable(reverse) (Swift.Float) -> Swift.Float
+// CHECK: (function differentiable=reverse
+