Support generator methods in object literals, e.g. o = {*g() {...}}

mozilla · Dec 4, 2024 · 6d7b642 · 6d7b642
1 parent 8b45873
commit 6d7b642
Show file tree

Hide file tree

Showing 7 changed files with 168 additions and 81 deletions.
diff --git a/rhino/src/main/java/org/mozilla/javascript/IRFactory.java b/rhino/src/main/java/org/mozilla/javascript/IRFactory.java
@@ -33,6 +33,7 @@
 import org.mozilla.javascript.ast.FunctionNode;
 import org.mozilla.javascript.ast.GeneratorExpression;
 import org.mozilla.javascript.ast.GeneratorExpressionLoop;
+import org.mozilla.javascript.ast.GeneratorMethodDefinition;
 import org.mozilla.javascript.ast.IfStatement;
 import org.mozilla.javascript.ast.InfixExpression;
 import org.mozilla.javascript.ast.Jump;
@@ -254,6 +255,9 @@ private Node transform(AstNode node) {
                 if (node instanceof XmlLiteral) {
                     return transformXmlLiteral((XmlLiteral) node);
                 }
+                if (node instanceof GeneratorMethodDefinition) {
+                    return transformGeneratorMethodDefinition((GeneratorMethodDefinition) node);
+                }
                 throw new IllegalArgumentException("Can't transform: " + node);
         }
     }
@@ -1288,6 +1292,11 @@ private Node transformDefaultXmlNamespace(UnaryExpression node) {
         return createUnary(Token.DEFAULTNAMESPACE, child);
     }
 
+    private Node transformGeneratorMethodDefinition(GeneratorMethodDefinition node) {
+        // Unwrap the "temporary" AST node
+        return transform(node.getMethodName());
+    }
+
     /** If caseExpression argument is null it indicates a default label. */
     private static void addSwitchCase(Node switchBlock, Node caseExpression, Node statements) {
         if (switchBlock.getType() != Token.BLOCK) throw Kit.codeBug();

diff --git a/rhino/src/main/java/org/mozilla/javascript/Parser.java b/rhino/src/main/java/org/mozilla/javascript/Parser.java
@@ -41,6 +41,7 @@
 import org.mozilla.javascript.ast.FunctionNode;
 import org.mozilla.javascript.ast.GeneratorExpression;
 import org.mozilla.javascript.ast.GeneratorExpressionLoop;
+import org.mozilla.javascript.ast.GeneratorMethodDefinition;
 import org.mozilla.javascript.ast.IdeErrorReporter;
 import org.mozilla.javascript.ast.IfStatement;
 import org.mozilla.javascript.ast.InfixExpression;
@@ -3750,6 +3751,11 @@ private ObjectLiteral objectLiteral() throws IOException {
                 if (pname instanceof Name || pname instanceof StringLiteral) {
                     // For complicated reasons, parsing a name does not advance the token
                     pname.setLineColumnNumber(lineNumber(), columnNumber());
+                } else if (pname instanceof GeneratorMethodDefinition) {
+                    // Same as above
+                    ((GeneratorMethodDefinition) pname)
+                            .getMethodName()
+                            .setLineColumnNumber(lineNumber(), columnNumber());
                 }
 
                 // This code path needs to handle both destructuring object
@@ -3795,14 +3801,22 @@ private ObjectLiteral objectLiteral() throws IOException {
                         propertyName = null;
                     } else {
                         propertyName = ts.getString();
-                        ObjectProperty objectProp = methodDefinition(ppos, pname, entryKind);
+                        ObjectProperty objectProp =
+                                methodDefinition(
+                                        ppos,
+                                        pname,
+                                        entryKind,
+                                        pname instanceof GeneratorMethodDefinition);
                         pname.setJsDocNode(jsdocNode);
                         elems.add(objectProp);
                     }
                 } else {
                     pname.setJsDocNode(jsdocNode);
                     elems.add(plainProperty(pname, tt));
                 }
+                if (pname instanceof GeneratorMethodDefinition && entryKind != METHOD_ENTRY) {
+                    reportError("msg.bad.prop");
+                }
             }
 
             if (this.inUseStrictDirective
@@ -3894,6 +3908,22 @@ private AstNode objliteralProperty() throws IOException {
                 }
                 break;
 
+            case Token.MUL:
+                if (compilerEnv.getLanguageVersion() >= Context.VERSION_ES6) {
+                    int pos = ts.tokenBeg;
+                    nextToken();
+                    int lineno = lineNumber();
+                    int column = columnNumber();
+                    pname = objliteralProperty();
+
+                    pname = new GeneratorMethodDefinition(pos, ts.tokenEnd - pos, pname);
+                    pname.setLineColumnNumber(lineno, column);
+                } else {
+                    reportError("msg.bad.prop");
+                    return null;
+                }
+                break;
+
             default:
                 if (compilerEnv.isReservedKeywordAsIdentifier()
                         && TokenStream.isKeyword(
@@ -3942,8 +3972,8 @@ private ObjectProperty plainProperty(AstNode property, int ptt) throws IOExcepti
         return pn;
     }
 
-    private ObjectProperty methodDefinition(int pos, AstNode propName, int entryKind)
-            throws IOException {
+    private ObjectProperty methodDefinition(
+            int pos, AstNode propName, int entryKind, boolean isGenerator) throws IOException {
         FunctionNode fn = function(FunctionNode.FUNCTION_EXPRESSION);
         // We've already parsed the function name, so fn should be anonymous.
         Name name = fn.getFunctionName();
@@ -3963,6 +3993,9 @@ private ObjectProperty methodDefinition(int pos, AstNode propName, int entryKind
             case METHOD_ENTRY:
                 pn.setIsNormalMethod();
                 fn.setFunctionIsNormalMethod();
+                if (isGenerator) {
+                    fn.setIsES6Generator();
+                }
                 break;
         }
         int end = getNodeEnd(fn);
@@ -4477,6 +4510,8 @@ static Object getPropKey(Node id) {
         } else if (id instanceof NumberLiteral) {
             double n = ((NumberLiteral) id).getNumber();
             key = ScriptRuntime.getIndexObject(n);
+        } else if (id instanceof GeneratorMethodDefinition) {
+            key = getPropKey(((GeneratorMethodDefinition) id).getMethodName());
         } else {
             key = null; // Filled later
         }

diff --git a/rhino/src/main/java/org/mozilla/javascript/ast/GeneratorMethodDefinition.java b/rhino/src/main/java/org/mozilla/javascript/ast/GeneratorMethodDefinition.java
@@ -0,0 +1,47 @@
+/* -*- Mode: java; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 4 -*-
+ *
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.javascript.ast;
+
+import org.mozilla.javascript.Token;
+
+/**
+ * AST node for a generator method definition in an object literal, i.e. `*key() {}` in an object
+ * literal.
+ */
+public class GeneratorMethodDefinition extends AstNode {
+
+    private AstNode methodName;
+
+    public GeneratorMethodDefinition(int pos, int len, AstNode methodName) {
+        super(pos, len);
+        setType(Token.MUL);
+        setMethodName(methodName);
+    }
+
+    public AstNode getMethodName() {
+        return methodName;
+    }
+
+    public void setMethodName(AstNode methodName) {
+        assertNotNull(methodName);
+        this.methodName = methodName;
+        methodName.setParent(this);
+    }
+
+    @Override
+    public String toSource(int depth) {
+        return makeIndent(depth) + '*' + methodName.toSource(depth);
+    }
+
+    /** Visits this node, then the name. */
+    @Override
+    public void visit(NodeVisitor v) {
+        if (v.visit(this)) {
+            methodName.visit(v);
+        }
+    }
+}
diff --git a/tests/src/test/java/org/mozilla/javascript/tests/ParserTest.java b/tests/src/test/java/org/mozilla/javascript/tests/ParserTest.java
@@ -33,6 +33,7 @@
 import org.mozilla.javascript.ast.ForLoop;
 import org.mozilla.javascript.ast.FunctionCall;
 import org.mozilla.javascript.ast.FunctionNode;
+import org.mozilla.javascript.ast.GeneratorMethodDefinition;
 import org.mozilla.javascript.ast.IfStatement;
 import org.mozilla.javascript.ast.InfixExpression;
 import org.mozilla.javascript.ast.LabeledStatement;
@@ -1402,6 +1403,27 @@ public void es6Generator() {
         assertTrue(f.isES6Generator());
     }
 
+    @Test
+    public void memberFunctionGenerator() {
+        environment.setLanguageVersion(Context.VERSION_ES6);
+        AstNode root = parse("o = { *g() { return true; } }");
+        ExpressionStatement expr = (ExpressionStatement) root.getFirstChild();
+        assertTrue(expr.getExpression() instanceof Assignment);
+        assertTrue(((Assignment) expr.getExpression()).getRight() instanceof ObjectLiteral);
+        ObjectLiteral obj = (ObjectLiteral) ((Assignment) expr.getExpression()).getRight();
+        assertEquals(1, obj.getElements().size());
+        ObjectProperty g = obj.getElements().get(0);
+
+        assertTrue(g.getLeft() instanceof GeneratorMethodDefinition);
+        assertLineColumnAre(0, 7, g.getLeft());
+        AstNode genMethodName = ((GeneratorMethodDefinition) g.getLeft()).getMethodName();
+        assertTrue(genMethodName instanceof Name);
+        assertLineColumnAre(0, 8, genMethodName);
+
+        assertTrue(g.getRight() instanceof FunctionNode);
+        assertTrue(((FunctionNode) g.getRight()).isES6Generator());
+    }
+
     @Test
     public void es6GeneratorNot() {
         expectParseErrors(

diff --git a/tests/src/test/java/org/mozilla/javascript/tests/es6/GeneratorMethodTest.java b/tests/src/test/java/org/mozilla/javascript/tests/es6/GeneratorMethodTest.java
@@ -0,0 +1,10 @@
+package org.mozilla.javascript.tests.es6;
+
+import org.mozilla.javascript.Context;
+import org.mozilla.javascript.drivers.LanguageVersion;
+import org.mozilla.javascript.drivers.RhinoTest;
+import org.mozilla.javascript.drivers.ScriptTestsBase;
+
+@RhinoTest("testsrc/jstests/es6/generator-method.js")
+@LanguageVersion(Context.VERSION_ES6)
+public class GeneratorMethodTest extends ScriptTestsBase {}
diff --git a/tests/testsrc/jstests/es6/generator-method.js b/tests/testsrc/jstests/es6/generator-method.js
@@ -0,0 +1,31 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+load("testsrc/assert.js");
+
+(function generatorAsShortHandMethod() {
+	const o = {
+		h() {},
+		*g() {
+			yield 1;
+			yield 2;
+		}
+	};
+
+	const iter = o.g();
+
+	let next = iter.next();
+	assertEquals(1, next.value);
+	assertEquals(false, next.done);
+
+	next = iter.next();
+	assertEquals(2, next.value);
+	assertEquals(false, next.done);
+
+	next = iter.next();
+	assertEquals(undefined, next.value);
+	assertEquals(true, next.done);
+})();
+
+"success";