Data flow: Rework lambda flow [DO NOT MERGE]

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowNodes.qll

@@ -6,6 +6,7 @@ private import semmle.code.java.dataflow.TypeFlow
 private import semmle.code.java.dataflow.FlowSteps
 private import DataFlowPrivate
 private import DataFlowUtil
+private import DataFlowDispatch
 private import FlowSummaryImpl as FlowSummaryImpl
 private import DataFlowImplCommon as DataFlowImplCommon
 private import semmle.code.java.controlflow.Guards
@@ -68,7 +69,9 @@ private module Cached {
     TMapKeyContent() or
     TMapValueContent() or
     TCapturedVariableContent(CapturedVariable v) or
-    TSyntheticFieldContent(SyntheticField s)
+    TSyntheticFieldContent(SyntheticField s) or
+    TLambdaReturn(Method m) or
+    TLambdaArgument(Method m, ArgumentPosition pos)
 
   cached
   newtype TContentApprox =

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowPrivate.qll

@@ -464,7 +464,8 @@
   TCall(Call c) or
   TSummaryCall(SummarizedCallable c, FlowSummaryImpl::Private::SummaryNode receiver) {
     FlowSummaryImpl::Private::summaryCallbackRange(c, receiver)
-  }
+  } or
+  TLambdaSynthCall(Node node) { lambdaCreationHelper(node, _, _) }
 
 /** A call relevant for data flow. Includes both source calls and synthesized calls. */
 class DataFlowCall extends TDataFlowCall {
@@ -526,6 +527,19 @@
   override Location getLocation() { result = c.getLocation() }
 }
 
+/** A synthesized call inside a `SummarizedCallable`. */
+class LambdaSynthCall extends DataFlowCall, TLambdaSynthCall {
+  private Node node;
+
+  LambdaSynthCall() { this = TLambdaSynthCall(node) }
+
+  override DataFlowCallable getEnclosingCallable() { result.asCallable() = node.getEnclosingCallable() }
+
+  override string toString() { result = "[synthetic] call to " + node }
+
+  override Location getLocation() { result = node.getLocation() }
+}
+
 class NodeRegion instanceof BasicBlock {
   string toString() { result = "NodeRegion" }
 
@@ -585,8 +599,7 @@
 
 class LambdaCallKind = Method; // the "apply" method in the functional interface
 
-/** Holds if `creation` is an expression that creates a lambda of kind `kind` for `c`. */
-predicate lambdaCreation(Node creation, LambdaCallKind kind, DataFlowCallable c) {
+predicate lambdaCreationHelper(Node creation, LambdaCallKind kind, DataFlowCallable c) {
   exists(ClassInstanceExpr func, Interface t, FunctionalInterface interface |
     creation.asExpr() = func and
     func.getAnonymousClass().getAMethod() = c.asCallable() and
@@ -597,6 +610,12 @@
   )
 }
 
+/** Holds if `creation` is an expression that creates a lambda of kind `kind` for `c`. */
+predicate lambdaCreation(Node creation, LambdaCallKind kind, DataFlowCallable c, DataFlowCall synthCall) {
+  synthCall = TLambdaSynthCall(creation) and
+  lambdaCreationHelper(creation, kind, c)
+}
+
 /** Holds if `call` is a lambda call of kind `kind` where `receiver` is the lambda expression. */
 predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver) {
   receiver.(FlowSummaryNode).getSummaryNode() = call.(SummaryCall).getReceiver() and
@@ -766,3 +785,15 @@
   c instanceof MapKeyContent or
   c instanceof MapValueContent
 }
+
+Content getLambdaReturnContent(LambdaCallKind kind, ReturnKind k) {
+  result = TLambdaReturn(kind) and exists(k)
+}
+
+Content getLambdaArgumentContent(LambdaCallKind kind, ArgumentPosition pos) {
+  result = TLambdaArgument(kind, pos)
+}
+
+predicate isLambdaInstanceParameter(ParameterNode p) {
+  exists(DataFlowCallable c | lambdaCreationHelper(_, _, c) and p.isParameterOf(c, -1))
+}

java/ql/lib/semmle/code/java/dataflow/internal/DataFlowUtil.qll

@@ -4,6 +4,7 @@
 
 private import java
 private import DataFlowPrivate
+private import DataFlowDispatch
 private import semmle.code.java.dataflow.SSA
 private import semmle.code.java.controlflow.Guards
 private import semmle.code.java.dataflow.ExternalFlow
@@ -359,6 +360,31 @@ class SyntheticFieldContent extends Content, TSyntheticFieldContent {
   override string toString() { result = s.toString() }
 }
 
+class LambdaReturnContent extends Content, TLambdaReturn {
+  Method m;
+
+  LambdaReturnContent() { this = TLambdaReturn(m) }
+
+  override DataFlowType getType() {
+    result = getErasedRepr(m.getReturnType())
+  }
+
+  override string toString() { result = "<lambda-return>" }
+}
+
+class LambdaArgumentContent extends Content, TLambdaArgument {
+  Method m;
+  ArgumentPosition pos;
+
+  LambdaArgumentContent() {
+    this = TLambdaArgument(m, pos)
+  }
+
+  override DataFlowType getType() {
+    result = getErasedRepr(m.getParameter(pos).getType())
+  }
+  override string toString() { result = "<lambda-argument>" + pos.toString() }
+}
 /**
  * An entity that represents a set of `Content`s.
  *

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -290,7 +290,9 @@ signature module InputSig<LocationSig Location> {
   class LambdaCallKind;
 
   /** Holds if `creation` is an expression that creates a lambda of kind `kind` for `c`. */
-  predicate lambdaCreation(Node creation, LambdaCallKind kind, DataFlowCallable c);
+  predicate lambdaCreation(
+    Node creation, LambdaCallKind kind, DataFlowCallable c, DataFlowCall synthCall
+  );
 
   /** Holds if `call` is a lambda call of kind `kind` where `receiver` is the lambda expression. */
   predicate lambdaCall(DataFlowCall call, LambdaCallKind kind, Node receiver);
@@ -344,6 +346,12 @@ signature module InputSig<LocationSig Location> {
 
   /** Holds if `fieldFlowBranchLimit` should be ignored for flow going into/out of `c`. */
   default predicate ignoreFieldFlowBranchLimit(DataFlowCallable c) { none() }
+
+  Content getLambdaReturnContent(LambdaCallKind kind, ReturnKind retKind);
+
+  Content getLambdaArgumentContent(LambdaCallKind kind, ArgumentPosition pos);
+
+  predicate isLambdaInstanceParameter(ParameterNode p);
 }
 
 module Configs<LocationSig Location, InputSig<Location> Lang> {

shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll

@@ -423,7 +423,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
 
     pragma[nomagic]
     private predicate readSetEx(NodeEx node1, ContentSet c, NodeEx node2) {
-      readSet(pragma[only_bind_into](node1.asNode()), c, pragma[only_bind_into](node2.asNode())) and
+      readSet(node1, c, node2) and
       stepFilter(node1, node2)
       or
       exists(Node n |
@@ -462,8 +462,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
     private predicate storeExUnrestricted(
       NodeEx node1, Content c, NodeEx node2, DataFlowType contentType, DataFlowType containerType
     ) {
-      store(pragma[only_bind_into](node1.asNode()), c, pragma[only_bind_into](node2.asNode()),
-        contentType, containerType) and
+      store(node1, c, node2, contentType, containerType) and
       stepFilter(node1, node2)
     }
 
@@ -478,16 +477,14 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       hasReadStep(c)
     }
 
-    pragma[nomagic]
-    private predicate viableReturnPosOutEx(DataFlowCall call, ReturnPosition pos, NodeEx out) {
-      viableReturnPosOut(call, pos, out.asNode())
-    }
-
-    pragma[nomagic]
-    private predicate viableParamArgEx(DataFlowCall call, ParamNodeEx p, ArgNodeEx arg) {
-      viableParamArg(call, p.asNode(), arg.asNode())
-    }
-
+    // pragma[nomagic]
+    // private predicate viableReturnPosOutEx(DataFlowCall call, ReturnPosition pos, NodeEx out) {
+    //   viableReturnPosOut(call, pos, out.asNode())
+    // }
+    // pragma[nomagic]
+    // private predicate viableParamArgEx(DataFlowCall call, ParamNodeEx p, ArgNodeEx arg) {
+    //   viableParamArg(call, p.asNode(), arg.asNode())
+    // }
     /**
      * Holds if field flow should be used for the given configuration.
      */