print types in subscripts

Author: benpious

lib/Demangling/NodePrinter.cpp

@@ -21,6 +21,7 @@
 #include <cassert>
 #include <cstdio>
 #include <cstdlib>
+#include <vector>
 
 using namespace swift;
 using namespace Demangle;
@@ -3211,6 +3212,32 @@ void NodePrinter::printEntityType(NodePointer Entity, NodePointer type,
   }
 }
 
+NodePointer matchSequenceOfKinds(NodePointer start,
+                                 std::vector<std::tuple<Node::Kind, size_t>> pattern) {
+  if (start != nullptr) {
+    NodePointer current = start;
+    size_t idx = 0;
+    while (idx < pattern.size()) {
+      std::tuple<Node::Kind, size_t> next = pattern[idx];
+      idx += 1;
+      NodePointer nextChild = current->getChild(std::get<1>(next));
+      if (nextChild != nullptr &&
+          nextChild->getKind() == std::get<0>(next)) {
+        current = nextChild;
+      } else {
+        return  nullptr;
+      }
+    }
+    if (idx == pattern.size()) {
+      return current;
+    } else {
+      return nullptr;
+    }
+  } else {
+    return nullptr;
+  }
+}
+
 std::string Demangle::keypathSourceString(const char *MangledName,
                                           size_t MangledNameLength) {
   std::string invalid = "";
@@ -3227,29 +3254,68 @@ std::string Demangle::keypathSourceString(const char *MangledName,
       switch (child->getKind()) {
         case Node::Kind::Subscript: {
           std::string subscriptText = "subscript(";
+          std::vector<std::string> argumentTypeNames;
+          // Multiple arguments case

+          NodePointer argList = matchSequenceOfKinds(child, {
+            std::make_pair(Node::Kind::Type, 2),
+            std::make_pair(Node::Kind::FunctionType, 0),
+            std::make_pair(Node::Kind::ArgumentTuple, 0),
+            std::make_pair(Node::Kind::Type, 0),
+            std::make_pair(Node::Kind::Tuple, 0),
+          });
+          if (argList != nullptr) {
+            size_t numArgumentTypes = argList->getNumChildren();
+            size_t idx = 0;
+            while (idx < numArgumentTypes) {
+              NodePointer argumentType = argList->getChild(idx);
+              idx += 1;
+              if (argumentType->getKind() == Node::Kind::TupleElement) {
+                argumentType = argumentType->getChild(0)->getChild(0)->getChild(1);
+                if (argumentType->getKind() == Node::Kind::Identifier) {
+                  argumentTypeNames.push_back(std::string(argumentType->getText()));
+                  continue;
+                }
+              }
+              argumentTypeNames.push_back("<Unknown>");
+            }
+          } else {
+            // Case where there is a single argument
+            argList = matchSequenceOfKinds(child, {
+              std::make_pair(Node::Kind::Type, 2),
+              std::make_pair(Node::Kind::FunctionType, 0),
+              std::make_pair(Node::Kind::ArgumentTuple, 0),
+              std::make_pair(Node::Kind::Type, 0),
+            });
+            if (argList != nullptr) {
+              argumentTypeNames.push_back(std::string(argList->getChild(0)->getChild(1)->getText()));
+            }
+          }
           child = child->getChild(1);
+          size_t idx = 0;
+          // There is an argument label:
           if (child != nullptr) {
-            // There is an argument label:
             if (child->getKind() == Node::Kind::LabelList) {
-              size_t idx = 0;
               size_t numChildren = child->getNumChildren();
               if (numChildren == 0) {
-                subscriptText += unlabelledArg;
+                subscriptText += unlabelledArg + argumentTypeNames[0];
               } else {
                 while (idx < numChildren) {
                   Node *argChild = child->getChild(idx);
                   idx += 1;
                   if (argChild->getKind() == Node::Kind::Identifier) {
-                    subscriptText += std::string(argChild->getText()) + ": ";
+                    subscriptText += std::string(argChild->getText()) + ": " + argumentTypeNames[idx - 1];
+                    if (idx != numChildren) {
+                      subscriptText += " ";
+                    }
                   } else if (argChild->getKind() == Node::Kind::FirstElementMarker ||
                              argChild->getKind() == Node::Kind::VariadicMarker) {
-                    subscriptText += unlabelledArg;
+                    subscriptText += unlabelledArg + argumentTypeNames[idx - 1];
                   }
                 }
               }
             }
           } else {
-            subscriptText += unlabelledArg;
+            subscriptText += unlabelledArg + argumentTypeNames[0];
           }
           return subscriptText + ")";
         }
@@ -3273,7 +3339,6 @@ std::string Demangle::keypathSourceString(const char *MangledName,
         }
         default:
           return invalid;
-          
       }
     }
   }

test/Interpreter/keypath.swift

@@ -76,11 +76,11 @@ generic_class_constrained_keypath(GenericController(5))
 // CHECK: {{\\Controller\.secondLabel!\.text|\\Controller\.<computed 0x.* \(Optional<MyLabel>\)>!\.<computed 0x.* \(String\)>}}
 print(\Controller.secondLabel!.text)
 
-// CHECK: {{\\Controller\.subscript\(_: \)|\\Controller\.<computed 0x.* \(String\)>}}
+// CHECK: {{\\Controller\.subscript\(_: String\)|\\Controller\.<computed 0x.* \(String\)>}}
 print(\Controller["abc"])
 // CHECK: \S.a
 print(\S.a)
-// CHECK: {{\\Controller\.subscript\(int: str: _: \)|\\Controller\.<computed 0x.* \(Int\)>}}
+// CHECK: {{\\Controller\.subscript\(int: Int str: String _: Int\)|\\Controller\.<computed 0x.* \(Int\)>}}
 print(\Controller[int: 0, str: "", 0])
 // CHECK: {{\\Controller\.thirdLabel|\\Controller\.<computed 0x.* \(Optional<MyLabel>\)>}}
 print(\Controller.thirdLabel)