unified: Share logic with type inference

Author: asgerf

unified/ql/lib/codeql/unified/internal/ExprPositions.qll

@@ -79,3 +79,13 @@ predicate hasResultValue(Expr e) {
   not isInAssignmentContext(e, any(AssignExpr n)) and // non-compound assignment target
   not e instanceof IdentifierLabel
 }
+
+/** Holds if `e` is returned from `callable` either because it is an operand of `return` or is implicitly returned due to language-specific semantics. */
+predicate isReturnedExpr(Callable callable, Expr e) {
+  exists(ReturnExpr ret |
+    callable = ret.getEnclosingCallable() and
+    e = ret.getValue()
+  )
+  or
+  any(ExprPositionsPlugin p).isReturnedExpr(callable, e)
+}

unified/ql/lib/codeql/unified/internal/ExprPositionsPlugin.qll

@@ -7,4 +7,6 @@ private module Plugins {
 
 class ExprPositionsPlugin extends Unit {
   predicate isInTypeContext(Expr e) { none() }
+
+  predicate isReturnedExpr(Callable callable, Expr e) { none() }
 }

unified/ql/lib/codeql/unified/internal/ExprPositionsPluginSwift.qll

@@ -3,4 +3,24 @@ private import ExprPositionsPlugin
 
 private class ExprPositionsPluginSwift extends ExprPositionsPlugin {
   override predicate isInTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
+
+  override predicate isReturnedExpr(Callable callable, Expr e) {
+    // Short-hand return from a function body, e.g. `func foo() -> String { "42" }`
+    // This happens when the body contains a single expression and this expression resolves to a value.
+    // We don't detect precisely whether the expression resolves to a value, but we restrict to cases
+    // where a spurious return would be harmless in practice.
+    (
+      callable instanceof FunctionDeclaration and
+      exists(Expr ret | ret = callable.getReturnType() | not ret.(Identifier).getValue() = "Void")
+      or
+      callable instanceof GetterDeclaration
+      or
+      callable instanceof FunctionExpr // We can't tell whether closures have a return value
+    ) and
+    exists(Block body |
+      body = callable.getBody() and
+      not exists(body.getStmt(1)) and // Check that only one statement exists
+      e = body.getStmt(0)
+    )
+  }
 }

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowGraph.qll

@@ -1,6 +1,7 @@
 private import unified
 private import AllDataFlow
 private import codeql.unified.internal.LocalNameBinding
+private import codeql.unified.internal.ExprPositions
 
 predicate step(BuilderNode node1, Step step, BuilderNode node2) {
   any(DataFlowPlugin p).step(node1, step, node2)
@@ -180,10 +181,11 @@ predicate step(BuilderNode node1, Step step, BuilderNode node2) {
     node2.isIncomingValue(fun.getNameNode())
   )
   or
-  exists(ReturnExpr ret |
-    node1.isResultValue(ret.getValue()) and
+  exists(Callable callable, Expr expr |
+    isReturnedExpr(callable, expr) and
+    node1.isResultValue(expr) and
     step.value() and
-    node2.isReturnNode(ret.getEnclosingCallable())
+    node2.isReturnNode(callable)
   )
   or
   none() // Temporarily disable compilation errors from unsatisfiable types

unified/ql/lib/codeql/unified/internal/dataflow/DataFlowPluginSwift.qll

@@ -76,22 +76,5 @@ private class SwiftDataFlowPlugin extends DataFlowPlugin {
       step.taint() and
       node2.isResultValue(call)
     )
-    or
-    // Short-hand return from a function body, e.g. `func foo() -> String { "42" }`
-    // This happens when the body contains a single expression and this expression resolves to a value.
-    // We don't detect precisely whether the expression resolves to a value, but we restrict to cases
-    // where a spurious flow edge would be harmless.
-    exists(Callable callable, Block body |
-      (
-        callable instanceof FunctionExpr or
-        callable instanceof FunctionDeclaration or
-        callable instanceof GetterDeclaration
-      ) and
-      body = callable.getBody() and
-      not exists(body.getStmt(1)) and // Check that only one statement exists
-      node1.isResultValue(body.getStmt(0)) and
-      step.value() and
-      node2.isReturnNode(callable)
-    )
   }
 }

unified/ql/lib/codeql/unified/internal/typeinference/TypeInference.qll

@@ -340,8 +340,10 @@ module Input3 implements InputSig3 {
     result = M3::inferFieldAccessReceiverTypeContextualDefault(_, receiver, path)
   }
 
-  class Return extends ReturnExpr {
-    Expr getExpr() { result = this.getValue() }
+  class Return extends Expr {
+    Return() { isReturnedExpr(_, this) }
+
+    Expr getExpr() { result = this }
   }
 
   class Parameter extends VariableDeclarationImpl, TParameterExDeclaration {