diff --git a/lib/flux-core/src/iter/adapters/mod.rs b/lib/flux-core/src/iter/adapters/mod.rs index 805b80bdd32..bfb19670877 100644 --- a/lib/flux-core/src/iter/adapters/mod.rs +++ b/lib/flux-core/src/iter/adapters/mod.rs @@ -1,3 +1,4 @@ mod enumerate; mod map; mod skip; +mod take; diff --git a/lib/flux-core/src/iter/adapters/take.rs b/lib/flux-core/src/iter/adapters/take.rs new file mode 100644 index 00000000000..08ebc01e68c --- /dev/null +++ b/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; + +#[extern_spec(core::iter)] +#[assoc( + fn size(x: Take) -> int { min(x.n, ::size(x.inner)) } + fn done(x: Take) -> bool { x.n <= 0 || ::done(x.inner) } + fn step(x: Take, y: Take) -> bool { + y.n == x.n - 1 && ::step(x.inner, y.inner) + } +)] +impl Iterator for Take { + #[spec( + fn(self: &mut Self[@curr_s]) -> Option<_>[!::done(curr_s)] + ensures self: Self{next_s: ::step(curr_s, next_s)} + )] + fn next(&mut self) -> Option; +} diff --git a/lib/flux-core/src/iter/traits/iterator.rs b/lib/flux-core/src/iter/traits/iterator.rs index 7d386d28b06..42cf129eb41 100644 --- a/lib/flux-core/src/iter/traits/iterator.rs +++ b/lib/flux-core/src/iter/traits/iterator.rs @@ -40,6 +40,11 @@ trait Iterator { where Self: Sized; + #[spec(fn(Self[@s], n: usize) -> Take[n, s])] + fn take(self, n: usize) -> Take + where + Self: Sized; + #[spec(fn(Self[@s], f: F) where F: FnMut(Self::Item{item: ::valid_item(s, item)}) -> () )] fn for_each(self, f: F) where diff --git a/tests/tests/with_deps/neg/extern_specs/extern_spec_iter00.rs b/tests/tests/with_deps/neg/extern_specs/extern_spec_iter00.rs index b4b0b088e13..23cba53fdb6 100644 --- a/tests/tests/with_deps/neg/extern_specs/extern_spec_iter00.rs +++ b/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 { } // 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(target: &mut [i32], iter: I) +where + I: IntoIterator, +{ + 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; + } +} diff --git a/tests/tests/with_deps/pos/extern_specs/flux_core_iter00.rs b/tests/tests/with_deps/pos/extern_specs/flux_core_iter00.rs index 5c176d220d8..93beead2875 100644 --- a/tests/tests/with_deps/pos/extern_specs/flux_core_iter00.rs +++ b/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(target: &mut [i32], iter: I) +where + I: IntoIterator, +{ + 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(target: &mut [i32], n: usize, iter: I) +where + I: IntoIterator, +{ + 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; + } +}