Swift: add expr types to for-in upgrade script

swift/ql/lib/upgrades/147e087e57e51b2eb41e75c9c97380d0e6c20ecb/for_in_upgrade.ql

@@ -136,6 +136,17 @@ Element getParent(Element type) {
   )
 }
 
+Element getNextMethod(ForEachStmt foreach) {
+  exists(@element sequence, @element exprType, @element parentType, @element typeDecl |
+    sequence = foreach.getSequence() and
+    expr_types(sequence, exprType) and
+    parentType = getParent*(exprType) and
+    any_generic_types(parentType, typeDecl) and
+    decl_members(typeDecl, _, result) and
+    callable_names(result, "next()")
+  )
+}
+
 // TODO: do we need a new_apply_expr_arguments
 query predicate new_decl_ref_exprs(NewElement id, NewElement decl) {
   decl_ref_exprs(id, decl)
@@ -145,18 +156,9 @@ query predicate new_decl_ref_exprs(NewElement id, NewElement decl) {
     Fresh::map(TIteratorVarConcreteDecl(foreach)) = decl
   )
   or
-  exists(
-    ForEachStmt foreach, @element sequence, @element exprType, @element parentType,
-    @element typeDecl
-  |
+  exists(ForEachStmt foreach |
     Fresh::map(TNextCallFuncRef(foreach)) = id and
-    // ForEachStmt.getSequence().getType().getMember(next)
-    sequence = foreach.getSequence() and
-    expr_types(sequence, exprType) and
-    parentType = getParent*(exprType) and
-    any_generic_types(parentType, typeDecl) and
-    decl_members(typeDecl, _, decl) and
-    callable_names(decl, "next()")
+    decl = getNextMethod(foreach)
   )
 }
 
@@ -217,4 +219,25 @@ query predicate new_expr_types(NewElement expr, NewElement type) {
     var_decl_parent_patterns(var_decl, pattern) and
     var_decls(var_decl, _, type)
   )
+  or
+  exists(ForEachStmt stmt |
+    expr = Fresh::map(TNextCallMethodLookup(stmt)) and
+    value_decls(getNextMethod(stmt), type)
+  )
+  or
+  exists(ForEachStmt stmt |
+    expr = Fresh::map(TNextCallVarRef(stmt)) and
+    expr_types(stmt.getSequence(), type)
+  )
+  or
+  exists(ForEachStmt stmt |
+    expr = Fresh::map(TNextCallFuncRef(stmt)) and
+    value_decls(getNextMethod(stmt), type)
+  )
+  or
+  exists(ForEachStmt stmt, NewElement plainType|
+    expr = Fresh::map(TNextCallInOutConversion(stmt)) and
+    expr_types(stmt.getSequence(), plainType) and
+    in_out_types(type, plainType)
+  )
 }

swift/ql/lib/upgrades/147e087e57e51b2eb41e75c9c97380d0e6c20ecb/upgrade.properties

@@ -16,4 +16,5 @@ in_out_exprs.rel: run for_in_upgrade.qlo new_in_out_exprs
 call_exprs.rel: run for_in_upgrade.qlo new_call_exprs
 locatable_locations.rel: run for_in_upgrade.qlo new_locatable_locations
 var_decls.rel: run for_in_upgrade.qlo new_var_decls
-concrete_var_decls.rel: run for_in_upgrade.qlo new_concrete_var_decls
+concrete_var_decls.rel: run for_in_upgrade.qlo new_concrete_var_decls
+expr_types.rel: run for_in_upgrade.qlo new_expr_types