Merge pull request #18633 from aschackmull/dataflow/refactor-flowstate

Dataflow: Refactor FlowState to be paired with Node

shared/dataflow/codeql/dataflow/DataFlow.qll

@@ -643,6 +643,7 @@ private module PathGraphSigMod {
 module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
   private import Lang
   private import internal.DataFlowImpl::MakeImpl<Location, Lang>
+  private import internal.DataFlowImplStage1::MakeImplStage1<Location, Lang>
   import Configs<Location, Lang>
 
   /**
@@ -700,7 +701,13 @@ module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
       }
     }
 
-    import Impl<C>
+    private module Stage1 = ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = Impl<C, Stage1::Stage1NoState>;
+
+    import Flow
   }
 
   /**
@@ -723,7 +730,13 @@ module DataFlowMake<LocationSig Location, InputSig<Location> Lang> {
       }
     }
 
-    import Impl<C>
+    private module Stage1 = ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = Impl<C, Stage1::Stage1WithState>;
+
+    import Flow
   }
 
   signature class PathNodeSig {

shared/dataflow/codeql/dataflow/TaintTracking.qll

@@ -5,6 +5,7 @@
 
 private import DataFlow as DF
 private import internal.DataFlowImpl
+private import internal.DataFlowImplStage1
 private import codeql.util.Location
 
 /**
@@ -47,6 +48,7 @@ module TaintFlowMake<
   private import TaintTrackingLang
   private import DF::DataFlowMake<Location, DataFlowLang> as DataFlow
   private import MakeImpl<Location, DataFlowLang> as DataFlowInternal
+  private import MakeImplStage1<Location, DataFlowLang> as DataFlowInternalStage1
 
   private module AddTaintDefaults<DataFlowInternal::FullStateConfigSig Config> implements
     DataFlowInternal::FullStateConfigSig
@@ -94,7 +96,13 @@ module TaintFlowMake<
       import AddTaintDefaults<Config0>
     }
 
-    import DataFlowInternal::Impl<C>
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1NoState>;
+
+    import Flow
   }
 
   /**
@@ -122,7 +130,13 @@ module TaintFlowMake<
       import AddTaintDefaults<Config0>
     }
 
-    import DataFlowInternal::Impl<C>
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1WithState>;
+
+    import Flow
   }
 
   signature int speculationLimitSig();
@@ -218,7 +232,13 @@ module TaintFlowMake<
       import AddTaintDefaults<AddSpeculativeTaintSteps<Config0, speculationLimit/0>>
     }
 
-    import DataFlowInternal::Impl<C>
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1WithState>;
+
+    import Flow
   }
 
   /**
@@ -250,6 +270,12 @@ module TaintFlowMake<
       import AddTaintDefaults<AddSpeculativeTaintSteps<Config0, speculationLimit/0>>
     }
 
-    import DataFlowInternal::Impl<C>
+    private module Stage1 = DataFlowInternalStage1::ImplStage1<C>;
+
+    import Stage1::PartialFlow
+
+    private module Flow = DataFlowInternal::Impl<C, Stage1::Stage1WithState>;
+
+    import Flow
   }
 }

shared/dataflow/codeql/dataflow/internal/DataFlowImplCommon.qll

@@ -67,7 +67,7 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
     /**
      * Holds if `source` is a relevant data flow source.
      */
-    signature predicate sourceNode(Node source);
+    signature predicate sourceNodeSig(Node source);
 
     /**
      * EXPERIMENTAL: This API is subject to change without notice.
@@ -75,7 +75,7 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
      * Given a source definition, this constructs a simple forward flow
      * computation with an access path limit of 1.
      */
-    module SimpleGlobal<sourceNode/1 source> {
+    module SimpleGlobal<sourceNodeSig/1 source> {
       import TypeTracking::TypeTrack<source/1>
     }
   }
@@ -853,6 +853,10 @@ module MakeImplCommon<LocationSig Location, InputSig<Location> Lang> {
   class SndLevelScopeOption = SndLevelScopeOption::Option;
 
   final class NodeEx extends TNodeEx {
+    NodeEx getNodeEx() { result = this }
+
+    Unit getState() { any() }
+
     string toString() {
       result = this.asNode().toString()
       or