Repository: python/mypy
Pull request: [mypyc] Add 'range' primitive type
### Description

Closes https://github.com/mypyc/mypyc/issues/770

## Test Plan

```py
def range_object() -> None:
    r = range(4, 12, 2)
    sum = 0
    for i in r:
        sum += I
```
This should generate a range object and a iterator when looping on it.

```py
def range_in_loop() -> None:
    sum = 0
    for i in range(4, 12, 2):
        sum += I
```
This should maintain original specialized IR result.


Problem statement:
Add primitive type for range
Since we optimize some code involving `range(...)`, it makes sense to make `range` a primitive type. Unfortunately there is very little C API support for `range` objects, but at least we could look up the type object efficiently and support runtime type checking.

We already optimize for loops over ranges. The only other primitives that seem worthwhile are these:
* Construct range object: `r = range(...)` (1-3 integer arguments)
* Construct a list from range: `list(range(...))`

Hints from issue discussion:
I tried to add a branch for `range_rprimitive` under `for_helpers/make_for_loop_generator`.  In that case, it would generate more specialized codes rather than generic codes based on iterator.

However, it seems that Python C API doesn't offer `rangeobject`, thus I cannot get its members by using `Loadmem`. 
The only thing we have seems to be the type object, which is available as `PyRange_Type`. Apparently `PyRange_New` was removed from the C API. This seems enough for minimal support -- we can at least perform type checks. Also construction operations like `range(5)` would be a bit faster, since we wouldn't need to look up `range` from builtins.

Though we could take a peek at the implementation and declare the equivalent struct ourselves, at least if it seems stable across Python versions:
```c
typedef struct {
    PyObject_HEAD
    PyObject *start;
    PyObject *stop;
    PyObject *step;
    PyObject *length;
} rangeobject;
```

This way we could perhaps construct an instance more quickly using the C API, but clearly it won't be very portable across Python implementations (but we don't really care about that yet).

Commit message: Load range type and add some tests

Patch to review:
diff --git a/mypyc/codegen/emit.py b/mypyc/codegen/emit.py
index e8b3c77be9d6..f51f56b523ad 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 2490da6fe51f..de8e7d6b60f2 100644
--- a/mypyc/ir/rtypes.py
+++ b/mypyc/ir/rtypes.py
@@ -324,6 +324,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
@@ -401,6 +405,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)
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/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
index 9b7eaeb8e58b..7fab79027af4 100644
--- a/mypyc/test-data/irbuild-basic.test
+++ b/mypyc/test-data/irbuild-basic.test
@@ -3646,3 +3646,69 @@ 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, r :: object
+    sum :: int
+    r5, r6 :: object
+    r7, i, r8 :: int
+    r9 :: 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)
+    r = r4
+    sum = 0
+    r5 = PyObject_GetIter(r)
+L1:
+    r6 = PyIter_Next(r5)
+    if is_error(r6) goto L4 else goto L2
+L2:
+    r7 = unbox(int, r6)
+    i = r7
+    r8 = CPyTagged_Add(sum, i)
+    sum = r8
+L3:
+    goto L1
+L4:
+    r9 = 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 329d186874f4..86e098f4529b 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..468f37a4148a 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,19 @@ def bar(x: Optional[str]) -> None:
 [file driver.py]
 from native import bar
 bar(None)
+
+[case testRangeObject]
+r1 = range(4, 12, 2)
+for i in r1:
+    print(i)
+try:
+    r2 = range(4, 12, 0)
+except ValueError as e:
+    assert 'range() arg 3 must not be zero' in str(e)
+else:
+    assert False
+[out]
+4
+6
+8
+10


Code review comments:
diff_hunk: @@ -452,3 +452,19 @@ def bar(x: Optional[str]) -> None:
 [file driver.py]
 from native import bar
 bar(None)
+
+[case testRangeObject]
+r1 = range(4, 12, 2)
path: mypyc/test-data/run-loops.test
text: @user1:
Can you change this to use a `test_` function without a custom driver, as that's the "modern" way to write tests.

Test also that coercion to `range` works, for example something like this:

```
def f(r: range) -> None:
    pass

r: Any = range(...)
f(r)  # OK

r = 'x'
f(r)  # Should be a TypeError

ff: Any = f
ff(r)  # Again should be a TypeError; this uses the wrapper function since the callee has type Any
```

@author:
I got this error after adding function `def f(r: range) -> None:`
```
run-loops.test:460: error: Function "builtins.range" is not valid as a type
run-loops.test:460: note: Perhaps you need "Callable[...]" or a callback protocol?
```

@user1:
You'll need to change the definition of `range` into a class in `mypyc/test-data/fixtures/ir.py`.