From 5fb79277dc3ee2c00a65e59a5bd06786c83edeae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=93lafur=20P=C3=A1ll=20Geirsson?= <olafurpg@gmail.com>
Date: Sun, 23 Oct 2016 17:01:57 +0200
Subject: [PATCH] Infix applications to count towards valign depth, fixes #531

---
 .../org/scalafmt/internal/FormatOps.scala      | 18 ++++++++++++++++++
 .../org/scalafmt/internal/FormatWriter.scala   |  2 +-
 .../main/scala/org/scalafmt/util/TreeOps.scala |  6 ------
 core/src/test/resources/align/AlignTokens.stat | 12 ++++++++++++
 4 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/core/src/main/scala/org/scalafmt/internal/FormatOps.scala b/core/src/main/scala/org/scalafmt/internal/FormatOps.scala
index f504e8a8f6..9eae2a9aeb 100644
--- a/core/src/main/scala/org/scalafmt/internal/FormatOps.scala
+++ b/core/src/main/scala/org/scalafmt/internal/FormatOps.scala
@@ -43,6 +43,7 @@ class FormatOps(val tree: Tree, val initStyle: ScalafmtConfig) {
   val dequeueSpots = getDequeueSpots(tree) ++ statementStarts.keys
   val matchingParentheses = getMatchingParentheses(tree.tokens)
   val styleMap = new StyleMap(tokens, initStyle)
+  private val vAlignDepthCache = mutable.Map.empty[Tree, Int]
 
   @inline
   def owners(token: Token): Tree = ownersMap(hash(token))
@@ -551,4 +552,21 @@ class FormatOps(val tree: Tree, val initStyle: ScalafmtConfig) {
       case d @ Decision(t @ FormatToken(_, `close`, _), s) =>
         d.onlyNewlines
     }, close.end)
+
+  // Returns the depth of this node in the AST, used to prevent false positives.
+  final def vAlignDepth(tree: Tree): Int = {
+    vAlignDepthCache.getOrElseUpdate(tree, vAlignDepthUnCached(tree))
+  }
+
+  private final def vAlignDepthUnCached(tree: Tree): Int = {
+    val count: Int = tree match {
+      // infix applications don't count towards the length, context #531
+      case _: Term.ApplyInfix => 0
+      case _ => 1
+    }
+    tree.parent match {
+      case Some(parent) => count + vAlignDepth(parent)
+      case None => count
+    }
+  }
 }
diff --git a/core/src/main/scala/org/scalafmt/internal/FormatWriter.scala b/core/src/main/scala/org/scalafmt/internal/FormatWriter.scala
index 60c9ae8b68..eec40c73ee 100644
--- a/core/src/main/scala/org/scalafmt/internal/FormatWriter.scala
+++ b/core/src/main/scala/org/scalafmt/internal/FormatWriter.scala
@@ -216,7 +216,7 @@ class FormatWriter(formatOps: FormatOps) {
         val row2Owner = getAlignOwner(row2.formatToken)
         val row1Owner = getAlignOwner(row1.formatToken)
         def sameLengthToRoot =
-          lengthToRoot(row1Owner) == lengthToRoot(row2Owner)
+          vAlignDepth(row1Owner) == vAlignDepth(row2Owner)
         key(row1.formatToken.right) == key(row2.formatToken.right) &&
         sameLengthToRoot && {
           val eofParents = parents(owners(endOfLine.right))
diff --git a/core/src/main/scala/org/scalafmt/util/TreeOps.scala b/core/src/main/scala/org/scalafmt/util/TreeOps.scala
index f7b441a2c9..b59a33b839 100644
--- a/core/src/main/scala/org/scalafmt/util/TreeOps.scala
+++ b/core/src/main/scala/org/scalafmt/util/TreeOps.scala
@@ -31,12 +31,6 @@ object TreeOps {
   import LoggerOps._
   import TokenOps._
 
-  @tailrec
-  final def lengthToRoot(tree: Tree, accum: Int = 0): Int = tree.parent match {
-    case Some(parent) => lengthToRoot(parent, accum + 1)
-    case None => 1 + accum
-  }
-
   def getEnumStatements(enums: Seq[Enumerator]): Seq[Enumerator] = {
     val ret = Seq.newBuilder[Enumerator]
     enums.zipWithIndex.foreach {
diff --git a/core/src/test/resources/align/AlignTokens.stat b/core/src/test/resources/align/AlignTokens.stat
index df4d2e1da9..da9576b39c 100644
--- a/core/src/test/resources/align/AlignTokens.stat
+++ b/core/src/test/resources/align/AlignTokens.stat
@@ -401,3 +401,15 @@ object M {
 object M {
   implicit class TimestampOps[@specialized(Int, Long) A](val i: A) {}
 }
+<<< #531
+Seq(
+"com.github.seratch" %% "awscala"                     % "0.5.7",
+"org.scalatest"      % "scalatest_2.11"               % "3.0.0" % Test,
+"org.scalamock"      %% "scalamock-scalatest-support" % "3.2.2" % Test
+)
+>>>
+Seq(
+    "com.github.seratch" %% "awscala"                     % "0.5.7",
+    "org.scalatest"      % "scalatest_2.11"               % "3.0.0" % Test,
+    "org.scalamock"      %% "scalamock-scalatest-support" % "3.2.2" % Test
+)