#!/usr/bin/ruby

#
## Lazy iterators
#

var iter = (^Inf -> lazy.grep{ .is_prime }.iter)

var a = 10.of(iter)
var b = 10.defs(iter)
var c = 10.by{.is_prime}

assert_eq(a, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29])
assert_eq(b, [31, 37, 41, 43, 47, 53, 59, 61, 67, 71])
assert_eq(a, c)

var jter = (^100 -> lazy.grep { .is_even }.map {|n| 4*n + 1 }.grep{ .is_prime }.iter)

var d = 100.of { jter() \\ break }

assert_eq(d, [17, 41, 73, 89, 97, 113, 137, 193, 233, 241, 257, 281, 313, 337, 353])

var kter = (^Inf -> lazy.grep { .is_odd }.map { |n| 12*n + 1 }.grep { .is_prime })

var e = []
for p in (kter) {
    e << p
    break if (p > 200)
}

assert_eq(e, [13, 37, 61, 109, 157, 181, 229])

for p in (kter) {
    break if (p == 13)
    die "kter error!"
}

do {
    var lz = (1..100 -> lazy)

    assert_eq(lz.map{ _**2 + 1}.grep{.is_prime}.first(10), [2, 5, 17, 37, 101, 197, 257, 401, 577, 677])
    assert_eq(lz.grep{.is_prime}.first(10), [2, 3, 5, 7, 11, 13, 17, 19, 23, 29])
}

do {
    assert_eq(10.by {.is_prime}, [2, 3, 5, 7, 11, 13, 17, 19, 23, 29])
    assert_eq(20.of({.is_prime ? _ : nil}, 20..1000).grep, [23, 29, 31, 37])
    assert_eq(10.by({.is_prime}, 50..1000), [53, 59, 61, 67, 71, 73, 79, 83, 89, 97])
}

do {
    var lz = (50 .. 1000 -> lazy)
    assert_eq(20.of({.is_prime ? _ : nil}, lz).grep, [53, 59, 61, 67])
    assert_eq(10.by({.is_prime}, lz), [53, 59, 61, 67, 71, 73, 79, 83, 89, 97])
    assert_eq(lz.grep{.is_prime}.first(10), [53, 59, 61, 67, 71, 73, 79, 83, 89, 97])
}

do {
    var a = [1, 2, nil, 3, nil, 4, nil, 5, 6, 7, 8, 21, 2, 31, 42, 5, 2]

    assert_eq(10.defs { a[_] }, [1, 2, 3, 4, 5, 6, 7, 8, 21, 2])
    assert_eq(10.by { a[_] }.map { a[_] }, [1, 2, 3, 4, 5, 6, 7, 8, 21, 2])

    var lz = (2..1000 -> lazy)

    assert_eq(10.defs({ a[_] }, lz), [3, 4, 5, 6, 7, 8, 21, 2, 31, 42])
    assert_eq(10.by({ a[_] }, lz).map { a[_] }, [3, 4, 5, 6, 7, 8, 21, 2, 31, 42])
}

say "** Test passed!"