From c4dad0c32123b0ddad315950782a1231b06ab483 Mon Sep 17 00:00:00 2001
From: Marcelo Fabri <me@marcelofabri.com>
Date: Sat, 29 Jul 2017 22:19:51 +0200
Subject: [PATCH] [SourceKit] Add subscript to doc structure (SR-5035)

---
 include/swift/IDE/SyntaxModel.h               |  4 ++-
 lib/IDE/SyntaxModel.cpp                       | 12 ++++++++-
 test/IDE/structure.swift                      | 25 +++++++++++++++++++
 .../lib/SwiftLang/SwiftLangSupport.cpp        |  2 ++
 tools/swift-ide-test/swift-ide-test.cpp       |  1 +
 5 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/include/swift/IDE/SyntaxModel.h b/include/swift/IDE/SyntaxModel.h
index a8de327a90b90..7a3a077511e23 100644
--- a/include/swift/IDE/SyntaxModel.h
+++ b/include/swift/IDE/SyntaxModel.h
@@ -93,6 +93,7 @@ enum class SyntaxStructureKind : uint8_t {
   EnumCase,
   EnumElement,
   TypeAlias,
+  Subscript,
 
   ForEachStatement,
   ForStatement,
@@ -139,13 +140,14 @@ struct SyntaxStructureNode {
   std::vector<CharSourceRange> InheritedTypeRanges;
   std::vector<SyntaxStructureElement> Elements;
 
-  bool isVariable() const {
+  bool hasSubstructure() const {
     switch (Kind) {
     case SyntaxStructureKind::GlobalVariable:
     case SyntaxStructureKind::InstanceVariable:
     case SyntaxStructureKind::StaticVariable:
     case SyntaxStructureKind::ClassVariable:
     case SyntaxStructureKind::Parameter:
+    case SyntaxStructureKind::Subscript:
       return true;
     default:
       return false;
diff --git a/lib/IDE/SyntaxModel.cpp b/lib/IDE/SyntaxModel.cpp
index 932f3f61acf30..c37129c979832 100644
--- a/lib/IDE/SyntaxModel.cpp
+++ b/lib/IDE/SyntaxModel.cpp
@@ -964,6 +964,16 @@ bool ModelASTWalker::walkToDeclPre(Decl *D) {
                                    TypeAliasD->getName().getLength());
     SN.Attrs = TypeAliasD->getAttrs();
     pushStructureNode(SN, TypeAliasD);
+  } else if (auto *SubscriptD = dyn_cast<SubscriptDecl>(D)) {
+    SyntaxStructureNode SN;
+    SN.Dcl = SubscriptD;
+    SN.Kind = SyntaxStructureKind::Subscript;
+    SN.Range = charSourceRangeFromSourceRange(SM,
+                                              SubscriptD->getSourceRange());
+    SN.BodyRange = innerCharSourceRangeFromSourceRange(SM,
+                                               SubscriptD->getBracesRange());
+    SN.Attrs = SubscriptD->getAttrs();
+    pushStructureNode(SN, SubscriptD);
   }
 
   return true;
@@ -1253,7 +1263,7 @@ bool ModelASTWalker::popStructureNode() {
 
   // VarDecls are popped before we see their TypeRepr, so if we pass the token
   // nodes now they will not change from identifier to a type-identifier.
-  if (!Node.isVariable()) {
+  if (!Node.hasSubstructure()) {
     if (!passTokenNodesUntil(Node.Range.getEnd(), IncludeNodeAtLocation))
       return false;
   }
diff --git a/test/IDE/structure.swift b/test/IDE/structure.swift
index 5938294e0696e..246eddd93eceb 100644
--- a/test/IDE/structure.swift
+++ b/test/IDE/structure.swift
@@ -187,3 +187,28 @@ class A {
 
 // CHECK: <typealias>typealias <name>OtherA</name> = A</typealias>
 typealias OtherA = A
+
+class SubscriptTest {
+  subscript(index: Int) -> Int {
+    return 0
+  }
+  // CHECK: <subscript>subscript(<param>index: Int</param>) -> Int {
+  // CHECK:  return 0
+  // CHECK: }</subscript>
+
+  subscript(string: String) -> Int {
+    get {
+      return 0
+    }
+    set(value) {
+      print(value)
+    }
+  }
+  // CHECK: <subscript>subscript(<param>string: String</param>) -> Int {
+  // CHECK: get {
+  // CHECK:   return 0
+  // CHECK: }
+  // CHECK: set(<param>value</param>) {
+  // CHECK:   <call><name>print</name>(value)</call>
+  // CHECK: }</subscript>
+}
diff --git a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp
index 298d4f215ad00..707edd4e2668e 100644
--- a/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp
+++ b/tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp
@@ -377,6 +377,8 @@ UIdent SwiftLangSupport::getUIDForSyntaxStructureKind(
       return KindDeclEnumElement;
     case SyntaxStructureKind::TypeAlias:
       return KindDeclTypeAlias;
+    case SyntaxStructureKind::Subscript:
+      return KindDeclSubscript;
     case SyntaxStructureKind::Parameter:
       return KindDeclVarParam;
     case SyntaxStructureKind::ForEachStatement:
diff --git a/tools/swift-ide-test/swift-ide-test.cpp b/tools/swift-ide-test/swift-ide-test.cpp
index 6b14f28584109..e0685cad132b7 100644
--- a/tools/swift-ide-test/swift-ide-test.cpp
+++ b/tools/swift-ide-test/swift-ide-test.cpp
@@ -1070,6 +1070,7 @@ class StructureAnnotator : public ide::SyntaxModelWalker {
       case SyntaxStructureKind::EnumCase: return "enum-case";
       case SyntaxStructureKind::EnumElement: return "enum-elem";
       case SyntaxStructureKind::TypeAlias: return "typealias";
+      case SyntaxStructureKind::Subscript: return "subscript";
       case SyntaxStructureKind::Parameter: return "param";
       case SyntaxStructureKind::ForEachStatement: return "foreach";
       case SyntaxStructureKind::ForStatement: return "for";