From 263bf5941df0481e5d74ec48cbf792dd20245b9a Mon Sep 17 00:00:00 2001
From: kevinoconnor <kevinoconnor@google.com>
Date: Wed, 9 Nov 2016 16:14:11 -0800
Subject: [PATCH] Propagate casts during direct inlining of returns

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=138699028
---
 .../javascript/jscomp/ConformanceRules.java   |  6 +--
 .../javascript/jscomp/FunctionInjector.java   |  3 ++
 .../google/javascript/jscomp/NodeUtil.java    | 11 +++++
 .../javascript/jscomp/NodeUtilTest.java       | 46 ++++++++++++++++++-
 4 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/src/com/google/javascript/jscomp/ConformanceRules.java b/src/com/google/javascript/jscomp/ConformanceRules.java
index 3d29f254589..4abf4ae55e4 100644
--- a/src/com/google/javascript/jscomp/ConformanceRules.java
+++ b/src/com/google/javascript/jscomp/ConformanceRules.java
@@ -310,13 +310,9 @@ protected boolean isAssertionCall(Node n) {
       return false;
     }
 
-    protected boolean wasCast(Node n) {
-      return n.getTypeIBeforeCast() != null;
-    }
-
     protected boolean isTypeImmediatelyTightened(Node n) {
       Node parent = n.getParent();
-      return wasCast(n) || isAssertionCall(parent);
+      return NodeUtil.wasCasted(n) || isAssertionCall(parent);
     }
 
     protected boolean isUsed(Node n) {
diff --git a/src/com/google/javascript/jscomp/FunctionInjector.java b/src/com/google/javascript/jscomp/FunctionInjector.java
index 694859db445..aa3c9bc0296 100644
--- a/src/com/google/javascript/jscomp/FunctionInjector.java
+++ b/src/com/google/javascript/jscomp/FunctionInjector.java
@@ -290,6 +290,9 @@ private Node inlineReturnValue(Reference ref, Node fnNode) {
       newExpression = safeReturnNode.removeFirstChild();
     }
 
+    // If the call site had a cast ensure it's persisted to the new expression that replaces it.
+    NodeUtil.maybePropagateCast(callNode, newExpression);
+
     callParentNode.replaceChild(callNode, newExpression);
     return newExpression;
   }
diff --git a/src/com/google/javascript/jscomp/NodeUtil.java b/src/com/google/javascript/jscomp/NodeUtil.java
index 8e3a9b6b671..18e03ec30d5 100644
--- a/src/com/google/javascript/jscomp/NodeUtil.java
+++ b/src/com/google/javascript/jscomp/NodeUtil.java
@@ -4868,4 +4868,15 @@ static boolean isGetterOrSetter(Node propNode) {
   static boolean isCallTo(Node n, String qualifiedName) {
     return n.isCall() && n.getFirstChild().matchesQualifiedName(qualifiedName);
   }
+
+  static boolean wasCasted(Node node) {
+    return node.getProp(Node.TYPE_BEFORE_CAST) != null;
+  }
+
+  static void maybePropagateCast(Node source, Node destination) {
+    if (wasCasted(source)) {
+      destination.putProp(Node.TYPE_BEFORE_CAST, source.getProp(Node.TYPE_BEFORE_CAST));
+      destination.setJSType(source.getJSType());
+    }
+  }
 }
diff --git a/test/com/google/javascript/jscomp/NodeUtilTest.java b/test/com/google/javascript/jscomp/NodeUtilTest.java
index 49eb0def2b6..789d68d756e 100644
--- a/test/com/google/javascript/jscomp/NodeUtilTest.java
+++ b/test/com/google/javascript/jscomp/NodeUtilTest.java
@@ -25,6 +25,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.javascript.jscomp.CompilerOptions.LanguageMode;
+import com.google.javascript.jscomp.type.ReverseAbstractInterpreter;
+import com.google.javascript.jscomp.type.SemanticReverseAbstractInterpreter;
 import com.google.javascript.rhino.IR;
 import com.google.javascript.rhino.JSTypeExpression;
 import com.google.javascript.rhino.Node;
@@ -40,15 +42,29 @@
  */
 public final class NodeUtilTest extends TestCase {
 
-  private static Node parse(String js) {
+  private static Node parse(String js, boolean typeCheck) {
     Compiler compiler = new Compiler();
     compiler.initCompilerOptionsIfTesting();
     compiler.getOptions().setLanguageIn(LanguageMode.ECMASCRIPT6);
     Node n = compiler.parseTestCode(js);
+    if (typeCheck) {
+      ReverseAbstractInterpreter rai =
+          new SemanticReverseAbstractInterpreter(compiler.getTypeRegistry());
+      new TypeCheck(compiler, rai, compiler.getTypeRegistry())
+          .processForTesting(n.getFirstChild(), n.getLastChild());
+    }
     assertThat(compiler.getErrors()).isEmpty();
     return n;
   }
 
+  private static Node parse(String js) {
+    return parse(js, false);
+  }
+
+  private static Node parseAndTypeCheck(String js) {
+    return parse(js, true);
+  }
+
   static Node getNode(String js) {
     Node root = parse("var a=(" + js + ");");
     Node expr = root.getFirstChild();
@@ -2443,6 +2459,34 @@ public void testIsObjectDefinePropertyDefinition() {
         getCallNode("Object.defineProperty();")));
   }
 
+  public void testWasCasted() {
+    assertFalse(NodeUtil.wasCasted(getCallNode(parseAndTypeCheck("var x = foo();"))));
+    assertTrue(
+        NodeUtil.wasCasted(
+            getCallNode(parseAndTypeCheck("var x = /** @type {string} */ (foo());"))));
+  }
+
+  public void testMaybePropagateCastTo() {
+    Node ast = parseAndTypeCheck("var x = /** @type {string} */ (foo());");
+    Node x = getNameNode(ast, "x");
+    Node call = getCallNode(ast);
+
+    assertFalse(NodeUtil.wasCasted(x));
+    NodeUtil.maybePropagateCast(call, x);
+    assertTrue(NodeUtil.wasCasted(x));
+  }
+
+  public void testMaybePropagateCastTo_noExistingCast() {
+    Node ast = parseAndTypeCheck("var x = foo();");
+    Node x = getNameNode(ast, "x");
+    Node call = getCallNode(ast);
+
+    assertFalse(NodeUtil.wasCasted(x));
+    NodeUtil.maybePropagateCast(call, x);
+    assertFalse(NodeUtil.wasCasted(x));
+  }
+
+
   private boolean executedOnceTestCase(String code) {
     Node ast = parse(code);
     Node nameNode = getNameNode(ast, "x");