[Demangler] Improve remangler error handling

include/swift/Demangling/Demangle.h

@@ -19,12 +19,13 @@
 #ifndef SWIFT_DEMANGLING_DEMANGLE_H
 #define SWIFT_DEMANGLING_DEMANGLE_H
 
-#include <memory>
-#include <string>
+#include "swift/Demangling/NamespaceMacros.h"
+#include "swift/Runtime/Config.h"
+#include "llvm/ADT/StringRef.h"
 #include <cassert>
 #include <cstdint>
-#include "llvm/ADT/StringRef.h"
-#include "swift/Demangling/NamespaceMacros.h"
+#include <memory>
+#include <string>
 
 namespace llvm {
   class raw_ostream;
@@ -512,42 +513,113 @@ enum class OperatorKind {
   Infix,
 };
 
+/// A mangling error, which consists of an error code and a Node pointer
+struct SWIFT_NODISCARD ManglingError {
+  enum Code {
+    Success = 0,
+    AssertionFailed,
+    Uninitialized,
+    TooComplex,
+    BadNodeKind,
+    BadNominalTypeKind,
+    NotAStorageNode,
+    UnsupportedNodeKind,
+    UnexpectedBuiltinVectorType,
+    UnexpectedBuiltinType,
+    MultipleChildNodes,
+    WrongNodeType,
+    WrongDependentMemberType,
+    BadDirectness,
+    UnknownEncoding,
+    InvalidImplCalleeConvention,
+    InvalidImplDifferentiability,
+    InvalidImplFunctionAttribute,
+    InvalidImplParameterConvention,
+    InvalidMetatypeRepresentation,
+    MultiByteRelatedEntity,
+    BadValueWitnessKind,
+    NotAContextNode,
+  };
+
+  Code        code;
+  NodePointer node;
+  unsigned    line;
+
+  ManglingError() : code(Uninitialized), node(nullptr) {}
+  ManglingError(Code c) : code(c), node(nullptr), line(0) {}
+  ManglingError(Code c, NodePointer n, unsigned l) : code(c), node(n), line(l) {}
+
+  bool isSuccess() const { return code == Success; }
+};
+
+#define MANGLING_ERROR(c,n)     ManglingError((c), (n), __LINE__)
+
+/// Used as a return type for mangling functions that may fail
+template <typename T>
+class SWIFT_NODISCARD ManglingErrorOr {
+private:
+  ManglingError err_;
+  T             value_;
+
+public:
+  ManglingErrorOr() : err_() {}
+  ManglingErrorOr(ManglingError::Code code,
+                  NodePointer node = nullptr,
+                  unsigned line = 0)
+  : err_(code, node, line) {}
+  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; }
+
+  const ManglingError &error() const { return err_; }
+
+  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.
 ///
@@ -626,7 +698,7 @@ bool nodeConsumesGenericArgs(Node *node);
 
 bool isSpecialized(Node *node);
 
-NodePointer getUnspecialized(Node *node, NodeFactory &Factory);
+ManglingErrorOr<NodePointer> getUnspecialized(Node *node, NodeFactory &Factory);
 
 /// Returns true if the node \p kind refers to a context node, e.g. a nominal
 /// type or a function.

include/swift/Demangling/TypeDecoder.h

@@ -608,8 +608,14 @@ 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:%u)",
+                                    mangling.error().code,
+                                    mangling.error().line);
+      }
+      return Builder.createBuiltinType(Node->getText().str(), mangling.result());
     }
     case NodeKind::Metatype:
     case NodeKind::ExistentialMetatype: {
@@ -1373,7 +1379,11 @@ return {}; // Not Implemented!
         parent = decodeMangledType(parentContext, depth + 1).getType();
         // Remove any generic arguments from the context node, producing a
         // node that references the nominal type declaration.
-        declNode = Demangle::getUnspecialized(node, Builder.getNodeFactory());
+        auto unspec =
+            Demangle::getUnspecialized(node, Builder.getNodeFactory());
+        if (!unspec.isSuccess())
+          return TypeLookupError("Failed to unspecialize type");
+        declNode = unspec.result();
         break;
       }
     }

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/Context.cpp

@@ -216,7 +216,11 @@ std::string Context::getModuleName(llvm::StringRef mangledName) {
     }
     default:
       if (isSpecialized(node)) {
-        node = getUnspecialized(node, *D);
+        auto unspec = getUnspecialized(node, *D);
+        if (!unspec.isSuccess())
+          node = nullptr;
+        else
+          node = unspec.result();
         break;
       }
       if (isContext(node->getKind())) {

lib/Demangling/Demangler.cpp

@@ -15,10 +15,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Demangling/Demangler.h"
-#include "swift/Demangling/ManglingUtils.h"
+#include "DemanglerAssert.h"
 #include "swift/Demangling/ManglingMacros.h"
+#include "swift/Demangling/ManglingUtils.h"
 #include "swift/Demangling/Punycode.h"
 #include "swift/Strings.h"
+#include <stdio.h>
+#include <stdlib.h>
 
 using namespace swift;
 using namespace Mangle;
@@ -87,6 +90,14 @@ static bool isRequirement(Node::Kind kind) {
 // Public utility functions    //
 //////////////////////////////////
 
+SWIFT_NORETURN void swift::Demangle::failAssert(const char *file, unsigned line,
+                                                NodePointer node,
+                                                const char *expr) {
+  fprintf(stderr, "%s:%u: assertion failed for Node %p: %s", file, line, node,
+          expr);
+  abort();
+}
+
 bool swift::Demangle::isContext(Node::Kind kind) {
   switch (kind) {
 #define NODE(ID)