Swift: use PostUpdateNodeImpl for more postupdate nodes

Author: Robert Marsh

swift/ql/lib/codeql/swift/dataflow/internal/DataFlowPrivate.qll

@@ -57,7 +57,7 @@ private class KeyPathComponentNodeImpl extends TKeyPathComponentNode, NodeImpl {
   KeyPathComponent getComponent() { result = component }
 }
 
-private class KeyPathComponentPostUpdateNode extends TKeyPathComponentPostUpdateNode, NodeImpl {
+private class KeyPathComponentPostUpdateNode extends TKeyPathComponentPostUpdateNode, NodeImpl, PostUpdateNodeImpl {
   KeyPathComponent component;
 
   KeyPathComponentPostUpdateNode() { this = TKeyPathComponentPostUpdateNode(component) }
@@ -70,6 +70,10 @@ private class KeyPathComponentPostUpdateNode extends TKeyPathComponentPostUpdate
     result.asSourceCallable() = component.getKeyPathExpr()
   }
 
+  override KeyPathComponentNodeImpl getPreUpdateNode() {
+    result.getComponent() = this.getComponent()
+  }
+
   KeyPathComponent getComponent() { result = component }
 }
 
@@ -448,7 +452,7 @@ class KeyPathReturnPostUpdateNode extends NodeImpl, ParameterNodeImpl, PostUpdat
 
   KeyPathReturnPostUpdateNode() { this = TKeyPathReturnPostUpdateNode(exit) }
 
-  override KeyPathParameterNode getPreUpdateNode() {
+  override KeyPathReturnNodeImpl getPreUpdateNode() {
     result.getKeyPathExpr() = this.getKeyPathExpr()
   }
 
@@ -831,24 +835,6 @@ predicate storeStep(Node node1, ContentSet c, Node node2) {
     c instanceof OptionalSomeContentSet
   )
   or
-  // store through a component in a key-path expression chain
-  exists(KeyPathComponent component |
-    component = node2.(KeyPathComponentPostUpdateNode).getComponent() and
-    (
-      c.isSingleton(any(Content::FieldContent ct | ct.getField() = component.getDeclRef()))
-      or
-      c.isSingleton(any(Content::ArrayContent ac)) and
-      component.isSubscript()
-    )
-  |
-    // the previous node is either the next element in the chain
-    node1.(KeyPathComponentPostUpdateNode).getComponent() = component.getNextComponent()
-    or
-    // or the return node, if this is the last component in the chain
-    not exists(component.getNextComponent()) and
-    node1.(KeyPathReturnPostUpdateNode).getKeyPathExpr() = component.getKeyPathExpr()
-  )
-  or
   FlowSummaryImpl::Private::Steps::summaryStoreStep(node1.(FlowSummaryNode).getSummaryNode(), c,
     node2.(FlowSummaryNode).getSummaryNode())
 }

swift/ql/test/library-tests/dataflow/dataflow/DataFlow.expected

@@ -1001,6 +1001,7 @@ subpaths
 | test.swift:765:29:765:29 | s [some:0, x] | test.swift:655:8:655:12 | s [some:0, x] | test.swift:655:3:657:3 | self[return] [s, some:0, x] | test.swift:765:14:765:30 | call to S2_Optional.init(s:) [s, some:0, x] |
 | test.swift:767:15:767:15 | s2 [s, some:0, x] | test.swift:766:13:766:29 | enter #keyPath(...) [s, some:0, x] | test.swift:766:13:766:29 | exit #keyPath(...) | test.swift:767:15:767:28 | \\...[...] |
 | test.swift:774:3:774:16 | \\...[...] | test.swift:773:11:773:17 | [post] exit #keyPath(...) | test.swift:773:11:773:17 | [post] enter #keyPath(...) [s, x] | test.swift:774:3:774:3 | [post] s2 [s, x] |
+| test.swift:774:3:774:16 | \\...[...] | test.swift:773:11:773:17 | [post] exit #keyPath(...) | test.swift:773:15:773:15 | [post] KeyPathComponent [s, x] | test.swift:774:3:774:3 | [post] s2 [s, x] |
 | test.swift:775:13:775:13 | s2 [s, x] | test.swift:632:7:632:7 | self [s, x] | file://:0:0:0:0 | .s [x] | test.swift:775:13:775:16 | .s [x] |
 | test.swift:775:13:775:16 | .s [x] | test.swift:615:7:615:7 | self [x] | file://:0:0:0:0 | .x | test.swift:775:13:775:18 | .x |
 #select