[Demangling] Add error handling to the remangler.

Mangling can fail, usually because the Node structure has been built
incorrectly or because something isn't supported with the old remangler.
We shouldn't just terminate the program when that happens, particularly
if it happens because someone has passed bad data to the demangler.

rdar://79725187

Author: al45tair

include/swift/Demangling/Demangle.h

@@ -512,42 +512,94 @@ enum class OperatorKind {
   Infix,
 };
 
+/// A mangling error, which consists of an error code and a Node pointer
+struct ManglingError {
+  enum Code {
+    Success = 0,
+    Uninitialized,
+    TooComplex,
+    BadNodeKind,
+    BadNominalTypeKind,
+    NotAStorageNode,
+    UnsupportedNodeKind,
+    UnexpectedBuiltinVectorType,
+    UnexpectedBuiltinType,
+  };
+
+  Code        code;
+  NodePointer node;
+
+  ManglingError() : code(Uninitialized), node(nullptr) {}
+  ManglingError(Code c, NodePointer n = nullptr) : code(c), node(n) {}
+
+  bool isSuccess() const { return code == Success; }
+};
+
+/// Used as a return type for mangling functions that may fail
+template <typename T>
+class ManglingErrorOr {
+private:
+  ManglingError err_;
+  T             value_;
+
+public:
+  ManglingErrorOr() : err_() {}
+  ManglingErrorOr(ManglingError::Code code, NodePointer node = nullptr)
+    : err_(code, node) {}
+  ManglingErrorOr(const ManglingError &err) : err_(err) {}
+  ManglingErrorOr(const T &t) : err_(ManglingError::Success), value_(t) {}
+  ManglingErrorOr(T &&t) : err_(ManglingError::Success), value_(std::move(t)) {}
+
+  bool isSuccess() const { return err_.code == ManglingError::Success; }
+
+  ManglingError::Code errorCode() const { return err_.code; }
+  NodePointer errorNode() const { return err_.node; }
+
+  const T &result() const {
+    assert(isSuccess());
+    return value_;
+  }
+};
+
 /// Remangle a demangled parse tree.
-std::string mangleNode(NodePointer root);
+ManglingErrorOr<std::string> mangleNode(NodePointer root);
 
 using SymbolicResolver =
   llvm::function_ref<Demangle::NodePointer (SymbolicReferenceKind,
                                             const void *)>;
 
 /// Remangle a demangled parse tree, using a callback to resolve
 /// symbolic references.
-std::string mangleNode(NodePointer root, SymbolicResolver resolver);
+ManglingErrorOr<std::string> mangleNode(NodePointer root, SymbolicResolver resolver);
 
 /// Remangle a demangled parse tree, using a callback to resolve
 /// symbolic references.
 ///
 /// The returned string is owned by \p Factory. This means \p Factory must stay
 /// alive as long as the returned string is used.
-llvm::StringRef mangleNode(NodePointer root, SymbolicResolver resolver,
-                           NodeFactory &Factory);
+ManglingErrorOr<llvm::StringRef> mangleNode(NodePointer root,
+                                            SymbolicResolver resolver,
+                                            NodeFactory &Factory);
 
 /// Remangle in the old mangling scheme.
 ///
 /// This is only used for objc-runtime names.
-std::string mangleNodeOld(NodePointer root);
+ManglingErrorOr<std::string> mangleNodeOld(NodePointer root);
 
 /// Remangle in the old mangling scheme.
 ///
 /// This is only used for objc-runtime names.
 /// The returned string is owned by \p Factory. This means \p Factory must stay
 /// alive as long as the returned string is used.
-llvm::StringRef mangleNodeOld(NodePointer node, NodeFactory &Factory);
+ManglingErrorOr<llvm::StringRef> mangleNodeOld(NodePointer node,
+                                               NodeFactory &Factory);
 
 /// Remangle in the old mangling scheme and embed the name in "_Tt<name>_".
 ///
 /// The returned string is null terminated and owned by \p Factory. This means
 /// \p Factory must stay alive as long as the returned string is used.
-const char *mangleNodeAsObjcCString(NodePointer node, NodeFactory &Factory);
+ManglingErrorOr<const char *> mangleNodeAsObjcCString(NodePointer node,
+                                                      NodeFactory &Factory);
 
 /// Transform the node structure to a string.
 ///

include/swift/Demangling/TypeDecoder.h

@@ -608,8 +608,13 @@ class TypeDecoder {
       return decodeMangledType(genericArgs->getChild(0), depth + 1);
     }
     case NodeKind::BuiltinTypeName: {
-      auto mangledName = Demangle::mangleNode(Node);
-      return Builder.createBuiltinType(Node->getText().str(), mangledName);
+      auto mangling = Demangle::mangleNode(Node);
+      if (!mangling.isSuccess()) {
+        return MAKE_NODE_TYPE_ERROR(Node,
+                                    "failed to mangle node (%d)",
+                                    mangling.errorCode());
+      }
+      return Builder.createBuiltinType(Node->getText().str(), mangling.result());
     }
     case NodeKind::Metatype:
     case NodeKind::ExistentialMetatype: {

include/swift/Reflection/TypeRefBuilder.h

@@ -289,7 +289,7 @@ class TypeRefBuilder {
   void clearNodeFactory() { Dem.clear(); }
 
   BuiltType decodeMangledType(Node *node);
-  
+
   ///
   /// Factory methods for all TypeRef kinds
   ///
@@ -300,12 +300,20 @@ class TypeRefBuilder {
   }
 
   llvm::Optional<std::string> createTypeDecl(Node *node, bool &typeAlias) {
-    return Demangle::mangleNode(node);
+    auto mangling = Demangle::mangleNode(node);
+    if (!mangling.isSuccess()) {
+      return llvm::None;
+    }
+    return mangling.result();
   }
 
   BuiltProtocolDecl
   createProtocolDecl(Node *node) {
-    return std::make_pair(Demangle::mangleNode(node), false);
+    auto mangling = Demangle::mangleNode(node);
+    if (!mangling.isSuccess()) {
+      return llvm::None;
+    }
+    return std::make_pair(mangling.result(), false);
   }
 
   BuiltProtocolDecl
@@ -395,12 +403,16 @@ class TypeRefBuilder {
 
       return underlyingTy->subst(*this, subs);
     }
-    
+
+    auto mangling = mangleNode(opaqueDescriptor,
+                               SymbolicResolver(),
+                               Dem);
+    if (!mangling.isSuccess())
+      return nullptr;
+
     // Otherwise, build a type ref that represents the opaque type.
     return OpaqueArchetypeTypeRef::create(*this,
-                                          mangleNode(opaqueDescriptor,
-                                                     SymbolicResolver(),
-                                                     Dem),
+                                          mangling.result(),
                                           nodeToString(opaqueDescriptor),
                                           ordinal,
                                           genericArgs);

include/swift/Remote/MetadataReader.h

@@ -933,7 +933,11 @@ class MetadataReader {
       if (!node || node->getKind() != Node::Kind::Type)
         return BuiltType();
 
-      auto name = Demangle::mangleNode(node);
+      auto mangling = Demangle::mangleNode(node);
+      if (!mangling.isSuccess())
+        return BuiltType();
+      auto name = mangling.result();
+
       auto BuiltForeign = Builder.createForeignClassType(std::move(name));
       TypeCache[MetadataAddress] = BuiltForeign;
       return BuiltForeign;

lib/AST/ASTDemangler.cpp

@@ -250,8 +250,11 @@ Type ASTBuilder::resolveOpaqueType(NodePointer opaqueDescriptor,
     auto definingDecl = opaqueDescriptor->getChild(0);
     auto definingGlobal = Factory.createNode(Node::Kind::Global);
     definingGlobal->addChild(definingDecl, Factory);
-    auto mangledName = mangleNode(definingGlobal);
-    
+    auto mangling = mangleNode(definingGlobal);
+    if (!mangling.isSuccess())
+      return Type();
+    auto mangledName = mangling.result();
+
     auto moduleNode = findModuleNode(definingDecl);
     if (!moduleNode)
       return Type();
@@ -935,7 +938,10 @@ ASTBuilder::findDeclContext(NodePointer node) {
       if (!module) return nullptr;
 
       // Look up the local type by its mangling.
-      auto mangledName = Demangle::mangleNode(node);
+      auto mangling = Demangle::mangleNode(node);
+      if (!mangling.isSuccess()) return nullptr;
+      auto mangledName = mangling.result();
+
       auto decl = module->lookupLocalType(mangledName);
       if (!decl) return nullptr;
 

lib/AST/ASTMangler.cpp

@@ -547,7 +547,11 @@ std::string ASTMangler::mangleSILDifferentiabilityWitness(StringRef originalName
     Demangler demangler;
     auto *node = mangleSILDifferentiabilityWitnessAsNode(
         originalName, kind, config, demangler);
-    return mangleNode(node);
+    auto mangling = mangleNode(node);
+    if (!mangling.isSuccess()) {
+      llvm_unreachable("unexpected mangling failure");
+    }
+    return mangling.result();
   }
   beginManglingWithoutPrefix();
   appendOperator(originalName);
@@ -721,8 +725,11 @@ std::string ASTMangler::mangleObjCRuntimeName(const NominalTypeDecl *Nominal) {
   TyMangling->addChild(Ty, Dem);
   Node *NewGlobal = Dem.createNode(Node::Kind::Global);
   NewGlobal->addChild(TyMangling, Dem);
-  std::string OldName = mangleNodeOld(NewGlobal);
-  return OldName;
+  auto mangling = mangleNodeOld(NewGlobal);
+  if (!mangling.isSuccess()) {
+    llvm_unreachable("unexpected mangling failure");
+  }
+  return mangling.result();
 #endif
 }
 

lib/Basic/Mangler.cpp

@@ -176,7 +176,12 @@ void Mangler::verify(StringRef nameStr) {
     llvm::errs() << "Can't demangle: " << nameStr << '\n';
     abort();
   }
-  std::string Remangled = mangleNode(Root);
+  auto mangling = mangleNode(Root);
+  if (!mangling.isSuccess()) {
+    llvm::errs() << "Can't remangle: " << nameStr << '\n';
+    abort();
+  }
+  std::string Remangled = mangling.result();
   if (Remangled == nameStr)
     return;
 

lib/Demangling/OldRemangler.cpp

@@ -74,19 +74,18 @@ namespace {
       };
     };
 
-    void mangle(Node *node, unsigned depth) {
+    ManglingError mangle(Node *node, unsigned depth) {
       if (depth > Remangler::MaxDepth) {
-        // FIXME: error handling needs doing properly (rdar://79725187)
-        unreachable("too complex to remangle");
+        return ManglingError(ManglingError::TooComplex, node);
       }
 
       switch (node->getKind()) {
 #define NODE(ID)                                                               \
   case Node::Kind::ID:                                                         \
-    return mangle##ID(node, depth);
+    mangle##ID(node, depth); return ManglingError(ManglingError::Success);
 #include "swift/Demangling/DemangleNodes.def"
       }
-      unreachable("bad demangling tree node");
+      return ManglingError(ManglingError::BadNodeKind, node);
     }
 
     void mangleGenericArgs(Node *node, EntityContext &ctx, unsigned depth);
@@ -2403,32 +2402,41 @@ void Remangler::mangleCanonicalPrespecializedGenericTypeCachingOnceToken(
 }
 
 /// The top-level interface to the remangler.
-std::string Demangle::mangleNodeOld(NodePointer node) {
-  if (!node) return "";
+ManglingErrorOr<std::string>
+Demangle::mangleNodeOld(NodePointer node) {
+  if (!node) return std::string();
 
   NodeFactory Factory;
   Remangler remangler(Factory);
-  remangler.mangle(node, 0);
+  ManglingError err = remangler.mangle(node, 0);
+  if (!err.isSuccess())
+    return err;
 
   return remangler.str();
 }
 
-llvm::StringRef Demangle::mangleNodeOld(NodePointer node, NodeFactory &Factory) {
-  if (!node) return "";
+ManglingErrorOr<llvm::StringRef>
+Demangle::mangleNodeOld(NodePointer node, NodeFactory &Factory) {
+  if (!node) return llvm::StringRef();
 
   Remangler remangler(Factory);
-  remangler.mangle(node, 0);
+  ManglingError err = remangler.mangle(node, 0);
+  if (!err.isSuccess())
+    return err;
 
   return remangler.getBufferStr();
 }
 
-const char *Demangle::mangleNodeAsObjcCString(NodePointer node,
-                                              NodeFactory &Factory) {
+ManglingErrorOr<const char *>
+Demangle::mangleNodeAsObjcCString(NodePointer node,
+                                  NodeFactory &Factory) {
   assert(node);
 
   Remangler remangler(Factory);
   remangler.append("_Tt");
-  remangler.mangle(node, 0);
+  ManglingError err = remangler.mangle(node, 0);
+  if (!err.isSuccess())
+    return err;
   remangler.append(StringRef("_", 2)); // Include the trailing 0 char.
 
   return remangler.getBufferStr().data();