[Distributed] RemoteCallTarget now pretty prints and hides mangled names

Author: ktoso

include/swift/Runtime/HeapObject.h

@@ -1101,6 +1101,15 @@ swift_getTypeName(const Metadata *type, bool qualified);
 ///   -> (UnsafePointer<UInt8>, Int)
 SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_API
 TypeNamePair
+swift_getFunctionFullNameFromMangledName(
+        const char *mangledNameStart, uintptr_t mangledNameLength);
+
+/// Return the human-readable full name of the mangled function name passed in.
+/// func _getMangledTypeName(_ mangledName: UnsafePointer<UInt8>,
+///                          mangledNameLength: UInt)
+///   -> (UnsafePointer<UInt8>, Int)
+SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_API
+TypeNamePair
 swift_getMangledTypeName(const Metadata *type);
 
 } // end namespace swift

lib/Sema/TypeCheckDistributed.cpp

@@ -310,7 +310,7 @@ bool swift::checkDistributedActorSystemAdHocProtocolRequirements(
     }
     if (checkAdHocRequirementAccessControl(decl, Proto, recordArgumentDecl))
       anyMissingAdHocRequirements = true;
-    
+
     // - recordReturnType
     auto recordReturnTypeDecl = C.getRecordReturnTypeOnDistributedInvocationEncoder(decl);
     if (!recordReturnTypeDecl) {

stdlib/public/Distributed/DistributedActorSystem.swift

@@ -172,23 +172,18 @@ extension DistributedActorSystem {
   /// latency sensitive-use-cases.
   public func executeDistributedTarget<Act, ResultHandler>(
     on actor: Act,
-    mangledTargetName: String,
+    target: RemoteCallTarget,
     invocationDecoder: inout InvocationDecoder,
     handler: ResultHandler
   ) async throws where Act: DistributedActor,
                        // Act.ID == ActorID, // FIXME(distributed): can we bring this back?
                        ResultHandler: DistributedTargetInvocationResultHandler {
-    // NOTE: this implementation is not the most efficient, nor final, version of this func
-    // we end up demangling the name multiple times, perform more heap allocations than
-    // we truly need to etc. We'll eventually move this implementation to a specialized one
-    // avoiding these issues.
-    guard mangledTargetName.count > 0 && mangledTargetName.first == "$" else {
+    // Get the expected parameter count of the func
+    guard let targetName = target.identifier else {
       throw ExecuteDistributedTargetError(
-        message: "Illegal mangledTargetName detected, must start with '$'")
+        message: "Unable to extract target identifier from remote call target: \(target)")
     }
-
-    // Get the expected parameter count of the func
-    let nameUTF8 = Array(mangledTargetName.utf8)
+    let nameUTF8 = Array(targetName.utf8)
 
     // Gen the generic environment (if any) associated with the target.
     let genericEnv = nameUTF8.withUnsafeBufferPointer { nameUTF8 in
@@ -207,7 +202,9 @@ extension DistributedActorSystem {
 
     if let genericEnv = genericEnv {
       let subs = try invocationDecoder.decodeGenericSubstitutions()
-
+      for item in subs {
+        print("SUB: \(item)")
+      }
       if subs.isEmpty {
         throw ExecuteDistributedTargetError(
           message: "Cannot call generic method without generic argument substitutions")
@@ -224,7 +221,7 @@ extension DistributedActorSystem {
                                                                      genericArguments: substitutionsBuffer!)
       if numWitnessTables < 0 {
         throw ExecuteDistributedTargetError(
-          message: "Generic substitutions \(subs) do not satisfy generic requirements of \(mangledTargetName)")
+          message: "Generic substitutions \(subs) do not satisfy generic requirements of \(target) (\(targetName))")
       }
     }
 
@@ -237,7 +234,7 @@ extension DistributedActorSystem {
         message: """
                  Failed to decode distributed invocation target expected parameter count,
                  error code: \(paramCount)
-                 mangled name: \(mangledTargetName)
+                 mangled name: \(targetName)
                  """)
     }
 
@@ -262,7 +259,7 @@ extension DistributedActorSystem {
         message: """
                  Failed to decode the expected number of params of distributed invocation target, error code: \(decodedNum)
                  (decoded: \(decodedNum), expected params: \(paramCount)
-                 mangled name: \(mangledTargetName)
+                 mangled name: \(targetName)
                  """)
     }
 
@@ -280,7 +277,7 @@ extension DistributedActorSystem {
       return UnsafeRawPointer(UnsafeMutablePointer<R>.allocate(capacity: 1))
     }
 
-    guard let returnTypeFromTypeInfo: Any.Type = _getReturnTypeInfo(mangledMethodName: mangledTargetName,
+    guard let returnTypeFromTypeInfo: Any.Type = _getReturnTypeInfo(mangledMethodName: targetName,
                                                                     genericEnv: genericEnv,
                                                                     genericArguments: substitutionsBuffer) else {
       throw ExecuteDistributedTargetError(
@@ -307,7 +304,7 @@ extension DistributedActorSystem {
       // Execute the target!
       try await _executeDistributedTarget(
         on: actor,
-        mangledTargetName, UInt(mangledTargetName.count),
+        targetName, UInt(targetName.count),
         argumentDecoder: &invocationDecoder,
         argumentTypes: argumentTypesBuffer.baseAddress!._rawValue,
         resultBuffer: resultBuffer._rawValue,
@@ -326,6 +323,52 @@ extension DistributedActorSystem {
   }
 }
 
+/// Represents a 'target' of a distributed call, such as a `distributed func` or
+/// `distributed` computed property. Identification schemes may vary between
+/// systems, and are subject to evolution.
+///
+/// Actor systems generally should treat the `identifier` as an opaque string,
+/// and pass it along to the remote system for in their `remoteCall`
+/// implementation. Alternative approaches are possible, where the identifiers
+/// are either compressed, cached, or represented in other ways, as long as the
+/// recipient system is able to determine which target was intended to be
+/// invoked.
+///
+/// The string representation will attempt to pretty print the target identifier,
+/// however its exact format is not specified and may change in future versions.
+@available(SwiftStdlib 5.7, *)
+public struct RemoteCallTarget: CustomStringConvertible {
+  private let _storage: _Storage
+  private enum _Storage {
+    case mangledName(String)
+  }
+
+  // Only intended to be created by the _Distributed library.
+  // TODO(distributed): make this internal and only allow calling by the synthesized code?
+  public init(_mangledName: String) {
+    self._storage = .mangledName(_mangledName)
+  }
+
+  /// The underlying identifier of the target, returned as-is.
+  public var identifier: String? {
+    switch self._storage {
+    case .mangledName(let name):
+      return name
+    }
+  }
+
+  public var description: String {
+    switch self._storage {
+    case .mangledName(let mangledName):
+      if let name = _getFunctionFullNameFromMangledName(mangledName: mangledName) {
+        return name
+      } else {
+        return "\(mangledName)"
+      }
+    }
+  }
+}
+
 @available(SwiftStdlib 5.7, *)
 @_silgen_name("swift_distributed_execute_target")
 func _executeDistributedTarget<D: DistributedTargetInvocationDecoder>(
@@ -339,29 +382,6 @@ func _executeDistributedTarget<D: DistributedTargetInvocationDecoder>(
   numWitnessTables: UInt
 ) async throws
 
-// ==== ----------------------------------------------------------------------------------------------------------------
-// MARK: Support types
-/// A distributed 'target' can be a `distributed func` or `distributed` computed property.
-@available(SwiftStdlib 5.7, *)
-public struct RemoteCallTarget { // TODO: ship this around always; make printing nice
-  let _mangledName: String // TODO: StaticString would be better here; no arc, codesize of cleanups
-
-  // Only intended to be created by the _Distributed library.
-  // TODO(distributed): make this internal and only allow calling by the synthesized code?
-  public init(_mangledName: String) {
-    self._mangledName = _mangledName
-  }
-
-  public var mangledName: String {
-    _mangledName
-  }
-
-  // <module>.Base.hello(hi:)
-  public var fullName: String { // TODO: make description
-    fatalError("NOT IMPLEMENTED YET: \(#function)")
-  }
-}
-
 /// Used to encode an invocation of a distributed target (method or computed property).
 ///
 /// ## Forming an invocation

stdlib/public/core/Misc.swift

@@ -46,6 +46,35 @@ public func _autorelease(_ x: AnyObject) {
 }
 #endif
 
+
+@available(SwiftStdlib 5.7, *)
+@_silgen_name("swift_getFunctionFullNameFromMangledName")
+public // SPI (Distributed)
+func _getFunctionFullNameFromMangledNameImpl(
+  _ mangledName: UnsafePointer<UInt8>, _ mangledNameLength: UInt
+) -> (UnsafePointer<UInt8>, UInt)
+
+/// Given a function's mangled name, return a human readable name.
+/// Used e.g. by Distributed.RemoteCallTarget to hide mangled names.
+@available(SwiftStdlib 5.7, *)
+public // SPI (Distributed)
+func _getFunctionFullNameFromMangledName(mangledName: String) -> String? {
+  let mangledNameUTF8 = Array(mangledName.utf8)
+  let (stringPtr, count) =
+    mangledNameUTF8.withUnsafeBufferPointer { (mangledNameUTF8) in
+    return _getFunctionFullNameFromMangledNameImpl(
+      mangledNameUTF8.baseAddress!,
+      UInt(mangledNameUTF8.endIndex))
+  }
+
+  guard count > 0 else {
+    return nil
+  }
+
+  return String._fromUTF8Repairing(
+    UnsafeBufferPointer(start: stringPtr, count: Int(count))).0
+}
+
 // FIXME(ABI)#51 : this API should allow controlling different kinds of
 // qualification separately: qualification with module names and qualification
 // with type names that we are nested in.

stdlib/public/runtime/Casting.cpp

@@ -146,13 +146,19 @@ using TypeNameCacheKey = llvm::PointerIntPair<const Metadata *, 2, TypeNameKind>
 
 #if SWIFT_CASTING_SUPPORTS_MUTEX
 static StaticReadWriteLock TypeNameCacheLock;
+static StaticReadWriteLock MangledToPrettyFunctionNameCacheLock;
 #endif
 
 /// Cache containing rendered names for Metadata.
 /// Access MUST be protected using `TypeNameCacheLock`.
 static Lazy<llvm::DenseMap<TypeNameCacheKey, std::pair<const char *, size_t>>>
   TypeNameCache;
 
+/// Cache containing rendered human-readable names for incoming mangled names.
+static Lazy<llvm::DenseMap<llvm::StringRef, std::pair<const char *, size_t>>>
+/// Access MUST be protected using `MangledToPrettyFunctionNameCache`.
+  MangledToPrettyFunctionNameCache;
+
 TypeNamePair
 swift::swift_getTypeName(const Metadata *type, bool qualified) {
   TypeNameCacheKey key = TypeNameCacheKey(type, qualified ? TypeNameKind::Qualified: TypeNameKind::NotQualified);
@@ -255,6 +261,141 @@ swift::swift_getMangledTypeName(const Metadata *type) {
   }
 }
 
+
+TypeNamePair
+swift::swift_getFunctionFullNameFromMangledName(
+    const char *mangledNameStart, uintptr_t mangledNameLength) {
+  llvm::StringRef mangledName(mangledNameStart, mangledNameLength);
+
+  auto &cache = MangledToPrettyFunctionNameCache.get();
+  // Attempt read-only lookup of cache entry.
+  {
+    #if SWIFT_CASTING_SUPPORTS_MUTEX
+    StaticScopedReadLock guard(MangledToPrettyFunctionNameCacheLock);
+    #endif
+
+    auto found = cache.find(mangledName);
+    if (found != cache.end()) {
+      auto result = found->second;
+      return TypeNamePair{result.first, result.second};
+    }
+  }
+
+  for (char c : mangledName) {
+    if (c >= '\x01' && c <= '\x1F')
+      return TypeNamePair{nullptr, 0};
+  }
+
+  // Read-only lookup failed, we may need to demangle and cache the entry.
+  // We have to copy the string to be able to refer to it "forever":
+  auto copy = (char *)malloc(mangledNameLength);
+  memcpy(copy, mangledNameStart, mangledNameLength);
+  mangledName = StringRef(copy, mangledNameLength);
+
+  std::string demangled;
+  StackAllocatedDemangler<1024> Dem;
+  NodePointer node = Dem.demangleSymbol(mangledName);
+  if (!node) {
+    return TypeNamePair{nullptr, 0};
+  }
+
+  // Form the demangled string from the node tree.
+  node = node->findByKind(Demangle::Node::Kind::Function, /*maxDepth=*/3);
+  if (!node || node->getNumChildren() < 3) {
+    // we normally expect Class/Identifier/Type, but don't need `Type`
+    return TypeNamePair{nullptr, 0};
+  }
+
+  // Class identifier:
+  auto clazz = node->findByKind(Demangle::Node::Kind::Class, 1);
+  if (clazz) {
+    if (auto module = clazz->findByKind(Demangle::Node::Kind::Module, 1)) {
+      demangled += module->getText();
+      demangled += ".";
+    }
+    if (auto clazzIdent = clazz->findByKind(Demangle::Node::Kind::Identifier, 1)) {
+      demangled += clazzIdent->getText();
+      demangled += ".";
+    }
+  }
+
+  // Function identifier:
+  NodePointer funcIdent = nullptr; // node == Function
+  for (size_t i = 0; i < node->getNumChildren(); ++i) {
+    if (node->getChild(i)->getKind() == Demangle::Node::Kind::Identifier) {
+      funcIdent = node->getChild(i);
+    }
+  }
+
+  // We always expect to work with functions here and they must have idents
+  if (!funcIdent) {
+    return TypeNamePair{nullptr, 0};
+  }
+  assert(funcIdent->getKind() == Demangle::Node::Kind::Identifier);
+  demangled += funcIdent->getText();
+  demangled += "(";
+
+  if (auto labelList = node->findByKind(Demangle::Node::Kind::LabelList, /*maxDepth=*/1)) {
+    if (labelList->getNumChildren()) {
+      size_t paramIdx = 0;
+      while (paramIdx < labelList->getNumChildren()) {
+        auto labelIdentifier = labelList->getChild(paramIdx++);
+        if (labelIdentifier) {
+          if (labelIdentifier->getKind() == Demangle::Node::Kind::Identifier) {
+            demangled += labelIdentifier->getText();
+            demangled += ":";
+          } else if (labelIdentifier->getKind() ==
+                     Demangle::Node::Kind::FirstElementMarker) {
+            demangled += "_:";
+          }
+        }
+      }
+    } else if (auto argumentTuple = node->findByKind(
+                   Demangle::Node::Kind::ArgumentTuple, /*maxDepth=*/5)) {
+      // LabelList was empty.
+        //
+        // The function has no labels at all, but could have some parameters...
+        // we need to check for their count, and render it as e.g. (::) for two
+        // anonymous parameters.
+        auto params = argumentTuple->getFirstChild();
+        if (auto paramsType = params->getFirstChild()) {
+          if (paramsType->getKind() != Demangle::Node::Kind::Tuple) {
+            // was a single, unnamed, parameter
+            demangled += "_:";
+          } else {
+            // there are a few parameters; find out how many
+            while (params && params->getFirstChild() &&
+                   params->getFirstChild()->getKind() !=
+                       Demangle::Node::Kind::TupleElement) {
+              params = params->getFirstChild();
+            }
+            if (params) {
+              for (size_t i = 0; i < params->getNumChildren(); ++i) {
+                demangled += "_:";
+              }
+            }
+          }
+        }
+    }
+  }
+  demangled += ")";
+
+  // We have to copy the string to be able to refer to it;
+  auto size = demangled.size();
+  auto result = (char *)malloc(size + 1);
+  memcpy(result, demangled.data(), size);
+  result[size] = 0; // 0-terminated string
+
+  {
+    #if SWIFT_CASTING_SUPPORTS_MUTEX
+    StaticScopedWriteLock guard(MangledToPrettyFunctionNameCacheLock);
+    #endif
+
+    cache.insert({mangledName, {result, size}});
+    return TypeNamePair{result, size};
+  }
+}
+
 /// Report a dynamic cast failure.
 // This is noinline to preserve this frame in stack traces.
 // We want "dynamicCastFailure" to appear in crash logs even we crash