Reduce generality of map and zip for performance

Author: joefarebrother

python/ql/lib/semmle/python/frameworks/Stdlib.qll

@@ -4540,8 +4540,11 @@ module StdlibPrivate {
           or
           input = "Argument[" + (i + 1).toString() + "].SetElement"
           or
+          // We reduce generality slightly by not tracking tuple contents on list arguments beyond the first, for performance.
+          // TODO: Once we have TupleElementAny, this generality can be increased.
+          i = 0 and
           exists(DataFlow::TupleElementContent tc, int j | j = tc.getIndex() |
-            input = "Argument[" + (i + 1).toString() + "].TupleElement[" + j.toString() + "]"
+            input = "Argument[1].TupleElement[" + j.toString() + "]"
           )
           // TODO: Once we have DictKeyContent, we need to transform that into ListElementContent
         ) and
@@ -4624,6 +4627,9 @@ module StdlibPrivate {
           or
           input = "Argument[" + i.toString() + "].SetElement"
           or
+          // We reduce generality slightly by not tracking tuple contents on arguments beyond the first two, for performance.
+          // TODO: Once we have TupleElementAny, this generality can be increased.
+          i in [0 .. 1] and
           exists(DataFlow::TupleElementContent tc, int j | j = tc.getIndex() |
             input = "Argument[" + i.toString() + "].TupleElement[" + j.toString() + "]"
           )

python/ql/test/library-tests/dataflow/coverage/test_builtins.py

@@ -414,6 +414,7 @@ def f(p1,p2):
     SINK(rl[0][0]) #$ flow="SOURCE, l:-10 -> rl[0][0]"
     SINK_F(rl[0][1])
 
+
 @expects(4)
 def test_map_dict():
     d1 = {SOURCE: "v1"}
@@ -429,6 +430,35 @@ def f(p1,p2):
     SINK(rl[0][0]) #$ MISSING: flow="SOURCE, l:-10 -> rl[0][0]"
     SINK_F(rl[0][1])
 
+@expects(6)
+def test_map_multi_list():
+    l1 = [SOURCE]
+    l2 = [SOURCE]
+
+    def f(p1,p2):
+        SINK(p1)  #$ flow="SOURCE, l:-4 -> p1"
+        SINK(p2)  #$ flow="SOURCE, l:-4 -> p2"
+
+        return p1,p2
+
+    rl = list(map(f, l1, l2))
+    SINK(rl[0][0]) #$ flow="SOURCE, l:-10 -> rl[0][0]"
+    SINK(rl[0][1]) #$ flow="SOURCE, l:-10 -> rl[0][1]"
+
+@expects(6)
+def test_map_multi_tuple():
+    l1 = (SOURCE,)
+    l2 = (SOURCE,)
+
+    def f(p1,p2):
+        SINK(p1)  #$ flow="SOURCE, l:-4 -> p1"
+        SINK(p2)  #$ MISSING: flow="SOURCE, l:-4 -> p2" # Tuples are not tracked beyond the first list argument for performance.
+        return p1,p2
+
+    rl = list(map(f, l1, l2))
+    SINK(rl[0][0]) #$ flow="SOURCE, l:-10 -> rl[0][0]"
+    SINK(rl[0][1]) #$ MISSING: flow="SOURCE, l:-10 -> rl[0][1]"
+
 ### filter 
 
 @expects(2)
@@ -474,3 +504,11 @@ def f(p):
 
     rl = list(filter(f,d))
     SINK(rl[0]) #$ MISSING: flow="SOURCE, l:-7 -> rl[0]"
+
+def test_enumerate_list():
+    # TODO
+    pass
+
+def test_zip_list():
+    # TODO
+    pass