From 95d579d9de7542ff5d31b60ca40445cbf28f71cb Mon Sep 17 00:00:00 2001
From: Tom Hvitved <hvitved@github.com>
Date: Wed, 24 Apr 2024 12:22:11 +0200
Subject: [PATCH] Data flow: Fix bad join

```
Evaluated relational algebra for predicate _DataFlowImpl::Impl<HardcodedDataInterpretedAsCodeQuery::HardcodedDataInterpretedAsCodeFlow::C>::ret__#count_range@d112335l with tuple counts:
            285176  ~2%    {3} r1 = SCAN `_DataFlowDispatch::DataFlowCall.getEnclosingCallable/0#dispred#b7b78b19_DataFlowImpl::Impl<Hardcoded__#shared` OUTPUT In.1, In.0, In.2
        3265592261  ~3%    {5}    | JOIN WITH `DataFlowImpl::Impl<HardcodedDataInterpretedAsCodeQuery::HardcodedDataInterpretedAsCodeFlow::C>::returnCallEdge1/4#d02cae42_2301#join_rhs` ON FIRST 2 OUTPUT Lhs.0, Lhs.2, Rhs.2, Lhs.1, Rhs.3
             39070  ~8%    {6}    | JOIN WITH `DataFlowImplCommon::Cached::viableImplInCallContextExt/2#58e931ad` ON FIRST 3 OUTPUT Lhs.0, Lhs.3, Lhs.1, Lhs.2, Lhs.4, _
             39070  ~0%    {6}    | REWRITE WITH Out.5 := 1
                           return r1
```
---
 .../codeql/dataflow/internal/DataFlowImpl.qll         | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll b/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
index 9dfbc3e0bcd..de9f718b347 100644
--- a/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
+++ b/shared/dataflow/codeql/dataflow/internal/DataFlowImpl.qll
@@ -1142,6 +1142,14 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         )
     }
 
+    pragma[nomagic]
+    private predicate returnCallEdgeInCtx1(
+      DataFlowCallable c, SndLevelScopeOption scope, DataFlowCall call, NodeEx out, DataFlowCall ctx
+    ) {
+      returnCallEdge1(c, scope, call, out) and
+      c = viableImplInCallContextExt(call, ctx)
+    }
+
     private int ctxDispatchFanoutOnReturn(NodeEx out, DataFlowCall ctx) {
       exists(DataFlowCall call, DataFlowCallable c |
         simpleDispatchFanoutOnReturn(call, out) > 1 and
@@ -1151,8 +1159,7 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         mayBenefitFromCallContextExt(call, _) and
         result =
           count(DataFlowCallable tgt, SndLevelScopeOption scope |
-            tgt = viableImplInCallContextExt(call, ctx) and
-            returnCallEdge1(tgt, scope, call, out)
+            returnCallEdgeInCtx1(tgt, scope, call, out, ctx)
           )
       )
     }