Merge pull request swiftlang#105 from dabelknap/function-decl

Add wrapping for `where` clauses and generic parameters to function declarations

Author: allevato

Sources/PrettyPrint/PrettyPrint.swift

@@ -97,6 +97,7 @@ public class PrettyPrinter {
   /// current line it is printing on. If a token exceeds the remaning space, we break to a new line,
   /// and apply the appropriate level of indentation.
   private func printToken(token: Token, length: Int) {
+    assert(length >= 0, "Token lengths must be positive")
     switch token {
 
     case .open(let breaktype, let offset):

Sources/PrettyPrint/TokenStreamCreator.swift

@@ -115,6 +115,7 @@ private final class TokenStreamCreator: SyntaxVisitor {
   override func visit(_ node: FunctionParameterSyntax) {
     before(node.firstToken, tokens: .open)
     after(node.colon, tokens: .break)
+    before(node.secondName, tokens: .break)
 
     if let trailingComma = node.trailingCommaWorkaround {
       after(trailingComma, tokens: .close, .break)
@@ -281,8 +282,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: GenericParameterClauseSyntax) {
-    after(node.leftAngleBracket, tokens: .open(.consistent, 2), .break(size: 0))
-    before(node.rightAngleBracket, tokens: .break(size: 0), .close)
+    after(node.leftAngleBracket, tokens: .break(size: 0, offset: 2), .open(.consistent, 0))
+    before(node.rightAngleBracket, tokens: .break(size: 0, offset: -2), .close)
     super.visit(node)
   }
 
@@ -466,8 +467,13 @@ private final class TokenStreamCreator: SyntaxVisitor {
     after(node.modifiers?.lastToken, tokens: .break)
     after(node.funcKeyword, tokens: .break)
 
+    before(node.genericWhereClause?.firstToken, tokens: .break, .open(.consistent, 0))
+    after(node.genericWhereClause?.lastToken, tokens: .break, .close)
+
     if let body = node.body {
-      before(body.leftBrace, tokens: .break)
+      if node.genericWhereClause == nil {
+        before(body.leftBrace, tokens: .break)
+      }
       after(body.leftBrace, tokens: .newline(offset: 2), .open(.consistent, 0))
       before(body.rightBrace, tokens: .newline(offset: -2), .close)
     }
@@ -657,7 +663,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: GenericParameterSyntax) {
+    before(node.firstToken, tokens: .open)
     after(node.colon, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -723,7 +731,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: GenericWhereClauseSyntax) {
+    before(node.whereKeyword, tokens: .open(.consistent, 2))
     after(node.whereKeyword, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -744,6 +754,10 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: SameTypeRequirementSyntax) {
+    before(node.firstToken, tokens: .open)
+    before(node.equalityToken, tokens: .break)
+    after(node.equalityToken, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 
@@ -800,6 +814,9 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: ConformanceRequirementSyntax) {
+    before(node.firstToken, tokens: .open)
+    after(node.colon, tokens: .break)
+    after(node.lastToken, tokens: .close)
     super.visit(node)
   }
 

Tests/PrettyPrinterTests/FunctionDeclTests.swift

@@ -67,6 +67,10 @@ public class FunctionDeclTests: PrettyPrintTestCase {
     func myFun<S, T>(var1: S, var2: T) {
       print("Hello World")
     }
+
+    func longerNameFun<ReallyLongTypeName: Conform, TypeName>(var1: ReallyLongTypeNAme, var2: TypeName) {
+      let a = 123
+    }
     """
 
     let expected =
@@ -75,8 +79,96 @@ public class FunctionDeclTests: PrettyPrintTestCase {
       print("Hello World")
     }
 
+    func longerNameFun<
+      ReallyLongTypeName: Conform,
+      TypeName
+    >(
+      var1: ReallyLongTypeNAme,
+      var2: TypeName
+    ) {
+      let a = 123
+    }
+
     """
 
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 40)
   }
+
+  public func testFunctionWhereClause() {
+    let input =
+    """
+    public func index<Elements: Collection, Element>(
+      of element: Element,
+      in collection: Elements
+    ) -> Elements.Index? where Elements.Element == Element {
+      let a = 123
+      let b = "abc"
+    }
+
+    public func index<Elements: Collection, Element>(
+      of element: Element,
+      in collection: Elements
+    ) -> Elements.Index? where Elements.Element == Element, Element: Equatable {
+      let a = 123
+      let b = "abc"
+    }
+    """
+
+    let expected =
+    """
+    public func index<Elements: Collection, Element>(
+      of element: Element,
+      in collection: Elements
+    ) -> Elements.Index? where Elements.Element == Element {
+      let a = 123
+      let b = "abc"
+    }
+
+    public func index<Elements: Collection, Element>(
+      of element: Element,
+      in collection: Elements
+    ) -> Elements.Index?
+    where
+      Elements.Element == Element,
+      Element: Equatable
+    {
+      let a = 123
+      let b = "abc"
+    }
+
+    """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 60)
+  }
+
+  public func testFunctionFullWrap() {
+    let input =
+    """
+    public func index<Elements: Collection, Element>(of element: Element, in collection: Elements) -> Elements.Index? where Elements.Element == Element, Element: Equatable {
+      let a = 123
+      let b = "abc"
+    }
+    """
+
+    let expected =
+    """
+    public func index<
+      Elements: Collection,
+      Element
+    >(
+      of element: Element,
+      in collection: Elements
+    ) -> Elements.Index?
+    where
+      Elements.Element == Element,
+      Element: Equatable
+    {
+      let a = 123
+      let b = "abc"
+    }
+
+    """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 30)
+  }
 }