Fix end offset of expressions in Skylark parser

This CL doesn't fix all of the problems. The end location of blocks
after an if, elif, else and for is still wrong. (I added TODOs for them)
The latter is not easy to fix: One might think that one could just
set the end location of a block to the end location of the last
statement. However, if the last statement is a "pass" statement,
this is not preserved in the AST, in particular, there's no location
for it. In the future, we probably want to preserve "pass" statements.

RELNOTES: none
PiperOrigin-RevId: 170028466

2 files changed, 75 insertions, 26 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
index b25f89be19..979c63b4e3 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Parser.java
@@ -500,7 +500,7 @@ public class Parser {
     expect(TokenKind.DOT);
     if (token.kind == TokenKind.IDENTIFIER) {
       Identifier ident = parseIdent();
-      return setLocation(new DotExpression(receiver, ident), start, token.right);
+      return setLocation(new DotExpression(receiver, ident), start, ident);
     } else {
       syntaxError(token, "expected identifier after dot");
       int end = syncTo(EXPR_TERMINATOR_SET);
@@ -677,15 +677,18 @@ public class Parser {
     }
     // This is an index/key access
     if (token.kind == TokenKind.RBRACKET) {
+      Expression expr = setLocation(new IndexExpression(receiver, startExpr), start, token.right);
       expect(TokenKind.RBRACKET);
-      return setLocation(new IndexExpression(receiver, startExpr), start, token.right);
+      return expr;
     }
     // This is a slice (or substring)
     Expression endExpr = parseSliceArgument(new Identifier("None"));
     Expression stepExpr = parseSliceArgument(new IntegerLiteral(1));
+    Expression expr =
+        setLocation(
+            new SliceExpression(receiver, startExpr, endExpr, stepExpr), start, token.right);
     expect(TokenKind.RBRACKET);
-    return setLocation(new SliceExpression(receiver, startExpr, endExpr, stepExpr),
-        start, token.right);
+    return expr;
   }
 
   /**
@@ -735,14 +738,16 @@ public class Parser {
       }
       tuple.add(parsePrimaryWithSuffix());
     }
-    return setLocation(ListLiteral.makeTuple(tuple), start, token.right);
+    return setLocation(ListLiteral.makeTuple(tuple), start, Iterables.getLast(tuple));
   }
 
   // comprehension_suffix ::= 'FOR' loop_variables 'IN' expr comprehension_suffix
   //                        | 'IF' expr comprehension_suffix
   //                        | ']'
   private Expression parseComprehensionSuffix(
-      AbstractComprehension.AbstractBuilder comprehensionBuilder, TokenKind closingBracket) {
+      AbstractComprehension.AbstractBuilder comprehensionBuilder,
+      TokenKind closingBracket,
+      int comprehensionStartOffset) {
     while (true) {
       if (token.kind == TokenKind.FOR) {
         nextToken();
@@ -758,12 +763,14 @@ public class Parser {
         // [x for x in li if (1, 2)]  # ok
         comprehensionBuilder.addIf(parseNonTupleExpression(0));
       } else if (token.kind == closingBracket) {
+        Expression expr = comprehensionBuilder.build();
+        setLocation(expr, comprehensionStartOffset, token.right);
         nextToken();
-        return comprehensionBuilder.build();
+        return expr;
       } else {
         syntaxError(token, "expected '" + closingBracket.getPrettyName() + "', 'for' or 'if'");
         syncPast(LIST_TERMINATOR_SET);
-        return makeErrorExpression(token.left, token.right);
+        return makeErrorExpression(comprehensionStartOffset, token.right);
       }
     }
   }
@@ -794,11 +801,10 @@ public class Parser {
         }
       case FOR:
         { // list comprehension
-          Expression result =
-              parseComprehensionSuffix(
-                  new ListComprehension.Builder().setOutputExpression(expression),
-                  TokenKind.RBRACKET);
-          return setLocation(result, start, token.right);
+          return parseComprehensionSuffix(
+              new ListComprehension.Builder().setOutputExpression(expression),
+              TokenKind.RBRACKET,
+              start);
         }
       case COMMA:
         {
@@ -843,12 +849,12 @@ public class Parser {
     DictionaryEntryLiteral entry = parseDictEntry();
     if (token.kind == TokenKind.FOR) {
       // Dict comprehension
-      Expression result = parseComprehensionSuffix(
+      return parseComprehensionSuffix(
           new DictComprehension.Builder()
               .setKeyExpression(entry.getKey())
               .setValueExpression(entry.getValue()),
-          TokenKind.RBRACE);
-      return setLocation(result, start, token.right);
+          TokenKind.RBRACE,
+          start);
     }
     List<DictionaryEntryLiteral> entries = new ArrayList<>();
     entries.add(entry);
@@ -959,7 +965,7 @@ public class Parser {
     // It's a tuple
     List<Expression> tuple = parseExprList(insideParens);
     tuple.add(0, expression);  // add the first expression to the front of the tuple
-    return setLocation(ListLiteral.makeTuple(tuple), start, token.right);
+    return setLocation(ListLiteral.makeTuple(tuple), start, Iterables.getLast(tuple));
   }
 
   // Equivalent to 'test' rule in Python grammar.
@@ -1000,7 +1006,7 @@ public class Parser {
     Expression expression = parseNonTupleExpression(prec + 1);
     UnaryOperatorExpression notExpression =
         new UnaryOperatorExpression(UnaryOperator.NOT, expression);
-    return setLocation(notExpression, start, token.right);
+    return setLocation(notExpression, start, expression);
   }
 
   // file_input ::= ('\n' | stmt)* EOF
@@ -1187,6 +1193,7 @@ public class Parser {
     } else {
       elseBlock = ImmutableList.of();
     }
+    // TODO(skylark-team): the end offset should be the *previous* token, not the current one.
     return setLocation(new IfStatement(thenBlocks, elseBlock), start, token.right);
   }
 
@@ -1198,6 +1205,7 @@ public class Parser {
     expect(TokenKind.COLON);
     List<Statement> thenBlock = parseSuite();
     ConditionalStatements stmt = new ConditionalStatements(expr, thenBlock);
+    // TODO(skylark-team): the end offset should be the *previous* token, not the current one.
     return setLocation(stmt, start, token.right);
   }
 
@@ -1211,6 +1219,7 @@ public class Parser {
     expect(TokenKind.COLON);
     List<Statement> block = parseSuite();
     Statement stmt = new ForStatement(new LValue(loopVar), collection, block);
+    // TODO(skylark-team): the end offset should be the *previous* token, not the current one.
     list.add(setLocation(stmt, start, token.right));
   }
 
@@ -1227,6 +1236,7 @@ public class Parser {
     expect(TokenKind.COLON);
     List<Statement> block = parseSuite();
     FunctionDefStatement stmt = new FunctionDefStatement(ident, params, signature, block);
+    // TODO(skylark-team): the end offset should be the *previous* token, not the current one.
     list.add(setLocation(stmt, start, token.right));
   }
 
diff --git a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
index edd7cc13c5..20923da4db 100644
--- a/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
+++ b/src/test/java/com/google/devtools/build/lib/syntax/ParserTest.java
@@ -21,6 +21,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.events.Location;
 import com.google.devtools.build.lib.syntax.DictionaryLiteral.DictionaryEntryLiteral;
+import com.google.devtools.build.lib.syntax.Parser.ParsingLevel;
 import com.google.devtools.build.lib.syntax.SkylarkImports.SkylarkImportSyntaxException;
 import com.google.devtools.build.lib.syntax.util.EvaluationTestCase;
 import com.google.devtools.build.lib.vfs.PathFragment;
@@ -488,8 +489,8 @@ public class ParserTest extends EvaluationTestCase {
   @Test
   public void testListPositions() throws Exception {
     String expr = "[0,f(1),2]";
+    assertExpressionLocationCorrect(expr);
     ListLiteral list = (ListLiteral) parseExpression(expr);
-    assertThat(getText(expr, list)).isEqualTo("[0,f(1),2]");
     assertThat(getText(expr, getElem(list, 0))).isEqualTo("0");
     assertThat(getText(expr, getElem(list, 1))).isEqualTo("f(1)");
     assertThat(getText(expr, getElem(list, 2))).isEqualTo("2");
@@ -498,8 +499,8 @@ public class ParserTest extends EvaluationTestCase {
   @Test
   public void testDictPositions() throws Exception {
     String expr = "{1:2,2:f(1),3:4}";
+    assertExpressionLocationCorrect(expr);
     DictionaryLiteral list = (DictionaryLiteral) parseExpression(expr);
-    assertThat(getText(expr, list)).isEqualTo("{1:2,2:f(1),3:4}");
     assertThat(getText(expr, getElem(list, 0))).isEqualTo("1:2");
     assertThat(getText(expr, getElem(list, 1))).isEqualTo("2:f(1)");
     assertThat(getText(expr, getElem(list, 2))).isEqualTo("3:4");
@@ -507,12 +508,50 @@ public class ParserTest extends EvaluationTestCase {
 
   @Test
   public void testArgumentPositions() throws Exception {
-    String stmt = "f(0,g(1,2),2)";
-    FuncallExpression f = (FuncallExpression) parseExpression(stmt);
-    assertThat(getText(stmt, f)).isEqualTo(stmt);
-    assertThat(getText(stmt, getArg(f, 0))).isEqualTo("0");
-    assertThat(getText(stmt, getArg(f, 1))).isEqualTo("g(1,2)");
-    assertThat(getText(stmt, getArg(f, 2))).isEqualTo("2");
+    String expr = "f(0,g(1,2),2)";
+    assertExpressionLocationCorrect(expr);
+    FuncallExpression f = (FuncallExpression) parseExpression(expr);
+    assertThat(getText(expr, getArg(f, 0))).isEqualTo("0");
+    assertThat(getText(expr, getArg(f, 1))).isEqualTo("g(1,2)");
+    assertThat(getText(expr, getArg(f, 2))).isEqualTo("2");
+  }
+
+  @Test
+  public void testSuffixPosition() throws Exception {
+    assertExpressionLocationCorrect("'a'.len");
+    assertExpressionLocationCorrect("'a'[0]");
+    assertExpressionLocationCorrect("'a'[0:1]");
+  }
+
+  @Test
+  public void testTuplePosition() throws Exception {
+    String input = "for a,b in []: pass";
+    ForStatement stmt = (ForStatement) parseStatement(ParsingLevel.LOCAL_LEVEL, input);
+    assertThat(getText(input, stmt.getVariable())).isEqualTo("a,b");
+    input = "for (a,b) in []: pass";
+    stmt = (ForStatement) parseStatement(ParsingLevel.LOCAL_LEVEL, input);
+    assertThat(getText(input, stmt.getVariable())).isEqualTo("(a,b)");
+    assertExpressionLocationCorrect("a, b");
+    assertExpressionLocationCorrect("(a, b)");
+  }
+
+  @Test
+  public void testComprehensionPosition() throws Exception {
+    assertExpressionLocationCorrect("[[] for x in []]");
+    assertExpressionLocationCorrect("{1: [] for x in []}");
+  }
+
+  @Test
+  public void testUnaryOperationPosition() throws Exception {
+    assertExpressionLocationCorrect("not True");
+  }
+
+  private void assertExpressionLocationCorrect(String exprStr) {
+    Expression expr = parseExpression(exprStr);
+    assertThat(getText(exprStr, expr)).isEqualTo(exprStr);
+    // Also try it with another token at the end (newline), which broke the location in the past.
+    expr = parseExpression(exprStr + "\n");
+    assertThat(getText(exprStr, expr)).isEqualTo(exprStr);
   }
 
   @Test