Preference key for newlines in nested anonymous functions

scalariform/src/main/scala/scalariform/formatter/ExprFormatter.scala

@@ -814,53 +814,83 @@ trait ExprFormatter { self: HasFormattingPreferences with AnnotationFormatter wi
         if (singleLineBlock)
           formatResult ++= format(statSeq)(newFormatterState)
         else {
-          if (statSeq.firstTokenOption.isDefined) {
-            statSeq.firstStatOpt match {
-              case Some(Expr(List(AnonymousFunction(params, _, subStatSeq)))) ⇒
-                def hasNestedAnonymousFunction(subStatSeq: StatSeq): Boolean =
-                  (for {
-                    firstStat <- subStatSeq.firstStatOpt
-                    head <- firstStat.immediateChildren.headOption
-                  } yield head.isInstanceOf[AnonymousFunction]).getOrElse(false)
-
-                val (instruction, subStatState) =
-                  if (hasNestedAnonymousFunction(subStatSeq))
-                    (CompactEnsuringGap, indentedState.indent(-1))
-                  else if (hiddenPredecessors(params.head.firstToken).containsNewline)
-                    (indentedInstruction, indentedState.indent)
-                  else
-                    (CompactEnsuringGap, indentedState)
-                formatResult = formatResult.before(statSeq.firstToken, instruction)
-                formatResult ++= format(params)
-                for (firstToken ← subStatSeq.firstTokenOption) {
-                  val instruction =
-                    if (hasNestedAnonymousFunction(subStatSeq))
-                      CompactEnsuringGap
-                    else if (hiddenPredecessors(firstToken).containsNewline || containsNewline(subStatSeq))
-                      statFormatterState(subStatSeq.firstStatOpt)(subStatState).currentIndentLevelInstruction
-                    else
-                      CompactEnsuringGap
-                  formatResult = formatResult.before(firstToken, instruction)
-                }
-                formatResult ++= format(subStatSeq)(subStatState)
-              case _ ⇒
-                val instruction = statSeq.selfReferenceOpt match {
-                  case Some((selfReference, _)) if !hiddenPredecessors(selfReference.firstToken).containsNewline ⇒
-                    CompactEnsuringGap
-                  case _ ⇒
-                    statFormatterState(statSeq.firstStatOpt)(indentedState).currentIndentLevelInstruction
-                }
-                formatResult = formatResult.before(statSeq.firstToken, instruction)
-                formatResult ++= format(statSeq)(indentedState)
-            }
-          }
+          formatResult ++= formatMultilineBlock(statSeq)(indentedInstruction)(indentedState)
           formatResult = formatResult.before(rbrace, newFormatterState.currentIndentLevelInstruction)
         }
     }
 
     formatResult
   }
 
+  private def formatMultilineBlock(statSeq: StatSeq)(indentedInstruction: IntertokenFormatInstruction)(implicit formatterState: FormatterState): FormatResult = {
+    if (statSeq.firstTokenOption.isDefined) {
+      val formatResult = asAnonymousFunction(statSeq)
+        .map(af ⇒ formatAnonymousFunction(af, statSeq)(indentedInstruction))
+        .getOrElse {
+          val instruction = statSeq.selfReferenceOpt match {
+            case Some((selfReference, _)) if !hiddenPredecessors(selfReference.firstToken).containsNewline ⇒
+              CompactEnsuringGap
+            case _ ⇒
+              statFormatterState(statSeq.firstStatOpt).currentIndentLevelInstruction
+          }
+          NoFormatResult.before(statSeq.firstToken, instruction) ++ format(statSeq)
+        }
+      formatResult
+    } else NoFormatResult
+  }
+
+  private def formatAnonymousFunction(anonymousFunction: AnonymousFunction, parentBlock: StatSeq)(indentedInstruction: IntertokenFormatInstruction)(implicit formatterState: FormatterState): FormatResult = {
+    import anonymousFunction.{ parameters, body }
+
+    var formatResult: FormatResult = NoFormatResult
+
+    val hasNestedAnonymousFunction = asAnonymousFunction(body).isDefined
+    val newlinePolicy = formattingPreferences(NewlinesAtNestedAnonymousFunctions)
+
+    val (instruction, subStatState) =
+      if (hasNestedAnonymousFunction && newlinePolicy == Prevent)
+        (CompactEnsuringGap, formatterState.indent(-1))
+      else if (hasNestedAnonymousFunction && newlinePolicy == Force)
+        (EnsureNewlineAndIndent(1, Some(parentBlock.firstToken)), formatterState.indent)
+      else if (hiddenPredecessors(parameters.head.firstToken).containsNewline)
+        (indentedInstruction, formatterState.indent)
+      else
+        (CompactEnsuringGap, formatterState)
+
+    formatResult = formatResult.before(parentBlock.firstToken, instruction)
+    formatResult ++= format(parameters)
+
+    var bodyState = subStatState
+
+    for (firstToken ← body.firstTokenOption) {
+      val (instruction, newState) =
+        if (hasNestedAnonymousFunction && newlinePolicy == Prevent)
+          (CompactEnsuringGap, subStatState)
+        else if (hasNestedAnonymousFunction && newlinePolicy == Preserve && !hiddenPredecessors(firstToken).containsNewline)
+          (CompactEnsuringGap, subStatState.indent(-1))
+        else if (hiddenPredecessors(firstToken).containsNewline || containsNewline(body))
+          (statFormatterState(body.firstStatOpt)(subStatState).currentIndentLevelInstruction, subStatState)
+        else
+          (CompactEnsuringGap, subStatState)
+      formatResult = formatResult.before(firstToken, instruction)
+      bodyState = newState
+    }
+
+    formatResult ++= format(body)(bodyState)
+
+    formatResult
+  }
+
+  private def asAnonymousFunction(subStatSeq: StatSeq): Option[AnonymousFunction] =
+    for {
+      firstStat <- subStatSeq.firstStatOpt
+      head <- firstStat.immediateChildren.headOption
+      result <- head match {
+        case af: AnonymousFunction => Some(af)
+        case _                     => None
+      }
+    } yield result
+
   private def statFormatterState(statOpt: Option[Stat])(implicit formatterState: FormatterState) = statOpt match {
     case Some(FullDefOrDcl(_, _, FunDefOrDcl(_, _, _, _, _, _, /* localDef = */ true))) if formattingPreferences(IndentLocalDefs) ⇒
       formatterState.indent

scalariform/src/main/scala/scalariform/formatter/preferences/PreferenceDescriptor.scala

@@ -94,9 +94,10 @@ object AllPreferences {
     AllowParamGroupsOnNewlines, CompactControlReadability, CompactStringConcatenation, DanglingCloseParenthesis,
     DoubleIndentClassDeclaration, DoubleIndentConstructorArguments, DoubleIndentMethodDeclaration, FirstArgumentOnNewline,
     FirstParameterOnNewline, FormatXml, IndentLocalDefs, IndentPackageBlocks, IndentSpaces, IndentWithTabs,
-    MultilineScaladocCommentsStartOnFirstLine, NewlineAtEndOfFile, PlaceScaladocAsterisksBeneathSecondAsterisk,
-    PreserveSpaceBeforeArguments, RewriteArrowSymbols, SingleCasePatternOnNewline, SpaceBeforeColon,
-    SpaceBeforeContextColon, SpaceInsideBrackets, SpaceInsideParentheses, SpacesAroundMultiImports, SpacesWithinPatternBinders
+    MultilineScaladocCommentsStartOnFirstLine, NewlineAtEndOfFile, NewlinesAtNestedAnonymousFunctions,
+    PlaceScaladocAsterisksBeneathSecondAsterisk, PreserveSpaceBeforeArguments, RewriteArrowSymbols, SingleCasePatternOnNewline,
+    SpaceBeforeColon, SpaceBeforeContextColon, SpaceInsideBrackets, SpaceInsideParentheses, SpacesAroundMultiImports,
+    SpacesWithinPatternBinders
   )
 
   val preferencesByKey: Map[String, PreferenceDescriptor[_]] =
@@ -228,6 +229,12 @@ case object NewlineAtEndOfFile extends BooleanPreferenceDescriptor {
   val defaultValue = false
 }
 
+case object NewlinesAtNestedAnonymousFunctions extends IntentPreferenceDescriptor {
+  val key = "newlinesAtNestedAnonymousFunctions"
+  val description = "When there are nested anonymous functions, add a newline after arrow at each nesting level"
+  val defaultValue: Prevent.type = Prevent
+}
+
 case object PlaceScaladocAsterisksBeneathSecondAsterisk extends BooleanPreferenceDescriptor {
   val key = "placeScaladocAsterisksBeneathSecondAsterisk"
   val description = "Place Scaladoc asterisks beneath the second asterisk in the opening '/**', as opposed to the first"

scalariform/src/test/scala/scalariform/formatter/NestedAnonymousFunctionsTest.scala

@@ -0,0 +1,85 @@
+package scalariform.formatter
+
+import scalariform.formatter.preferences._
+import scalariform.parser._
+
+
+class NestedAnonymousFunctionsTest extends AbstractFormatterTest {
+  // format: OFF
+  {
+    implicit val formattingPreferences = FormattingPreferences
+      .setPreference(NewlinesAtNestedAnonymousFunctions, Force)
+
+    """def foo: Int => Int => String = { x: Int => y: Int =>
+      |  "x"
+      |}""" ==>
+      """def foo: Int => Int => String = {
+        |  x: Int =>
+        |    y: Int =>
+        |      "x"
+        |}"""
+  }
+
+  {
+    implicit val formattingPreferences = FormattingPreferences
+      .setPreference(NewlinesAtNestedAnonymousFunctions, Prevent)
+
+    """def bar: Int => Int => String = {
+      |  x: Int =>
+      |    y: Int =>
+      |      "x"
+      |}""" ==>
+      """def bar: Int => Int => String = { x: Int => y: Int =>
+        |  "x"
+        |}"""
+  }
+
+  {
+    implicit val formattingPreferences = FormattingPreferences
+      .setPreference(NewlinesAtNestedAnonymousFunctions, Preserve)
+
+    """def baz: Int => Int => String = { x: Int =>
+      |  y: Int =>
+      |    "x"
+      |}""" ==>
+      """def baz: Int => Int => String = { x: Int =>
+        |  y: Int =>
+        |    "x"
+        |}"""
+
+    """def quux: Int => Int => String = { x: Int => y: Int =>
+      |  "x"
+      |}""" ==>
+      """def quux: Int => Int => String = { x: Int => y: Int =>
+        |  "x"
+        |}"""
+  }
+
+  {
+    implicit val formattingPreferences = FormattingPreferences
+      .setPreference(NewlinesAtNestedAnonymousFunctions, Force)
+      .setPreference(PreserveSpaceBeforeArguments, true)
+      .setPreference(AlignParameters, true)
+      .setPreference(AlignSingleLineCaseStatements, true)
+      .setPreference(DanglingCloseParenthesis, Preserve)
+
+    """def foo() = Seq.tabulate[String => String](
+      |  5) { baz: Int => _ =>
+      |    baz.toString
+      |  }""" ==>
+      """def foo() = Seq.tabulate[String => String](5) { baz: Int =>
+        |  _ =>
+        |    baz.toString
+        |}"""
+  }
+  //  format: ON
+
+  override val debug = false
+
+  def parse(parser: ScalaParser) = parser.nonLocalDefOrDcl()
+
+  type Result = FullDefOrDcl
+
+  def format(formatter: ScalaFormatter, result: Result) = formatter.format(result)(FormatterState())
+
+}