Supporting range() with one argument

marcoeilers · Feb 16, 2024 · 760e8c3 · 760e8c3
1 parent f6a369f
commit 760e8c3
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 7 deletions.
diff --git a/src/nagini_translation/resources/preamble.index b/src/nagini_translation/resources/preamble.index
@@ -804,7 +804,7 @@
         "__getitem__": {
             "args": ["range", "__prim__int"],
             "type": "__prim__int",
-            "requires": ["__val__"]
+            "requires": ["__val__", "__len__"]
         },
         "__getitem_slice__": {
             "args": ["range", "slice"],

diff --git a/src/nagini_translation/translators/call.py b/src/nagini_translation/translators/call.py
@@ -416,13 +416,22 @@ def _translate_set(self, node: ast.Call, ctx: Context) -> StmtsAndExpr:
         return stmts, result_var
 
     def _translate_range(self, node: ast.Call, ctx: Context) -> StmtsAndExpr:
-        if len(node.args) != 2:
-            msg = 'range() is currently only supported with two args.'
+        if node.keywords:
+            msg = 'range() with keyword args is currently not supported.'
+            raise UnsupportedException(node, msg)
+        if len(node.args) == 2:
+            start_arg = node.args[0]
+            end_arg = node.args[1]
+        elif len(node.args) == 1:
+            start_arg = None
+            end_arg = node.args[0]
+        else:
+            msg = 'range() step is currently not supported.'
             raise UnsupportedException(node, msg)
         range_class = ctx.module.global_module.classes[RANGE_TYPE]
-        start_stmt, start = self.translate_expr(node.args[0], ctx,
-                                                self.viper.Int)
-        end_stmt, end = self.translate_expr(node.args[1], ctx, self.viper.Int)
+        start_stmt, start = (self.translate_expr(start_arg, ctx, self.viper.Int) if start_arg
+                             else ([], self.viper.IntLit(0, self.to_position(node, ctx), self.no_info(ctx))))
+        end_stmt, end = self.translate_expr(end_arg, ctx, self.viper.Int)
         # Add unique integer to make new instance different from other ranges.
         args = [start, end, self.get_fresh_int_lit(ctx)]
         arg_types = [None, None, None]

diff --git a/tests/functional/verification/test_range.py b/tests/functional/verification/test_range.py
@@ -16,10 +16,22 @@ def rangetest() -> None:
 
 
 def rangetest2() -> None:
-    a = range(0, 5)
+    a = range(5)
     for i in a:
         assert i < 6
     assert i == 4
     for b in a:
         #:: ExpectedOutput(assert.failed:assertion.false)
         assert b > 2
+
+
+def rangetest3() -> None:
+    a = range(1, 5)
+    Assert(a[2] == 3)
+    Assert(3 in a)
+    Assert(7 not in a)
+    Assert(5 not in a)
+    Assert(Forall(a, lambda x: (x < 5, [])))
+    Assert(Forall(a, lambda x: (x > 0, [])))
+    #:: ExpectedOutput(assert.failed:assertion.false)
+    Assert(Forall(a, lambda x: (x > 1, [])))