Swift: move keypath dataflow writes to fix types

Author: rdmarsh2

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

@@ -263,11 +263,18 @@ private module Cached {
     or
     // Flow from the unique parameter of a key path expression to
     // the first component in the chain.
-    nodeTo.(KeyPathComponentNodeImpl).getComponent() =
-      nodeFrom.(KeyPathParameterNode).getComponent(0)
+    exists(KeyPathExpr keyPath |
+      nodeFrom.(KeyPathComponentNodeImpl).getComponent() = keyPath.getComponent(keyPath.getNumberOfComponents() - 1) and
+      nodeTo.(KeyPathReturnNodeImpl).getKeyPathExpr() = keyPath
+    )
     or
-    nodeFrom.(KeyPathComponentPostUpdateNode).getComponent() =
-      nodeTo.(KeyPathParameterPostUpdateNode).getComponent(0)
+    exists(KeyPathExpr keyPath |
+      nodeTo.(KeyPathComponentPostUpdateNode).getComponent() = keyPath.getComponent(keyPath.getNumberOfComponents()-1) and
+      nodeFrom.(KeyPathReturnPostUpdateNode).getKeyPathExpr() = keyPath
+    )
+    
+  //  nodeFrom.(KeyPathComponentPostUpdateNode).getComponent() =
+  //    nodeTo.(KeyPathParameterPostUpdateNode).getComponent(0)
     or
     // Flow to the result of a keypath assignment
     exists(KeyPathApplicationExpr apply, AssignExpr assign |
@@ -864,8 +871,8 @@ predicate storeStep(Node node1, ContentSet c, Node node2) {
   // creation of an optional via implicit wrapping keypath component
   exists(KeyPathComponent component |
     component.isOptionalWrapping() and
-    node1.(KeyPathComponentNodeImpl).getComponent() = component and
-    node2.(KeyPathReturnNodeImpl).getKeyPathExpr() = component.getKeyPathExpr() and
+    node1.(KeyPathComponentNodeImpl).getComponent().getNextComponent() = component and
+    node2.(KeyPathComponentNodeImpl).getComponent() = component and
     c instanceof OptionalSomeContentSet
   )
   or
@@ -951,7 +958,14 @@ predicate readStep(Node node1, ContentSet c, Node node2) {
   or
   // read of a component in a key-path expression chain
   exists(KeyPathComponent component |
-    component = node1.(KeyPathComponentNodeImpl).getComponent() and
+    // the first node is either the previous element in the chain
+    node1.(KeyPathComponentNodeImpl).getComponent().getNextComponent() = component
+    or
+    // or the start node, if this is the first component in the chain 
+    component =
+      node1.(KeyPathParameterNode).getComponent(0)
+  |
+    component = node2.(KeyPathComponentNodeImpl).getComponent() and
     (
       c.isSingleton(any(Content::FieldContent ct | ct.getField() = component.getDeclRef()))
       or
@@ -965,13 +979,6 @@ predicate readStep(Node node1, ContentSet c, Node node2) {
         component.isOptionalChaining()
       )
     )
-  |
-    // the next node is either the next element in the chain
-    node2.(KeyPathComponentNodeImpl).getComponent() = component.getNextComponent()
-    or
-    // or the return node, if this is the last component in the chain
-    not exists(component.getNextComponent()) and
-    node2.(KeyPathReturnNodeImpl).getKeyPathExpr() = component.getKeyPathExpr()
   )
   or
   // read of array or collection content via subscript operator