Implementing zip lookaside for interpreter (#1259)

Co-authored-by: Rany Kamel <[email protected]>

thunder/core/interpreter.py

@@ -1413,6 +1413,47 @@ def impl(obj, start):
     return _interpret_call(impl, obj, wrap_const(start))
 
 
+def _zip_lookaside(*obj: Iterable, strict=False):
+
+    if not obj:
+        return
+
+    def zip(*obj, strict=False):
+        # zip('ABCD', 'xy') --> Ax By
+        sentinel = object()
+        iterators = [iter(it) for it in obj]
+        while iterators:
+            result = []
+            break_loop = False
+            for it in iterators:
+                elem = next(it, sentinel)
+                if elem is sentinel:
+                    if not strict:
+                        return
+                    else:
+                        break_loop = True
+                        break
+                result.append(elem)
+
+            if break_loop:
+                break
+
+            yield tuple(result)
+        if result:
+            i = len(result)
+            plural = " " if i == 1 else "s 1-"
+            msg = f"zip() argument {i+1} is shorter than argument{plural}{i}"
+            raise ValueError(msg)
+        sentinel = object()
+        for i, iterator in enumerate(iterators[1:], 1):
+            if next(iterator, sentinel) is not sentinel:
+                plural = " " if i == 1 else "s 1-"
+                msg = f"zip() argument {i+1} is longer than argument{plural}{i}"
+                raise ValueError(msg)
+
+    return _interpret_call(zip, *obj, strict=wrap_const(strict))
+
+
 @interpreter_needs_wrap
 def eval_lookaside(
     source: str | bytes | bytearray | CodeType,  # A python expression
@@ -2743,6 +2784,7 @@ def _type_call_lookaside(wrapped_typ, *args, **kwargs):
     any: _any_lookaside,
     bool: _bool_lookaside,
     enumerate: _enumerate_lookaside,
+    zip: _zip_lookaside,
     exec: exec_lookaside,
     eval: eval_lookaside,
     getattr: _getattr_lookaside,

thunder/tests/test_interpreter.py

@@ -1691,6 +1691,42 @@ def fn():
     jit(fn)()
 
 
+def test_zip_lookaside(jit):
+    import re
+
+    jitting = False
+
+    def foo(*a, strict=False):
+        return list(zip(*a, strict=strict))
+
+    jfoo = jit(foo)
+    jitting = False
+
+    res1 = foo([1, 2, 3], [4, 5, 6])
+    res2 = foo([1, 2, 3], [4, 5, 6], [7, 8, 9])
+    res3 = foo([1, 2], [4, 5, 6])
+    res4 = foo("abc", "xyz")
+    # , match="zip() argument 2 is longer than argument 1"
+
+    with pytest.raises(ValueError, match=re.escape("zip() argument 2 is longer than argument 1")):
+        res5 = foo([1, 2], [4, 5, 6], strict=True)
+
+    jitting = True
+    jres1 = jfoo([1, 2, 3], [4, 5, 6])
+    jres2 = jfoo([1, 2, 3], [4, 5, 6], [7, 8, 9])
+    jres3 = jfoo([1, 2], [4, 5, 6])
+    jres4 = jfoo("abc", "xyz")
+
+    # , match=" zip() argument 2 is longer than argument 1"
+    with pytest.raises(ValueError, match=re.escape("zip() argument 2 is longer than argument 1")):
+        jres5 = jfoo([1, 2], [4, 5, 6], strict=True)
+
+    assert res1 == jres1
+    assert res2 == jres2
+    assert res3 == jres3
+    assert res4 == jres4
+
+
 def test_enumerate_lookaside(jit):
     jitting = False