JS: Also use contextual types directly for underlying types

Author: asgerf

javascript/ql/lib/semmle/javascript/dataflow/DataFlow.qll

@@ -229,7 +229,10 @@ module DataFlow {
     predicate hasUnderlyingType(string globalName) {
       Stages::TypeTracking::ref() and
       exists(NameResolution::Node type |
-        TypeResolution::valueHasType(this.getNameResolutionNode(), type) and
+        TypeResolution::valueHasType(this.getNameResolutionNode(), type)
+        or
+        TypeResolution::valueHasContextualType(this.getNameResolutionNode(), type)
+      |
         UnderlyingTypes::nodeHasUnderlyingType(type, "global", globalName)
       )
     }
@@ -243,7 +246,10 @@ module DataFlow {
       Stages::TypeTracking::ref() and
       moduleName != "global" and
       exists(NameResolution::Node type |
-        TypeResolution::valueHasType(this.getNameResolutionNode(), type) and
+        TypeResolution::valueHasType(this.getNameResolutionNode(), type)
+        or
+        TypeResolution::valueHasContextualType(this.getNameResolutionNode(), type)
+      |
         UnderlyingTypes::nodeHasUnderlyingType(type, moduleName, typeName)
       )
     }

javascript/ql/lib/semmle/javascript/internal/TypeResolution.qll

@@ -114,11 +114,22 @@ module TypeResolution {
     )
   }
 
-  private predicate contextualType(Node value, Node contextualType) {
+  predicate valueHasContextualType(Node value, Node type) {
     exists(InvokeExpr call, Function target, int i |
       callTarget(call, target) and
       value = call.getArgument(i) and
-      contextualType = target.getParameter(i).getTypeAnnotation()
+      type = target.getParameter(i).getTypeAnnotation()
+    )
+    or
+    exists(VariableDeclarator decl |
+      value = decl.getInit() and
+      type = decl.getTypeAnnotation()
+    )
+    or
+    exists(Function functionValue, Function functionType |
+      valueHasContextualType(functionValue, trackFunctionType(functionType)) and
+      value = functionValue.getAReturnedExpr() and
+      type = functionType.getReturnTypeAnnotation()
     )
   }
 
@@ -154,10 +165,10 @@ module TypeResolution {
     or
     // Contextual typing for parameters
     exists(Function lambda, Function functionType, int i |
-      contextualType(lambda, trackFunctionType(functionType))
+      valueHasContextualType(lambda, trackFunctionType(functionType))
       or
       exists(InterfaceDefinition interface |
-        contextualType(lambda, trackType(interface)) and
+        valueHasContextualType(lambda, trackType(interface)) and
         functionType = interface.getACallSignature().getBody()
       )
     |