For loop named iter (#65)

This showed up in one of the files in `rust_verify/example`. See the
[for-loop PR](verus-lang/verus#954) for a few
details on the syntax.

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Unreleased
 
 * Support automatic formatting of the `calc!` macro
+* Support for named iterators in for loops
 
 # v0.3.1
 

examples/wip.rs

@@ -1,49 +1,32 @@
 verus! {
 
 fn test() {
-    let mut i: i8 = 0;
-    loop
+    for x in iter: 0..end
         invariant
-            i <= 9,
-        invariant_except_break
-            i <= 9,
-        invariant
-            0 <= i <= 10,
-        invariant_ensures
-            0 <= i <= 10,
-        ensures
-            1 <= i,
+            end == 10,
     {
-        assert(i <= 9);
-        i = i + 1;
-        if i == 10 {
-            break ;
-        }
+        n += 3;
     }
-    assert(1 <= i <= 10);
-}
-
-fn test() {
-    let mut i: i8 = 0;
-    while true
+    let x = 2;
+    for x in iter: vec_iter_copy(v)
         invariant
-            i <= 9,
-        invariant_except_break
-            i <= 9,
+            b <==> (forall|i: int| 0 <= i < iter.cur ==> v[i] > 0),
+    {
+        b = b && x > 0;
+    }
+    let y = 3;
+    for x in iter: 0..({
+        let z = end;
+        non_spec();
+        z
+    })
         invariant
-            0 <= i <= 10,
-        invariant_ensures
-            0 <= i <= 10,
-        ensures
-            1 <= i,
+            n == iter.cur * 3,
+            end == 10,
     {
-        assert(i <= 9);
-        i = i + 1;
-        if i == 10 {
-            break ;
-        }
+        n += 3;
+        end = end + 0;  // causes end to be non-constant
     }
-    assert(1 <= i <= 10);
 }
 
 } // verus!

src/lib.rs

@@ -459,18 +459,24 @@ fn loop_to_doc<'a>(
     // We need to add a newline after the very last clause,
     // so that the opening brace of the loop body is on a fresh line
     let mut last_clause_span = None;
+    // We also need to supress the space after the `for x in y`
+    // if we see a colon indicating `for x in y: e`
+    let mut colon_expr = false;
     pair.clone().into_inner().for_each(|p| {
         if p.as_rule() == Rule::loop_clause {
             last_clause_span = Some(p.as_span())
         }
+        if p.as_rule() == Rule::colon_str {
+            colon_expr = true;
+        }
     });
     arena.concat(pair.into_inner().map(|p| {
         if p.as_rule() == Rule::condition {
             to_doc(ctx, p, arena).append(arena.space())
         } else if p.as_rule() == Rule::in_str {
             // Used for for-loops
             arena.space().append(to_doc(ctx, p, arena))
-        } else if p.as_rule() == Rule::expr_no_struct {
+        } else if p.as_rule() == Rule::expr_no_struct && !colon_expr {
             to_doc(ctx, p, arena).append(arena.space())
         } else if let Some(c) = last_clause_span {
             if p.as_span() == c {

src/verus.pest

@@ -1351,7 +1351,8 @@ loop_expr = {
 
 for_expr = {
     attr* ~ label? ~ for_str ~ pat ~ in_str ~ expr_no_struct ~
-    loop_clause* ~
+    (colon_str ~ expr_no_struct)?
+    ~ loop_clause* ~
     fn_block_expr
 }
 

tests/verus-consistency.rs

@@ -1735,6 +1735,35 @@ fn reverse(v: &mut Vec<u64>)
     }
 }
 
+fn test() {
+    for x in iter: 0..end
+        invariant
+            end == 10,
+    {
+        n += 3;
+    }
+    let x = 2;
+    for x in iter: vec_iter_copy(v)
+        invariant
+            b <==> (forall|i: int| 0 <= i < iter.cur ==> v[i] > 0),
+    {
+        b = b && x > 0;
+    }
+    let y = 3;
+    for x in iter: 0..({
+        let z = end;
+        non_spec();
+        z
+    })
+        invariant
+            n == iter.cur * 3,
+            end == 10,
+    {
+        n += 3;
+        end = end + 0;  // causes end to be non-constant
+    }
+}
+
 }
 "#;
     assert_snapshot!(parse_and_format(file).unwrap(), @r###"
@@ -1761,6 +1790,35 @@ fn reverse(v: &mut Vec<u64>)
         }
     }
 
+    fn test() {
+        for x in iter: 0..end
+            invariant
+                end == 10,
+        {
+            n += 3;
+        }
+        let x = 2;
+        for x in iter: vec_iter_copy(v)
+            invariant
+                b <==> (forall|i: int| 0 <= i < iter.cur ==> v[i] > 0),
+        {
+            b = b && x > 0;
+        }
+        let y = 3;
+        for x in iter: 0..({
+            let z = end;
+            non_spec();
+            z
+        })
+            invariant
+                n == iter.cur * 3,
+                end == 10,
+        {
+            n += 3;
+            end = end + 0;  // causes end to be non-constant
+        }
+    }
+
     } // verus!
     "###);
 }