Python: Clean up various jump/local data flow steps

Removes steps from `ModuleVariableNode`s from `essaFlowStep`, and
instead puts them only in `jumpStep`. This cleans up the logic a bit.

This slightly broke the type tracker implementation (as it relied on
`essaFlowStep` being fairly liberal), so I have rewritten it to
explicitly rely on just familiar predicates for local and jump steps.

Additionally, we disallow Essa-to-Essa steps where exactly one of the
two nodes corresponds to a global variable (i.e. only local-local and
global-global steps).

python/ql/src/experimental/dataflow/TypeTracker.qll

@@ -47,11 +47,11 @@ class StepSummary extends TStepSummary {
 module StepSummary {
   cached
   predicate step(Node nodeFrom, Node nodeTo, StepSummary summary) {
-    exists(Node mid | EssaFlow::essaFlowStep*(nodeFrom, mid) and smallstep(mid, nodeTo, summary))
+    exists(Node mid | typePreservingStep*(nodeFrom, mid) and smallstep(mid, nodeTo, summary))
   }
 
   predicate smallstep(Node nodeFrom, Node nodeTo, StepSummary summary) {
-    EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+    typePreservingStep(nodeFrom, nodeTo) and
     summary = LevelStep()
     or
     callStep(nodeFrom, nodeTo) and summary = CallStep()
@@ -68,6 +68,12 @@ module StepSummary {
   }
 }
 
+/** Holds if it's reasonable to expect the data flow step from `nodeFrom` to `nodeTo` to preserve types. */
+private predicate typePreservingStep(Node nodeFrom, Node nodeTo) {
+  EssaFlow::essaFlowStep(nodeFrom, nodeTo) or
+  jumpStep(nodeFrom, nodeTo)
+}
+
 /** Holds if `nodeFrom` steps to `nodeTo` by being passed as a parameter in a call. */
 predicate callStep(ArgumentNode nodeFrom, ParameterNode nodeTo) {
   // TODO: Support special methods?
@@ -111,7 +117,7 @@ predicate returnStep(ReturnNode nodeFrom, Node nodeTo) {
 predicate basicStoreStep(Node nodeFrom, Node nodeTo, string attr) {
   exists(AttributeAssignment a, Node var |
     a.getName() = attr and
-    EssaFlow::essaFlowStep*(nodeTo, var) and
+    simpleLocalFlowStep*(nodeTo, var) and
     var.asVar() = a.getInput() and
     nodeFrom.asCfgNode() = a.getValue()
   )
@@ -276,7 +282,7 @@ class TypeTracker extends TTypeTracker {
       result = this.append(summary)
     )
     or
-    EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+    typePreservingStep(nodeFrom, nodeTo) and
     result = this
   }
 }

python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll

@@ -127,12 +127,6 @@ module EssaFlow {
       nodeTo.(EssaNode).getVar() = p.getVariable() and
       nodeFrom.(EssaNode).getVar() = p.getAnInput()
     )
-    or
-    // Module variable read
-    nodeFrom.(ModuleVariableNode).getARead() = nodeTo
-    or
-    // Module variable write
-    nodeFrom = nodeTo.(ModuleVariableNode).getAWrite()
   }
 
   predicate useToNextUse(NameNode nodeFrom, NameNode nodeTo) {
@@ -153,13 +147,31 @@ module EssaFlow {
  * excludes SSA flow through instance fields.
  */
 predicate simpleLocalFlowStep(Node nodeFrom, Node nodeTo) {
-  not nodeFrom instanceof ModuleVariableNode and
-  not nodeTo instanceof ModuleVariableNode and
   // If there is ESSA-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)
+    nodeFrom = update(node) and
+    (
+      not node instanceof EssaNode or
+      not nodeTo instanceof EssaNode or
+      localEssaStep(node, nodeTo)
+    )
+  )
+}
+
+/**
+ * Holds if there is an Essa flow step from `nodeFrom` to `nodeTo` that does not switch between
+ * local and global SSA variables.
+ */
+private predicate localEssaStep(EssaNode nodeFrom, EssaNode nodeTo) {
+  EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
+  (
+    nodeFrom.getVar() instanceof GlobalSsaVariable and
+    nodeTo.getVar() instanceof GlobalSsaVariable
+    or
+    not nodeFrom.getVar() instanceof GlobalSsaVariable and
+    not nodeTo.getVar() instanceof GlobalSsaVariable
   )
 }
 
@@ -413,12 +425,11 @@ string ppReprType(DataFlowType t) { none() }
  * taken into account.
  */
 predicate jumpStep(Node nodeFrom, Node nodeTo) {
-  EssaFlow::essaFlowStep(nodeFrom, nodeTo) and
-  (
-    nodeFrom instanceof ModuleVariableNode
-    or
-    nodeTo instanceof ModuleVariableNode
-  )
+  // Module variable read
+  nodeFrom.(ModuleVariableNode).getARead() = nodeTo
+  or
+  // Module variable write
+  nodeFrom = nodeTo.(ModuleVariableNode).getAWrite()
 }
 
 //--------

python/ql/test/experimental/dataflow/strange-essaflow/testFlow.expected

@@ -4,4 +4,3 @@ flowstep
 jumpStep
 | test.py:2:8:2:9 | GSSA Variable os | test.py:0:0:0:0 | ModuleVariableNode for Global Variable os in Module test |
 essaFlowStep
-| test.py:2:8:2:9 | GSSA Variable os | test.py:0:0:0:0 | ModuleVariableNode for Global Variable os in Module test |