Fix lineSpan (#193)

Fix lineSpan function to not throw NPE.

Author: dansanduleac

changelog/@unreleased/pr-193.v2.yml

@@ -0,0 +1,6 @@
+type: fix
+fix:
+  description: Fix bug in 0.3.24 that could sometimes cause a NPE when encountering
+    a newline right after a string concatenation expression.
+  links:
+  - https://github.com/palantir/palantir-java-format/pull/193

palantir-java-format/src/main/java/com/palantir/javaformat/java/JavaInputAstVisitor.java

@@ -49,6 +49,7 @@
 import com.google.common.base.Verify;
 import com.google.common.collect.HashMultiset;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableRangeMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
@@ -84,7 +85,6 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.IntFunction;
 import java.util.regex.Pattern;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
@@ -3304,13 +3304,24 @@ private Integer actualColumn(ExpressionTree expression) {
 
     /** How many lines does this node take up in the input. Returns at least 1. */
     int lineSpan(Tree node) {
-        IntFunction<Integer> lineNumberAt = tokenPosition -> {
-            Input.Token token = builder.getInput().getPositionTokenMap().get(tokenPosition);
-            return builder.getInput().getLineNumber(token.getTok().getPosition());
-        };
-        return lineNumberAt.apply(getEndPosition(node, getCurrentPath()))
-                - lineNumberAt.apply(getStartPosition(node))
-                + 1;
+        ImmutableRangeMap<Integer, ? extends Input.Token> positionTokenMap =
+                builder.getInput().getPositionTokenMap();
+
+        int startPosition = getStartPosition(node);
+        int endPosition = getEndPosition(node, getCurrentPath());
+        // The last token will not be in the range map if it's whitespace, because of JavaOutput.endTok's filtering.
+        // Thus, we go back until we find a "real" last token.
+        // If all tokens down from endPosition are null, then we are guaranteed stop at startPosition.
+        while (endPosition > startPosition && positionTokenMap.get(endPosition) == null) {
+            endPosition--;
+        }
+        Input.Token startToken = positionTokenMap.get(startPosition);
+        Input.Token endToken = positionTokenMap.get(endPosition);
+        return lineNumberAt(endToken) - lineNumberAt(startToken) + 1;
+    }
+
+    private int lineNumberAt(Input.Token token) {
+        return builder.getInput().getLineNumber(token.getTok().getPosition());
     }
 
     /** Returns true if {@code atLeastM} of the expressions in the given column are the same kind. */

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-10.input

@@ -16,5 +16,12 @@ class Palantir10 {
                     + "bar" + THING + "baz");
         foo("foo"
                 + "bar" + THING + "baz");
+
+        // This causes an NPE if we're not careful when getting the last token of the binary expression
+        if (!partialDataSourceType.equals(dataSourceType)) {
+            throw new ClientException("You may not use a " + partialDataSourceType
+                    + " partial in a " + dataSourceType + " query. Caused by partial \"" + partialName + "\""
+            );
+        }
     }
 }

palantir-java-format/src/test/resources/com/palantir/javaformat/java/testdata/palantir-10.output

@@ -13,5 +13,11 @@ class Palantir10 {
         // Test that we still reflow short string concatenations (spanning exactly 2 lines)
         foo("foo" + "bar" + THING + "baz");
         foo("foo" + "bar" + THING + "baz");
+
+        // This causes an NPE if we're not careful when getting the last token of the binary expression
+        if (!partialDataSourceType.equals(dataSourceType)) {
+            throw new ClientException("You may not use a " + partialDataSourceType + " partial in a " + dataSourceType
+                    + " query. Caused by partial \"" + partialName + "\"");
+        }
     }
 }