Add specs for Iterator/Take

lib/flux-core/src/iter/adapters/mod.rs

@@ -1,3 +1,4 @@
 mod enumerate;
 mod map;
 mod skip;
+mod take;

lib/flux-core/src/iter/adapters/take.rs

@@ -0,0 +1,29 @@
+use flux_attrs::*;
+
+defs! {
+    fn min(a: int, b: int) -> int { if a < b { a } else { b } }
+}
+
+/// `Take` is refined by the number of elements left to take (`n`) together with the inner
+/// iterator, so that `done` can be stated exactly: a `Take` is exhausted when it has taken its
+/// quota *or* the inner iterator has run out. Refining by `n` alone would let `next` claim a
+/// `Some` that the inner iterator cannot supply.
+#[extern_spec(core::iter)]
+#[refined_by(n: int, inner: I)]
+struct Take<I>;
+
+#[extern_spec(core::iter)]
+#[assoc(
+    fn size(x: Take<I>) -> int { min(x.n, <I as Iterator>::size(x.inner)) }
+    fn done(x: Take<I>) -> bool { x.n <= 0 || <I as Iterator>::done(x.inner) }
+    fn step(x: Take<I>, y: Take<I>) -> bool {
+        y.n == x.n - 1 && <I as Iterator>::step(x.inner, y.inner)
+    }
+)]
+impl<I: Iterator> Iterator for Take<I> {
+    #[spec(
+        fn(self: &mut Self[@curr_s]) -> Option<_>[!<Self as Iterator>::done(curr_s)]
+        ensures self: Self{next_s: <Self as Iterator>::step(curr_s, next_s)}
+    )]
+    fn next(&mut self) -> Option<I::Item>;
+}

lib/flux-core/src/iter/traits/iterator.rs

@@ -40,6 +40,11 @@ trait Iterator {
     where
         Self: Sized;
 
+    #[spec(fn(Self[@s], n: usize) -> Take<Self>[n, s])]
+    fn take(self, n: usize) -> Take<Self>
+    where
+        Self: Sized;
+
     #[spec(fn(Self[@s], f: F) where F: FnMut(Self::Item{item: <Self as Iterator>::valid_item(s, item)}) -> () )]
     fn for_each<F>(self, f: F)
     where

tests/tests/with_deps/neg/extern_specs/extern_spec_iter00.rs

@@ -1,6 +1,8 @@
 #![allow(unused)]
 use std::slice::Iter;
 
+use flux_rs::{assert, macros::qualifier};
+
 extern crate flux_core;
 
 #[flux_rs::extern_spec(std::slice)]
@@ -24,9 +26,6 @@ trait Iterator {
         Self: Sized;
 }
 
-#[flux_rs::sig(fn (bool[true]))]
-fn assert(_b: bool) {}
-
 #[flux_rs::extern_spec]
 #[flux_rs::assoc(fn done(x: Iter) -> bool { x.idx >= x.len })]
 #[flux_rs::assoc(fn step(x: Iter, y: Iter) -> bool { x.idx + 1 == y.idx && x.len == y.len})]
@@ -146,3 +145,18 @@ pub fn find_index_of_3(slice: &[usize]) -> Option<usize> {
 }
 
 // TODO: implement IntoIter so I can use these with `for` loops
+
+#[flux_rs::spec(fn(target: &mut [i32][4], iter: I))]
+pub fn test_take_easy<I>(target: &mut [i32], iter: I)
+where
+    I: IntoIterator<Item = i32>,
+{
+    let iter = iter.into_iter();
+    let mut pushed = 0;
+    for element in iter.take(5) {
+        // Relates `pushed` to the `Take`'s remaining count, which the `for` desugaring hides.
+        qualifier!(pushed: int, k: int ; pushed + k == 5);
+        target[pushed] = element; //~ ERROR: assertion might fail
+        pushed += 1;
+    }
+}

tests/tests/with_deps/pos/extern_specs/flux_core_iter00.rs

@@ -1,5 +1,5 @@
 extern crate flux_core;
-use flux_rs::assert;
+use flux_rs::{assert, macros::qualifier};
 
 // --- position ---
 
@@ -16,3 +16,37 @@ pub fn test_position_safe_index(xs: &[i32]) {
         let _ = xs[i];
     }
 }
+
+// --- take ---
+
+#[flux_rs::spec(fn(target: &mut [i32][100], iter: I))]
+pub fn test_take_easy<I>(target: &mut [i32], iter: I)
+where
+    I: IntoIterator<Item = i32>,
+{
+    let iter = iter.into_iter();
+    let mut pushed = 0;
+    for element in iter.take(5) {
+        // Relates `pushed` to the `Take`'s remaining count, which the `for` desugaring hides.
+        qualifier!(pushed: int, k: int ; pushed + k == 5);
+        target[pushed] = element;
+        pushed += 1;
+    }
+}
+
+#[flux_rs::spec(fn(target: &mut [i32][@len], n: usize{n <= len}, iter: I))]
+pub fn test_take_loop<I>(target: &mut [i32], n: usize, iter: I)
+where
+    I: IntoIterator<Item = i32>,
+{
+    let iter = iter.into_iter();
+    let mut pushed = n;
+    let to_take = target.len() - pushed;
+    for element in iter.take(to_take) {
+        // As above, but the sum is `len` rather than a literal: `pushed` starts at `n` and
+        // `to_take` is `len - n`.
+        qualifier!(pushed: int, k: int, len: int ; pushed + k == len);
+        target[pushed] = element;
+        pushed += 1;
+    }
+}