From 0a6a48c4278f99417fe81f7d3f0194edeaec7e46 Mon Sep 17 00:00:00 2001
From: 97littleleaf11 <97littleleaf11@users.noreply.github.com>
Date: Tue, 8 Jun 2021 19:07:00 +0800
Subject: [PATCH] [mypyc] Add 'range' primitive type (#10307)

Closes mypyc/mypyc#770.
---
 mypyc/codegen/emit.py              | 12 +++---
 mypyc/ir/rtypes.py                 | 10 ++++-
 mypyc/irbuild/mapper.py            |  4 +-
 mypyc/primitives/misc_ops.py       |  6 +++
 mypyc/test-data/fixtures/ir.py     |  7 ++-
 mypyc/test-data/irbuild-basic.test | 68 ++++++++++++++++++++++++++++++
 mypyc/test-data/run-dicts.test     |  2 +-
 mypyc/test-data/run-floats.test    |  2 +
 mypyc/test-data/run-loops.test     | 33 ++++++++++++++-
 9 files changed, 134 insertions(+), 10 deletions(-)

diff --git a/mypyc/codegen/emit.py b/mypyc/codegen/emit.py
index 769888e5ad6e..e1ad024bf9f2 100644
--- a/mypyc/codegen/emit.py
+++ b/mypyc/codegen/emit.py
@@ -15,7 +15,7 @@
     is_list_rprimitive, is_dict_rprimitive, is_set_rprimitive, is_tuple_rprimitive,
     is_none_rprimitive, is_object_rprimitive, object_rprimitive, is_str_rprimitive,
     int_rprimitive, is_optional_type, optional_value_type, is_int32_rprimitive,
-    is_int64_rprimitive, is_bit_rprimitive
+    is_int64_rprimitive, is_bit_rprimitive, is_range_rprimitive
 )
 from mypyc.ir.func_ir import FuncDecl
 from mypyc.ir.class_ir import ClassIR, all_concrete_classes
@@ -410,8 +410,8 @@ def emit_cast(self, src: str, dest: str, typ: RType, declare_dest: bool = False,
 
         # TODO: Verify refcount handling.
         if (is_list_rprimitive(typ) or is_dict_rprimitive(typ) or is_set_rprimitive(typ)
-                or is_float_rprimitive(typ) or is_str_rprimitive(typ) or is_int_rprimitive(typ)
-                or is_bool_rprimitive(typ)):
+                or is_str_rprimitive(typ) or is_range_rprimitive(typ) or is_float_rprimitive(typ)
+                or is_int_rprimitive(typ) or is_bool_rprimitive(typ)):
             if declare_dest:
                 self.emit_line('PyObject *{};'.format(dest))
             if is_list_rprimitive(typ):
@@ -420,10 +420,12 @@ def emit_cast(self, src: str, dest: str, typ: RType, declare_dest: bool = False,
                 prefix = 'PyDict'
             elif is_set_rprimitive(typ):
                 prefix = 'PySet'
-            elif is_float_rprimitive(typ):
-                prefix = 'CPyFloat'
             elif is_str_rprimitive(typ):
                 prefix = 'PyUnicode'
+            elif is_range_rprimitive(typ):
+                prefix = 'PyRange'
+            elif is_float_rprimitive(typ):
+                prefix = 'CPyFloat'
             elif is_int_rprimitive(typ):
                 prefix = 'PyLong'
             elif is_bool_rprimitive(typ) or is_bit_rprimitive(typ):
diff --git a/mypyc/ir/rtypes.py b/mypyc/ir/rtypes.py
index 39dce61c2c76..d0148090b1e3 100644
--- a/mypyc/ir/rtypes.py
+++ b/mypyc/ir/rtypes.py
@@ -326,6 +326,10 @@ def __hash__(self) -> int:
 tuple_rprimitive = RPrimitive('builtins.tuple', is_unboxed=False,
                               is_refcounted=True)  # type: Final
 
+# Python range object.
+range_rprimitive = RPrimitive('builtins.range', is_unboxed=False,
+                              is_refcounted=True)  # type: Final
+
 
 def is_tagged(rtype: RType) -> bool:
     return rtype is int_rprimitive or rtype is short_int_rprimitive
@@ -405,6 +409,10 @@ def is_tuple_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and rtype.name == 'builtins.tuple'
 
 
+def is_range_rprimitive(rtype: RType) -> bool:
+    return isinstance(rtype, RPrimitive) and rtype.name == 'builtins.range'
+
+
 def is_sequence_rprimitive(rtype: RType) -> bool:
     return isinstance(rtype, RPrimitive) and (
         is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype) or is_str_rprimitive(rtype)
@@ -805,5 +813,5 @@ def deserialize(cls, data: JsonDict, ctx: 'DeserMaps') -> 'RArray':
 PyListObject = RStruct(
     name='PyListObject',
     names=['ob_base', 'ob_item', 'allocated'],
-    types=[PyObject, pointer_rprimitive, c_pyssize_t_rprimitive]
+    types=[PyVarObject, pointer_rprimitive, c_pyssize_t_rprimitive]
 )
diff --git a/mypyc/irbuild/mapper.py b/mypyc/irbuild/mapper.py
index 5b2521d2e76c..8e9dd8a9c578 100644
--- a/mypyc/irbuild/mapper.py
+++ b/mypyc/irbuild/mapper.py
@@ -12,7 +12,7 @@
 from mypyc.ir.rtypes import (
     RType, RUnion, RTuple, RInstance, object_rprimitive, dict_rprimitive, tuple_rprimitive,
     none_rprimitive, int_rprimitive, float_rprimitive, str_rprimitive, bool_rprimitive,
-    list_rprimitive, set_rprimitive
+    list_rprimitive, set_rprimitive, range_rprimitive
 )
 from mypyc.ir.func_ir import FuncSignature, FuncDecl, RuntimeArg
 from mypyc.ir.class_ir import ClassIR
@@ -58,6 +58,8 @@ def type_to_rtype(self, typ: Optional[Type]) -> RType:
                 return set_rprimitive
             elif typ.type.fullname == 'builtins.tuple':
                 return tuple_rprimitive  # Varying-length tuple
+            elif typ.type.fullname == 'builtins.range':
+                return range_rprimitive
             elif typ.type in self.type_to_ir:
                 inst = RInstance(self.type_to_ir[typ.type])
                 # Treat protocols as Union[protocol, object], so that we can do fast
diff --git a/mypyc/primitives/misc_ops.py b/mypyc/primitives/misc_ops.py
index 4f3a49c14a9f..3c8f8e4fa9dd 100644
--- a/mypyc/primitives/misc_ops.py
+++ b/mypyc/primitives/misc_ops.py
@@ -15,6 +15,12 @@
     type=object_rprimitive,
     src='PyBool_Type')
 
+# Get the 'range' type object.
+load_address_op(
+    name='builtins.range',
+    type=object_rprimitive,
+    src='PyRange_Type')
+
 # Get the boxed Python 'None' object
 none_object_op = load_address_op(
     name='Py_None',
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index b13deed90f4a..bd98a07949f2 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/mypyc/test-data/fixtures/ir.py
@@ -193,6 +193,12 @@ def __or__(self, s: Set[S]) -> Set[Union[T, S]]: ...
 
 class slice: pass
 
+class range(Iterable[int]):
+    def __init__(self, x: int, y: int = ..., z: int = ...) -> None: pass
+    def __iter__(self) -> Iterator[int]: pass
+    def __len__(self) -> int: pass
+    def __next__(self) -> int: pass
+
 class property:
     def __init__(self, fget: Optional[Callable[[Any], Any]] = ...,
                  fset: Optional[Callable[[Any, Any], None]] = ...,
@@ -245,7 +251,6 @@ def id(o: object) -> int: pass
 # This type is obviously wrong but the test stubs don't have Sized anymore
 def len(o: object) -> int: pass
 def print(*object) -> None: pass
-def range(x: int, y: int = ..., z: int = ...) -> Iterator[int]: pass
 def isinstance(x: object, t: object) -> bool: pass
 def iter(i: Iterable[T]) -> Iterator[T]: pass
 @overload
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
index 40fe1f5aa809..98a86e3f7ee9 100644
--- a/mypyc/test-data/irbuild-basic.test
+++ b/mypyc/test-data/irbuild-basic.test
@@ -3649,3 +3649,71 @@ L0:
     r2 = r1 >= 0 :: signed
     r3 = truncate r1: int32 to builtins.bool
     return r3
+
+[case testRangeObject]
+def range_object() -> None:
+    r = range(4, 12, 2)
+    sum = 0
+    for i in r:
+        sum += i
+
+def range_in_loop() -> None:
+    sum = 0
+    for i in range(4, 12, 2):
+        sum += i
+[out]
+def range_object():
+    r0, r1, r2, r3, r4 :: object
+    r5, r :: range
+    sum :: int
+    r6, r7 :: object
+    r8, i, r9 :: int
+    r10 :: bit
+L0:
+    r0 = load_address PyRange_Type
+    r1 = box(short_int, 8)
+    r2 = box(short_int, 24)
+    r3 = box(short_int, 4)
+    r4 = PyObject_CallFunctionObjArgs(r0, r1, r2, r3, 0)
+    r5 = cast(range, r4)
+    r = r5
+    sum = 0
+    r6 = PyObject_GetIter(r)
+L1:
+    r7 = PyIter_Next(r6)
+    if is_error(r7) goto L4 else goto L2
+L2:
+    r8 = unbox(int, r7)
+    i = r8
+    r9 = CPyTagged_Add(sum, i)
+    sum = r9
+L3:
+    goto L1
+L4:
+    r10 = CPy_NoErrOccured()
+L5:
+    return 1
+def range_in_loop():
+    sum :: int
+    r0 :: short_int
+    i :: int
+    r1 :: bit
+    r2 :: int
+    r3 :: short_int
+L0:
+    sum = 0
+    r0 = 8
+    i = r0
+L1:
+    r1 = r0 < 24 :: signed
+    if r1 goto L2 else goto L4 :: bool
+L2:
+    r2 = CPyTagged_Add(sum, i)
+    sum = r2
+L3:
+    r3 = r0 + 4
+    r0 = r3
+    i = r3
+    goto L1
+L4:
+    return 1
diff --git a/mypyc/test-data/run-dicts.test b/mypyc/test-data/run-dicts.test
index 49deb0854483..b25b5237e80f 100644
--- a/mypyc/test-data/run-dicts.test
+++ b/mypyc/test-data/run-dicts.test
@@ -1,4 +1,4 @@
-# Dict test cases (compile and run)
+# Test cases for dicts (compile and run)
 
 [case testDictStuff]
 from typing import Dict, Any, List, Set, Tuple
diff --git a/mypyc/test-data/run-floats.test b/mypyc/test-data/run-floats.test
index e716949d69c4..1b40a598bf39 100644
--- a/mypyc/test-data/run-floats.test
+++ b/mypyc/test-data/run-floats.test
@@ -1,3 +1,5 @@
+# Test cases for floats (compile and run)
+
 [case testStrToFloat]
 def str_to_float(x: str) -> float:
     return float(x)
diff --git a/mypyc/test-data/run-loops.test b/mypyc/test-data/run-loops.test
index b83853bc6d16..994b30b42347 100644
--- a/mypyc/test-data/run-loops.test
+++ b/mypyc/test-data/run-loops.test
@@ -1,4 +1,4 @@
-# Test cases for "for" and "while" loops (compile and run)
+# Test cases for "range" objects, "for" and "while" loops (compile and run)
 
 [case testFor]
 from typing import List, Tuple
@@ -452,3 +452,34 @@ def bar(x: Optional[str]) -> None:
 [file driver.py]
 from native import bar
 bar(None)
+
+[case testRangeObject]
+from typing import Any
+
+def f(x: range) -> int:
+    sum = 0
+    for i in x:
+        sum += i
+    return sum
+
+def test_range_object() -> None:
+    r1 = range(4, 12, 2)
+    tmp_list = [x for x in r1]
+    assert tmp_list == [4, 6, 8, 10]
+    assert f(r1) == 28
+    r2: Any = range(10)
+    assert f(r2) == 45
+    r3: Any = 'x'
+    try:
+        f(r3)
+    except TypeError as e:
+        assert "range object expected; got str" in str(e)
+    try:
+        ff: Any = f
+        ff(r3)
+    except TypeError as e:
+        assert "range object expected; got str" in str(e)
+    try:
+        r4 = range(4, 12, 0)
+    except ValueError as e:
+        assert "range() arg 3 must not be zero" in str(e)