[Distributed] Augment parameter/result type decoded to support generic substitutions

xedin · xedin · commit 206969a70521 · 2022-01-24T10:17:59.000-08:00
diff --git a/stdlib/public/Distributed/DistributedActorSystem.swift b/stdlib/public/Distributed/DistributedActorSystem.swift
@@ -175,6 +175,24 @@ extension DistributedActorSystem {
                                                 UInt(nameUTF8.endIndex))
     }
 
+    var substitutionsBuffer: UnsafeMutablePointer<Any.Type>? = nil
+
+    if genericEnv != nil {
+      let subs = try invocationDecoder.decodeGenericSubstitutions()
+
+      if subs.isEmpty {
+        throw ExecuteDistributedTargetError(
+          message: "Cannot call generic method without generic argument substitutions")
+      }
+
+      substitutionsBuffer = .allocate(capacity: subs.count)
+
+      for (offset, substitution) in subs.enumerated() {
+        let element = substitutionsBuffer?.advanced(by: offset)
+        element?.initialize(to: substitution)
+      }
+    }
+
     let paramCount = nameUTF8.withUnsafeBufferPointer { nameUTF8 in
       __getParameterCount(nameUTF8.baseAddress!, UInt(nameUTF8.endIndex))
     }
@@ -201,6 +219,7 @@ extension DistributedActorSystem {
       __getParameterTypeInfo(
         nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
         genericEnv,
+        substitutionsBuffer,
         paramTypesBuffer.baseAddress!._rawValue, Int(paramCount))
     }
 
@@ -229,7 +248,8 @@ extension DistributedActorSystem {
     }
 
     guard let returnTypeFromTypeInfo: Any.Type = _getReturnTypeInfo(mangledMethodName: mangledTargetName,
-                                                                    genericEnv: genericEnv) else {
+                                                                    genericEnv: genericEnv,
+                                                                    genericArguments: substitutionsBuffer) else {
       throw ExecuteDistributedTargetError(
         message: "Failed to decode distributed target return type")
     }
diff --git a/stdlib/public/Distributed/DistributedMetadata.swift b/stdlib/public/Distributed/DistributedMetadata.swift
@@ -43,13 +43,14 @@ public // SPI _Distributed
 func _getParameterTypeInfo(
   mangledMethodName name: String,
   genericEnv: UnsafeRawPointer?, // GenericEnvironmentDescriptor *
+  genericArguments: UnsafeRawPointer?,
   into typesBuffer: Builtin.RawPointer, length typesLength: Int
 ) -> Int32 {
   let nameUTF8 = Array(name.utf8)
   return nameUTF8.withUnsafeBufferPointer { nameUTF8 in
     return __getParameterTypeInfo(
         nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
-        genericEnv, typesBuffer, typesLength)
+        genericEnv, genericArguments, typesBuffer, typesLength)
   }
 }
 
@@ -61,19 +62,21 @@ public // SPI _Distributed
 func __getParameterTypeInfo(
     _ typeNameStart: UnsafePointer<UInt8>, _ typeNameLength: UInt,
     _ genericEnv: UnsafeRawPointer?, // GenericEnvironmentDescriptor *
+    _ genericArguments: UnsafeRawPointer?,
     _ types: Builtin.RawPointer, _ typesLength: Int
 ) -> Int32
 
 @available(SwiftStdlib 5.6, *)
 public // SPI _Distributed
 func _getReturnTypeInfo(
   mangledMethodName name: String,
-  genericEnv: UnsafeRawPointer? // GenericEnvironmentDescriptor *
+  genericEnv: UnsafeRawPointer?, // GenericEnvironmentDescriptor *
+  genericArguments: UnsafeRawPointer?
 ) -> Any.Type? {
   let nameUTF8 = Array(name.utf8)
   return nameUTF8.withUnsafeBufferPointer { nameUTF8 in
     return __getReturnTypeInfo(nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
-                               genericEnv)
+                               genericEnv, genericArguments)
   }
 }
 
@@ -83,7 +86,8 @@ public // SPI _Distributed
 func __getReturnTypeInfo(
     _ typeNameStart: UnsafePointer<UInt8>,
     _ typeNameLength: UInt,
-    _ genericEnv: UnsafeRawPointer? // GenericEnvironmentDescriptor *
+    _ genericEnv: UnsafeRawPointer?, // GenericEnvironmentDescriptor *
+    _ genericArguments: UnsafeRawPointer?
 ) -> Any.Type?
 
 
diff --git a/stdlib/public/runtime/MetadataLookup.cpp b/stdlib/public/runtime/MetadataLookup.cpp
@@ -1967,6 +1967,22 @@ static NodePointer getParameterList(NodePointer funcType) {
   return parameterContainer;
 }
 
+static const Metadata *decodeType(TypeDecoder<DecodedMetadataBuilder> &decoder,
+                                  NodePointer type) {
+  assert(type->getKind() == Node::Kind::Type);
+
+  auto builtTypeOrError = decoder.decodeMangledType(type);
+
+  if (builtTypeOrError.isError()) {
+    auto err = builtTypeOrError.getError();
+    char *errStr = err->copyErrorString();
+    err->freeErrorString(errStr);
+    return nullptr;
+  }
+
+  return builtTypeOrError.getType();
+}
+
 SWIFT_CC(swift)
 SWIFT_RUNTIME_STDLIB_SPI
 unsigned swift_func_getParameterCount(const char *typeNameStart,
@@ -1985,7 +2001,8 @@ unsigned swift_func_getParameterCount(const char *typeNameStart,
 SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_SPI
 const Metadata *_Nullable
 swift_func_getReturnTypeInfo(const char *typeNameStart, size_t typeNameLength,
-                             GenericEnvironmentDescriptor *genericEnv) {
+                             GenericEnvironmentDescriptor *genericEnv,
+                             const void * const *genericArguments) {
   StackAllocatedDemangler<1024> demangler;
 
   auto *funcType =
@@ -1999,31 +2016,30 @@ swift_func_getReturnTypeInfo(const char *typeNameStart, size_t typeNameLength,
 
   assert(resultType->getKind() == Node::Kind::ReturnType);
 
+  SubstGenericParametersFromMetadata substFn(genericEnv, genericArguments);
+
   DecodedMetadataBuilder builder(
       demangler,
       /*substGenericParam=*/
-      [](unsigned, unsigned) { return nullptr; },
+      [&substFn](unsigned depth, unsigned index) {
+        return substFn.getMetadata(depth, index);
+      },
       /*SubstDependentWitnessTableFn=*/
-      [](const Metadata *, unsigned) { return nullptr; });
+      [&substFn](const Metadata *type, unsigned index) {
+        return substFn.getWitnessTable(type, index);
+      });
 
   TypeDecoder<DecodedMetadataBuilder> decoder(builder);
-  auto builtTypeOrError =
-      decoder.decodeMangledType(resultType->getFirstChild());
-  if (builtTypeOrError.isError()) {
-    auto err = builtTypeOrError.getError();
-    char *errStr = err->copyErrorString();
-    err->freeErrorString(errStr);
-    return nullptr;
-  }
 
-  return builtTypeOrError.getType();
+  return decodeType(decoder, resultType->getFirstChild());
 }
 
 SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_SPI
 unsigned
 swift_func_getParameterTypeInfo(
     const char *typeNameStart, size_t typeNameLength,
     GenericEnvironmentDescriptor *genericEnv,
+    const void * const *genericArguments,
     Metadata const **types, unsigned typesLength) {
   if (typesLength < 0) return -1;
 
@@ -2041,35 +2057,36 @@ swift_func_getParameterTypeInfo(
   if (!(parameterList && parameterList->getNumChildren() == typesLength))
     return -2;
 
+  SubstGenericParametersFromMetadata substFn(genericEnv, genericArguments);
+
   DecodedMetadataBuilder builder(
       demangler,
       /*substGenericParam=*/
-      [](unsigned, unsigned) { return nullptr; },
+      [&substFn](unsigned depth, unsigned index) {
+        return substFn.getMetadata(depth, index);
+      },
       /*SubstDependentWitnessTableFn=*/
-      [](const Metadata *, unsigned) { return nullptr; });
+      [&substFn](const Metadata *type, unsigned index) {
+        return substFn.getWitnessTable(type, index);
+      });
   TypeDecoder<DecodedMetadataBuilder> decoder(builder);
 
-  auto typeIdx = 0;
   // for each parameter (TupleElement), store it into the provided buffer
-  for (auto *parameter : *parameterList) {
+  for (unsigned index = 0; index != typesLength; ++index) {
+    auto *parameter = parameterList->getChild(index);
+
     if (parameter->getKind() == Node::Kind::TupleElement) {
       assert(parameter->getNumChildren() == 1);
       parameter = parameter->getFirstChild();
     }
 
     assert(parameter->getKind() == Node::Kind::Type);
 
-    auto builtTypeOrError = decoder.decodeMangledType(parameter);
-    if (builtTypeOrError.isError()) {
-      auto err = builtTypeOrError.getError();
-      char *errStr = err->copyErrorString();
-      err->freeErrorString(errStr);
-      typeIdx += 1;
-      continue;
-    }
+    auto type = decodeType(decoder, parameter);
+    if (!type)
+      return -3; // Failed to decode a type.
 
-    types[typeIdx] = builtTypeOrError.getType();
-    ++typeIdx;
+    types[index] = type;
   } // end foreach parameter
 
   return typesLength;
diff --git a/test/Distributed/Runtime/distributed_func_metadata.swift b/test/Distributed/Runtime/distributed_func_metadata.swift
@@ -118,10 +118,10 @@ typealias DefaultDistributedActorSystem = FakeActorSystem
   static func test_returnType() {
     print("~~ \(#function)")
     // CHECK: _getReturnTypeInfo: empty() = ()
-    print("_getReturnTypeInfo: empty() = \(String(reflecting: _getReturnTypeInfo(mangledMethodName: empty, genericEnv: nil)!))")
+    print("_getReturnTypeInfo: empty() = \(String(reflecting: _getReturnTypeInfo(mangledMethodName: empty, genericEnv: nil, genericArguments: nil)!))")
 
     // CHECK: _getReturnTypeInfo: one(s:) = Swift.Int
-    print("_getReturnTypeInfo: one(s:) = \(String(reflecting: _getReturnTypeInfo(mangledMethodName: one, genericEnv: nil)!))")
+    print("_getReturnTypeInfo: one(s:) = \(String(reflecting: _getReturnTypeInfo(mangledMethodName: one, genericEnv: nil, genericArguments: nil)!))")
   }
 
   static func test_paramTypes() {
@@ -179,6 +179,7 @@ func _withParameterTypeInfo(
     let decodedNum = __getParameterTypeInfo(
         nameUTF8.baseAddress!, UInt(nameUTF8.endIndex),
         /*genericEnvironment=*/nil,
+        /*genericArguments=*/nil,
         infoBuffer.baseAddress!._rawValue, Int(paramCount))
 
     // if we failed demangling the types, return an empty array
