Merge pull request swiftlang#158 from dabelknap/subscripts

Add support for subscripts

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -383,8 +383,10 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: AccessorBlockSyntax) {
-    after(node.leftBrace, tokens: .break(offset: 2), .open(.consistent, 0))
-    before(node.rightBrace, tokens: .break(offset: -2), .close)
+    if !(node.parent is SubscriptDeclSyntax) {
+      after(node.leftBrace, tokens: .break(offset: 2), .open(.consistent, 0))
+      before(node.rightBrace, tokens: .break(offset: -2), .close)
+    }
     super.visit(node)
   }
 
@@ -833,6 +835,38 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
+  override func visit(_ node: SubscriptDeclSyntax) {
+    before(node.firstToken, tokens: .open(.inconsistent, 0))
+
+    if let attributes = node.attributes {
+      before(node.firstToken, tokens: .space(size: 0), .open(.consistent, 0))
+      after(attributes.lastToken, tokens: .open)
+    } else {
+      before(node.firstToken, tokens: .space(size: 0), .open(.consistent, 0), .open)
+    }
+
+    before(node.result.firstToken, tokens: .break)
+
+    before(
+      node.genericWhereClause?.firstToken,
+      tokens: .break, .open(.inconsistent, 0), .break(size: 0), .open(.consistent, 0)
+    )
+    after(node.genericWhereClause?.lastToken, tokens: .break, .close, .close)
+
+    if let accessorBlock = node.accessor {
+      if node.genericWhereClause == nil {
+        before(accessorBlock.leftBrace, tokens: .break)
+      }
+      after(accessorBlock.leftBrace, tokens: .close, .close, .break(offset: 2), .open(.consistent, 0))
+      before(accessorBlock.rightBrace, tokens: .break(offset: -2), .close)
+    } else {
+      after(node.lastToken, tokens: .close, .close)
+    }
+
+    after(node.lastToken, tokens: .close)
+    super.visit(node)
+  }
+
   override func visit(_ node: FunctionSignatureSyntax) {
     before(node.throwsOrRethrowsKeyword, tokens: .break)
     before(node.output?.firstToken, tokens: .break)
@@ -951,10 +985,6 @@ private final class TokenStreamCreator: SyntaxVisitor {
     super.visit(node)
   }
 
-  override func visit(_ node: SubscriptDeclSyntax) {
-    super.visit(node)
-  }
-
   override func visit(_ node: TypealiasDeclSyntax) {
     super.visit(node)
   }

Tests/SwiftFormatPrettyPrintTests/SubscriptDeclTests.swift

@@ -0,0 +1,236 @@
+public class SubscriptDeclTests: PrettyPrintTestCase {
+  public func testBasicSubscriptDeclarations() {
+    let input =
+      """
+      struct MyStruct {
+        subscript(index: Int) -> Int {
+          return self.values[index]
+        }
+        subscript(row: Int, col: Int) -> Int {
+          return self.values[row][col]
+        }
+        subscript(index: Int) -> Int {
+          get { return self.value[index] }
+          set(newValue) { self.value[index] = newValue } 
+        }
+      }
+      """
+
+    let expected =
+      """
+      struct MyStruct {
+        subscript(index: Int) -> Int {
+          return self.values[index]
+        }
+        subscript(row: Int, col: Int) -> Int {
+          return self.values[row][col]
+        }
+        subscript(index: Int) -> Int {
+          get { return self.value[index] }
+          set(newValue) { self.value[index] = newValue }
+        }
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 50)
+  }
+
+  public func testSubscriptGenerics() {
+    let input =
+      """
+      struct MyStruct {
+        subscript<T>(index: T) -> Double {
+          return 1.23
+        }
+        subscript<S, T>(row: S, col: T) -> Double {
+          return self.values[row][col]
+        }
+        subscript<LongTypeName1, LongTypeName2, LongTypeName3>(
+          var1: LongTypeName1,
+          var2: LongTypeName2,
+          var3: LongTypeName3
+        ) -> Int {
+          return self.values[var1][var2][var3]
+        }
+      }
+      """
+
+    let expected =
+      """
+      struct MyStruct {
+        subscript<T>(index: T) -> Double { return 1.23 }
+        subscript<S, T>(row: S, col: T) -> Double {
+          return self.values[row][col]
+        }
+        subscript<
+          LongTypeName1,
+          LongTypeName2,
+          LongTypeName3
+        >(
+          var1: LongTypeName1,
+          var2: LongTypeName2,
+          var3: LongTypeName3
+        ) -> Int {
+          return self.values[var1][var2][var3]
+        }
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 50)
+  }
+
+  public func testSubscriptGenericWhere() {
+    let input =
+      """
+      struct MyStruct {
+        subscript<Elements: Collection, Element>(
+          var1: Element,
+          var2: Elements
+        ) -> Double where Elements.Element == Element {
+          return 1.23
+        }
+        subscript<Elements: Collection, Element>(
+          var1: Element,
+          var2: Elements
+        ) -> Double
+        where
+          Elements.Element == Element,
+          Element: Equatable
+        { return 1.23 }
+      }
+      """
+
+    let expected =
+      """
+      struct MyStruct {
+        subscript<Elements: Collection, Element>(
+          var1: Element,
+          var2: Elements
+        ) -> Double where Elements.Element == Element {
+          return 1.23
+        }
+        subscript<Elements: Collection, Element>(
+          var1: Element,
+          var2: Elements
+        ) -> Double
+        where
+          Elements.Element == Element,
+          Element: Equatable
+        { return 1.23 }
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 50)
+  }
+
+  public func testSubscriptAttributes() {
+    let input =
+      """
+      struct MyStruct {
+        @discardableResult subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult @objc subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult @objc @inlinable subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult
+        @objc
+        @inlinable
+        @available(swift 4.0)
+        subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+      }
+      """
+
+    let expected =
+      """
+      struct MyStruct {
+        @discardableResult subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult @objc subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult @objc @inlinable subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+        @discardableResult
+        @objc
+        @inlinable
+        @available(swift 4.0)
+        subscript(index: Int) -> Int {
+          let a = 123
+          return a
+        }
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 70)
+  }
+
+  public func testSubscriptFullWrap() {
+    let input =
+    """
+    struct MyStruct {
+      @discardableResult @objc
+      subscript<ManyElements: Collection, Element>(var1: Element, var2: ManyElements) -> ManyElements.Index? where ManyElements.Element == Element, Element: Equatable {
+        get {
+          let out = vals[var1][var2]
+          return out
+        }
+        set(newValue) {
+          let tmp = compute(newValue)
+          vals[var1][var2] = tmp
+        }
+      }
+    }
+    """
+
+    let expected =
+    """
+    struct MyStruct {
+      @discardableResult
+      @objc
+      subscript<
+        ManyElements: Collection,
+        Element
+      >(
+        var1: Element,
+        var2: ManyElements
+      ) -> ManyElements.Index?
+      where
+        ManyElements.Element ==
+        Element,
+        Element: Equatable
+      {
+        get {
+          let out = vals[var1][var2]
+          return out
+        }
+        set(newValue) {
+          let tmp = compute(newValue)
+          vals[var1][var2] = tmp
+        }
+      }
+    }
+
+    """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 34)
+  }
+}