JavaScript: Refactor parsing of generator/array comprehensions.

Max Schaefer · Max Schaefer · commit 92c8501e6741 · 2019-02-24T20:45:41.000Z
diff --git a/javascript/extractor/src/com/semmle/jcorn/CustomParser.java b/javascript/extractor/src/com/semmle/jcorn/CustomParser.java
@@ -28,6 +28,7 @@
 import com.semmle.js.ast.MemberExpression;
 import com.semmle.js.ast.NewExpression;
 import com.semmle.js.ast.Node;
+import com.semmle.js.ast.ParenthesizedExpression;
 import com.semmle.js.ast.Position;
 import com.semmle.js.ast.SourceLocation;
 import com.semmle.js.ast.Statement;
@@ -137,7 +138,9 @@ protected Expression parseExprAtom(DestructuringErrors refDestructuringErrors) {
 			this.next();
 			// check whether this is array comprehension or regular array
 			if (this.type == TokenType._for) {
-				return (Expression) this.parseComprehension(startLoc, false, null);
+				ComprehensionExpression c = this.parseComprehension(startLoc, false, null);
+				this.expect(TokenType.bracketR);
+				return this.finishNode(c);
 			}
 			List<Expression> elements;
 			if (this.type == TokenType.comma || this.type == TokenType.bracketR ||
@@ -147,7 +150,9 @@ protected Expression parseExprAtom(DestructuringErrors refDestructuringErrors) {
 				Expression firstExpr = this.parseMaybeAssign(true, refDestructuringErrors, null);
 				// check whether this is a postfix array comprehension
 				if (this.type == TokenType._for || this.type == TokenType._if) {
-					return (Expression) this.parseComprehension(startLoc, false, firstExpr);
+					ComprehensionExpression c = this.parseComprehension(startLoc, false, firstExpr);
+					this.expect(TokenType.bracketR);
+					return this.finishNode(c);
 				} else {
 					this.eat(TokenType.comma);
 					elements = new ArrayList<Expression>();
@@ -265,7 +270,7 @@ protected Expression parseMaybeAssign(boolean noIn,
 	}
 
 	// add parsing of comprehensions
-	protected Node parseComprehension(Position startLoc, boolean isGenerator, Expression body) {
+	protected ComprehensionExpression parseComprehension(Position startLoc, boolean isGenerator, Expression body) {
 		List<ComprehensionBlock> blocks = new ArrayList<ComprehensionBlock>();
 		while (this.type == TokenType._for) {
 			SourceLocation blockStart = new SourceLocation(this.startLoc);
@@ -287,9 +292,8 @@ protected Node parseComprehension(Position startLoc, boolean isGenerator, Expres
 		Expression filter = this.eat(TokenType._if) ? this.parseParenExpression() : null;
 		if (body == null)
 			body = this.parseExpression(false, null);
-		this.expect(isGenerator ? TokenType.parenR : TokenType.bracketR);
 
-		return this.finishNode(new ComprehensionExpression(new SourceLocation(startLoc), body, blocks, filter, isGenerator));
+		return new ComprehensionExpression(new SourceLocation(startLoc), body, blocks, filter, isGenerator);
 	}
 
 	@Override
@@ -303,7 +307,9 @@ protected Expression parseParenAndDistinguishExpression(boolean canBeArrow) {
 						"for".equals(input.substring(m.end(), m.end()+3)) &&
 						!Identifiers.isIdentifierChar(input.charAt(m.end()+3), true)) {
 					next();
-					return (Expression) parseComprehension(startLoc, true, null);
+					ComprehensionExpression c = parseComprehension(startLoc, true, null);
+					this.expect(TokenType.parenR);
+					return this.finishNode(c);
 				}
 			}
 		}
