Python: rewrite simpleLocalFlowStep

to take into account the split between
import time and runtime.

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll

@@ -152,6 +152,7 @@ class DataFlowExpr = Expr;
  * Flow comes from definitions, uses and refinements.
  */
 // TODO: Consider constraining `nodeFrom` and `nodeTo` to be in the same scope.
+// If they have different enclosing callables, we get consistency errors.
 module EssaFlow {
   predicate essaFlowStep(Node nodeFrom, Node nodeTo) {
     // Definition
@@ -225,41 +226,64 @@ module EssaFlow {
 //--------
 /**
  * This is the local flow predicate that is used as a building block in global
- * data flow. It is a strict subset of the `localFlowStep` predicate, as it
- * excludes SSA flow through instance fields.
+ * data flow.
+ *
+ * Local flow can happen either at import time, when the module is initialised
+ * or at runtime when callables in the module are called.
  */
 predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
-  // If there is ESSA-flow out of a node `node`, we want flow
+  // If there is local flow out of a node `node`, we want flow
   // both out of `node` and any post-update node of `node`.
   exists(Node node |
-    EssaFlow::essaFlowStep(node, nodeTo) and
     nodeFrom = update(node) and
     (
-      not node instanceof EssaNode or
-      not nodeTo instanceof EssaNode or
-      localEssaStep(node, nodeTo)
+      importTimeLocalFlowStep(node, nodeTo) or
+      runtimeLocalFlowStep(node, nodeTo)
     )
   )
 }
 
 /**
- * Holds if there is an Essa flow step from `nodeFrom` to `nodeTo` that does not switch between
- * local and global SSA variables.
- *
- * This predicate is currently empty, since `EssaFlow::essaFlowStep` never goes between `EssaNode`s.
- * (It only starts in an `EssaNode` in a single case, namely `defToFirstUse` which ends in a `CfgNode`.)
+ * Holds if `node` is found at the top level of a module.
  */
-private predicate localEssaStep(EssaNode nodeFrom, EssaNode nodeTo) {
-  EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+pragma[inline]
+predicate isTopLevel(Node node) { node.getScope() instanceof Module }
+
+/** Holds if there is local flow from `nodeFrom` to `nodeTo` at import time. */
+predicate importTimeLocalFlowStep(Node nodeFrom, Node nodeTo) {
+  // As a proxy for whether statements can be executed at import time,
+  // we check if they appear at the top level.
+  // This will miss statements inside functions called from the top level.
+  isTopLevel(nodeFrom) and
+  isTopLevel(nodeTo) and
   (
-    nodeFrom.getVar() instanceof GlobalSsaVariable and
-    nodeTo.getVar() instanceof GlobalSsaVariable
+    EssaFlow::essaFlowStep(nodeFrom, nodeTo)
     or
-    not nodeFrom.getVar() instanceof GlobalSsaVariable and
-    not nodeTo.getVar() instanceof GlobalSsaVariable
+    exists(SsaVariable def |
+      def = any(SsaVariable var).getAnUltimateDefinition() and
+      def.getDefinition() = nodeFrom.asCfgNode() and
+      def.getVariable() = nodeTo.(ModuleVariableNode).getVariable()
+    )
   )
 }
 
+/** Holds if there is local flow from `nodeFrom` to `nodeTo` at runtime. */
+predicate runtimeLocalFlowStep(Node nodeFrom, Node nodeTo) {
+  // Anything not at the top level can be executed at runtime.
+  not isTopLevel(nodeFrom) and
+  not isTopLevel(nodeTo) and
+  EssaFlow::essaFlowStep(nodeFrom, nodeTo)
+}
+
+/** `ModuleVariable`s are accessed via jump steps at runtime. */
+predicate runtimeJumpStep(Node nodeFrom, Node nodeTo) {
+  // Module variable read
+  nodeFrom.(ModuleVariableNode).getARead() = nodeTo
+  or
+  // Module variable write
+  nodeFrom = nodeTo.(ModuleVariableNode).getAWrite()
+}
+
 /**
  * Holds if `result` is either `node`, or the post-update node for `node`.
  */
@@ -860,11 +884,7 @@ string ppReprType(DataFlowType t) { none() }
  * taken into account.
  */
 predicate jumpStep(Node nodeFrom, Node nodeTo) {
-  // Module variable read
-  nodeFrom.(ModuleVariableNode).getARead() = nodeTo
-  or
-  // Module variable write
-  nodeFrom = nodeTo.(ModuleVariableNode).getAWrite()
+  runtimeJumpStep(nodeFrom, nodeTo)
   or
   // Read of module attribute:
   exists(AttrRead r, ModuleValue mv |