C++: Count return dispatch based on 2nd level scopes.

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -308,6 +308,24 @@ signature module InputSig<LocationSig Location> {
    */
   default int getAdditionalFlowIntoCallNodeTerm(ArgumentNode arg, ParameterNode p) { none() }
 
+  /**
+   * A second-level control-flow scope in a callable.
+   *
+   * This is used to provide a more fine-grained separation of a callable
+   * context for the purpose of identifying uncertain control flow. For most
+   * languages, this is not needed, as this separation is handled through
+   * virtual dispatch, but for some cases (for example, C++) this can be used to
+   * identify, for example, large top-level switch statements acting like
+   * virtual dispatch.
+   */
+  class DataFlowSecondLevelScope {
+    /** Gets a textual representation of this element. */
+    string toString();
+  }
+
+  /** Gets the second-level scope containing the node `n`, if any. */
+  default DataFlowSecondLevelScope getSecondLevelScope(Node n) { none() }
+
   bindingset[call, p, arg]
   default predicate golangSpecificParamArgFilter(
     DataFlowCall call, ParameterNode p, ArgumentNode arg

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

@@ -1113,15 +1113,33 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
       result = getAdditionalFlowIntoCallNodeTerm(arg.projectToNode(), p.projectToNode())
     }
 
+    private module SndLevelScopeOption = Option<DataFlowSecondLevelScope>;
+
+    private class SndLevelScopeOption = SndLevelScopeOption::Option;
+
     pragma[nomagic]
-    private predicate returnCallEdge1(DataFlowCallable c, DataFlowCall call, NodeEx out) {
+    private SndLevelScopeOption getScope(RetNodeEx ret) {
+      result = SndLevelScopeOption::some(getSecondLevelScope(ret.asNode()))
+      or
+      result instanceof SndLevelScopeOption::None and not exists(getSecondLevelScope(ret.asNode()))
+    }
+
+    pragma[nomagic]
+    private predicate returnCallEdge1(
+      DataFlowCallable c, SndLevelScopeOption scope, DataFlowCall call, NodeEx out
+    ) {
       exists(RetNodeEx ret |
-        flowOutOfCallNodeCand1(call, ret, _, out) and c = ret.getEnclosingCallable()
+        flowOutOfCallNodeCand1(call, ret, _, out) and
+        c = ret.getEnclosingCallable() and
+        scope = getScope(ret)
       )
     }
 
     private int simpleDispatchFanoutOnReturn(DataFlowCall call, NodeEx out) {
-      result = strictcount(DataFlowCallable c | returnCallEdge1(c, call, out))
+      result =
+        strictcount(DataFlowCallable c, SndLevelScopeOption scope |
+          returnCallEdge1(c, scope, call, out)
+        )
     }
 
     private int ctxDispatchFanoutOnReturn(NodeEx out, DataFlowCall ctx) {
@@ -1129,12 +1147,12 @@ module MakeImpl<LocationSig Location, InputSig<Location> Lang> {
         simpleDispatchFanoutOnReturn(call, out) > 1 and
         not Stage1::revFlow(out, false) and
         call.getEnclosingCallable() = c and
-        returnCallEdge1(c, ctx, _) and
+        returnCallEdge1(c, _, ctx, _) and
         mayBenefitFromCallContextExt(call, _) and
         result =
-          count(DataFlowCallable tgt |
+          count(DataFlowCallable tgt, SndLevelScopeOption scope |
             tgt = viableImplInCallContextExt(call, ctx) and
-            returnCallEdge1(tgt, call, out)
+            returnCallEdge1(tgt, scope, call, out)
           )
       )
     }