Add optional chaining support in PeepholeFoldConstants pass

PiperOrigin-RevId: 325926033

Author: rishipal

src/com/google/javascript/jscomp/NodeUtil.java

@@ -50,6 +50,7 @@
 import com.google.javascript.rhino.jstype.JSType;
 import com.google.javascript.rhino.jstype.TernaryValue;
 import java.math.BigInteger;
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.EnumSet;
@@ -2027,6 +2028,45 @@ static boolean isEndOfOptChainSegment(Node n) {
     }
   }
 
+  /**
+   * Given the end of an optional chain segment changes all nodes from the end down to the start
+   * into non-optional nodes. e.g `({a})?.a.b.c.d()?.x.y.z` gets converted to
+   * `({a}).a.b.c.d()?.x.y.z` when the passed in endOfOptChainSegment is `({a})?.a.b.c.d()`.
+   */
+  static void convertToNonOptionalChainSegment(Node endOfOptChainSegment) {
+    checkArgument(isEndOfOptChainSegment(endOfOptChainSegment), endOfOptChainSegment);
+    // Since part of changing the nodes removes the isOptionalChainStart() marker we look for to
+    // know we're done, this logic is easier to read if we just find all the nodes first, then
+    // change them.
+    ArrayDeque<Node> segmentNodes = new ArrayDeque<>();
+    Node segmentNode = endOfOptChainSegment;
+    while (true) {
+      checkState(NodeUtil.isOptChainNode(segmentNode), segmentNode);
+      segmentNodes.add(segmentNode);
+      if (segmentNode.isOptionalChainStart()) {
+        break;
+      }
+      segmentNode = segmentNode.getFirstChild();
+    }
+    for (Node n : segmentNodes) {
+      n.setIsOptionalChainStart(false);
+      n.setToken(getNonOptChainToken(n.getToken()));
+    }
+  }
+
+  private static Token getNonOptChainToken(Token optChainToken) {
+    switch (optChainToken) {
+      case OPTCHAIN_CALL:
+        return Token.CALL;
+      case OPTCHAIN_GETELEM:
+        return Token.GETELEM;
+      case OPTCHAIN_GETPROP:
+        return Token.GETPROP;
+      default:
+        throw new IllegalStateException("Should be an OPTCHAIN token: " + optChainToken);
+    }
+  }
+
   /**
    * Is this node the name of a block-scoped declaration? Checks for let, const, class, or
    * block-scoped function declarations.
@@ -4025,6 +4065,9 @@ public static List<Node> findLhsNodesInNode(Node assigningParent) {
 
   /** Returns {@code true} if the node is a definition with Object.defineProperties */
   public static boolean isObjectDefinePropertiesDefinition(Node n) {
+    // We intentionally don't check optional Object.defineProperties?.() because we expect it to be
+    // defined or polyfill-ed, and to avoid changing the current invariant that this method returns
+    // true only for non-optional `Object.defineProperties()`.
     if (!n.isCall() || !n.hasXChildren(3)) {
       return false;
     }

src/com/google/javascript/jscomp/OptionalChainRewriter.java

@@ -22,7 +22,6 @@
 import com.google.javascript.jscomp.parsing.parser.FeatureSet.Feature;
 import com.google.javascript.rhino.IR;
 import com.google.javascript.rhino.Node;
-import com.google.javascript.rhino.Token;
 import java.util.ArrayDeque;
 
 /**
@@ -190,7 +189,7 @@ private Node rewriteInitialSegment(final Node fullChainStart, final Node fullCha
     checkArgument(!NodeUtil.isOptChainNode(receiverNode), receiverNode);
 
     // change the initial chain's nodes to be non-optional
-    convertToNonOptionalChainSegment(initialChainEnd);
+    NodeUtil.convertToNonOptionalChainSegment(initialChainEnd);
 
     final Node placeholder = IR.empty();
     fullChainEnd.replaceWith(placeholder);
@@ -268,41 +267,7 @@ String declareTempVarName(Node valueNode) {
     return tempVarName;
   }
 
-  /**
-   * Given the end of an optional chain segment. Change all nodes from the end down to the start
-   * into non-optional nodes.
-   */
-  private static void convertToNonOptionalChainSegment(Node endOfOptChainSegment) {
-    // Since part of changing the nodes removes the isOptionalChainStart() marker we look for to
-    // know we're done, this logic is easier to read if we just find all the nodes first, then
-    // change them.
-    final ArrayDeque<Node> segmentNodes = new ArrayDeque<>();
-    Node segmentNode = endOfOptChainSegment;
-    while (true) {
-      checkState(NodeUtil.isOptChainNode(segmentNode), segmentNode);
-      segmentNodes.add(segmentNode);
-      if (segmentNode.isOptionalChainStart()) {
-        break;
-      } else {
-        segmentNode = segmentNode.getFirstChild();
-      }
-    }
-    for (Node n : segmentNodes) {
-      n.setIsOptionalChainStart(false);
-      n.setToken(getNonOptChainToken(n.getToken()));
-    }
-  }
 
-  private static Token getNonOptChainToken(Token optChainToken) {
-    switch (optChainToken) {
-      case OPTCHAIN_CALL:
-        return Token.CALL;
-      case OPTCHAIN_GETELEM:
-        return Token.GETELEM;
-      case OPTCHAIN_GETPROP:
-        return Token.GETPROP;
-      default:
-        throw new IllegalStateException("Should be an OPTCHAIN token: " + optChainToken);
-    }
-  }
+
+
 }

src/com/google/javascript/jscomp/PeepholeFoldConstants.java

@@ -63,8 +63,10 @@ class PeepholeFoldConstants extends AbstractPeepholeOptimization {
   @Override
   Node optimizeSubtree(Node subtree) {
     switch (subtree.getToken()) {
+      case OPTCHAIN_CALL:
       case CALL:
         return tryFoldCall(subtree);
+
       case NEW:
         return tryFoldCtorCall(subtree);
 
@@ -106,10 +108,12 @@ private Node tryFoldBinaryOperator(Node subtree) {
 
     // If we've reached here, node is truly a binary operator.
     switch (subtree.getToken()) {
+      case OPTCHAIN_GETPROP:
       case GETPROP:
         return tryFoldGetProp(subtree, left, right);
 
       case GETELEM:
+      case OPTCHAIN_GETELEM:
         return tryFoldGetElem(subtree, left, right);
 
       case INSTANCEOF:
@@ -1284,7 +1288,7 @@ private Node tryFoldCtorCall(Node n) {
    *   Object.defineProperties(o, {})  ->  o
    */
   private Node tryFoldCall(Node n) {
-    checkArgument(n.isCall());
+    checkArgument(n.isCall() || n.isOptChainCall());
 
     if (NodeUtil.isObjectDefinePropertiesDefinition(n)) {
       Node srcObj = n.getLastChild();
@@ -1344,10 +1348,11 @@ private Node tryFoldInForcedStringContext(Node n) {
   }
 
   /**
-   * Try to fold array-element. e.g [1, 2, 3][10];
+   * For element access using GETLEM/OPTCHAIN_GETELEM on object literals, arrays or strings, tries
+   * to fold the prop access. e.g. folds array-element [1, 2, 3][1];
    */
   private Node tryFoldGetElem(Node n, Node left, Node right) {
-    checkArgument(n.isGetElem());
+    checkArgument(n.isGetElem() || n.isOptChainGetElem());
 
     if (left.isObjectLit()) {
       return tryFoldObjectPropAccess(n, left, right);
@@ -1364,10 +1369,12 @@ private Node tryFoldGetElem(Node n, Node left, Node right) {
   }
 
   /**
-   * Try to fold array-length. e.g [1, 2, 3].length ==> 3, [x, y].length ==> 2
+   * For prop access using GETPROP/OPTCHAIN_GETPROP on object literals, tries to fold their property
+   * access. For prop access on arrays, only tries to fold array-length. e.g [1, 2, 3].length ==> 3,
+   * [x, y].length ==> 2
    */
   private Node tryFoldGetProp(Node n, Node left, Node right) {
-    checkArgument(n.isGetProp());
+    checkArgument(n.isGetProp() || n.isOptChainGetProp());
 
     if (left.isObjectLit()) {
       return tryFoldObjectPropAccess(n, left, right);
@@ -1420,6 +1427,7 @@ private Node tryFoldArrayAccess(Node n, Node left, Node right) {
     double index = right.getDouble();
     int intIndex = (int) index;
     if (intIndex != index) {
+      // Ideally this should be caught in the check passes.
       report(INVALID_GETELEM_INDEX_ERROR, right);
       return n;
     }
@@ -1535,8 +1543,25 @@ private Node tryFoldStringArrayAccess(Node n, Node left, Node right) {
     return elem;
   }
 
+  /**
+   * Tries to fold the node `n` that's accessing an object literal's property. Bails out (skips
+   * folding and returns the same `n` node) with certainty when any of the following holds true:
+   *
+   * <ul>
+   *   <li>the access `n` is L-value
+   *   <li>the property references super
+   *   <li>the object has side-effects other than those preserved by folding the property
+   *   <li>property accessed does not exist on the object (might exist on prototype)
+   * </ul>
+   *
+   * <ul>
+   *   Examples of folding:
+   *   <li>`({a() { return 1; }})?.a` ---> `(function() { return 1; })`
+   *   <li>`({a() { return 1; }}).a()` ---> `(function() { return 1; }())`
+   * </ul>
+   */
   private Node tryFoldObjectPropAccess(Node n, Node left, Node right) {
-    checkArgument(NodeUtil.isNormalGet(n));
+    checkArgument(NodeUtil.isNormalOrOptChainGet(n));
 
     if (!left.isObjectLit() || !right.isString()) {
       return n;
@@ -1546,6 +1571,7 @@ private Node tryFoldObjectPropAccess(Node n, Node left, Node right) {
       // If GETPROP/GETELEM is used as assignment target the object literal is
       // acting as a temporary we can't fold it here:
       //    "{a:x}.a += 1" is not "x += 1"
+      checkState(!NodeUtil.isOptChainNode(n)); // optional chains can not be targets of assign
       return n;
     }
 
@@ -1590,7 +1616,7 @@ private Node tryFoldObjectPropAccess(Node n, Node left, Node right) {
     }
 
     // Didn't find a definition of the name in the object literal, it might
-    // be coming from the Object prototype
+    // be coming from the Object prototype.
     if (value == null) {
       return n;
     }
@@ -1635,6 +1661,28 @@ private Node tryFoldObjectPropAccess(Node n, Node left, Node right) {
 
     value.detach();
 
+    Node parent = n.getParent();
+
+    if (NodeUtil.isOptChainNode(parent)) {
+
+      /**
+       * If the chain continues after `n`, simply doing `n.replaceWith(value)` below would leave the
+       * subsequent nodes in the current chain segment optional, with their start `n` replaced.
+       *
+       * <p>So, we must ensure that all nodes in the chain's current segment are made non-optional.
+       *
+       * <p>This can happen for e.g.
+       *
+       * <ul>
+       *   <li>`({a() { return 1; }})?.a()` ---> `(function() { return 1; })()`. Here, parent
+       *       OPTCHAIN_CALL node must be converted to CALL.
+       *   <li>`({a() { return 1; }})?.a().b.c?.d` ---> `(function() { return 1; })().b.c?.d`. Here,
+       *       all nodes upto `({a() { return 1; }})?.a().b.c` must become non-optional
+       * </ul>
+       */
+      Node endOfCurrentChain = NodeUtil.getEndOfOptChainSegment(parent);
+      NodeUtil.convertToNonOptionalChainSegment(endOfCurrentChain);
+    }
     if (keyIsGetter) {
       value = IR.call(value);
       value.putBooleanProp(Node.FREE_CALL, true);

test/com/google/javascript/jscomp/NodeUtilTest.java

@@ -3721,6 +3721,48 @@ public void testIsBundledGoogModule_onlyIfInScript() {
     }
   }
 
+  @RunWith(JUnit4.class)
+  public static class ConvertToNonOptChainTests {
+    @Test
+    public void simpleChain() {
+      // `expr?.prop`
+      Node origChain = IR.startOptChainGetprop(IR.name("expr"), IR.string("prop"));
+      NodeUtil.convertToNonOptionalChainSegment(origChain);
+      assertThat(isChainConverted(origChain)).isTrue();
+    }
+
+    @Test
+    public void continuedChain() {
+      // `expr?.prop1.prop2`
+      Node innerGetProp = IR.startOptChainGetprop(IR.name("expr"), IR.string("pro1"));
+      Node outterGetProp = IR.continueOptChainGetprop(innerGetProp, IR.string("prop2"));
+      NodeUtil.convertToNonOptionalChainSegment(outterGetProp);
+      assertThat(isChainConverted(outterGetProp)).isTrue();
+    }
+
+    @Test
+    public void nestedChain() {
+      // `expr2[expr1 ?.prop1]?.prop2`
+      Node innerGetProp = IR.startOptChainGetprop(IR.name("expr1"), IR.string("pro1"));
+      Node getElem = IR.getelem(IR.name("expr2"), innerGetProp);
+      Node outterGetProp = IR.startOptChainGetprop(getElem, IR.string("prop2"));
+
+      NodeUtil.convertToNonOptionalChainSegment(outterGetProp);
+      assertThat(isChainConverted(outterGetProp)).isTrue();
+
+      // nested (inner) optional chain is unchanged
+      assertThat(outterGetProp.getFirstChild().getSecondChild().isOptChainGetProp()).isTrue();
+    }
+
+    // All nodes till the current chain's start are converted to non-optional.
+    private boolean isChainConverted(Node node) {
+      if (node == null) {
+        return true;
+      }
+      return !NodeUtil.isOptChainNode(node) && isChainConverted(node.getFirstChild());
+    }
+  }
+
   @RunWith(JUnit4.class)
   public static class GetStartOfOptChainTests {