scurry 0.5.0

decl StopIteration {}

decl Iterator {
    next_fn: Fn

    fn $new(self, next_fn: Fn) {
        self.next_fn = next_fn;
    }

    exp fn next(self) {
        return self.next_fn();
    }

    exp fn map(self, func: Fn) > Iterator {
        new_fn = fn() {
            next = self.next();
            if next == StopIteration() {
                return next;
            } else {
                return func(next);
            }
        };
        return Iterator(new_fn);
    }

    exp fn contains(self, item) > Bool {
        next = self.next();
        while next != item {
            if next == StopIteration() {
                return False;
            }
            next = self.next();
        }
        return True;
    }

    exp fn all(self, predicate: Fn) > Bool {
        next = self.next();
        while next != StopIteration() {
            if !predicate(next) {
                return False;
            }
            next = self.next();
        }
        return True;
    }

    exp fn any(self, predicate: Fn) > Bool {
        next = self.next();
        while next != StopIteration() {
            if predicate(next) {
                return True;
            }
            next = self.next();
        }
        return False;
    }

    exp fn skip(self, amount: Int) > Iterator {
        for _ in Range(0, amount) {
            self.next();
        }
        return Iterator(self.next);
    }

    exp fn filter(self, predicate: Fn) > Iterator {
        new_fn = fn() {
            next = self.next();
            while next != StopIteration() && !predicate(next) {
                next = self.next();
            }
            return next;
        };
        return Iterator(new_fn);
    }

    exp fn fold(self, init: Any, func: Fn) {
        accum = init;
        next = self.next();
        while next != StopIteration() {
            accum = func(accum, next);
            next = self.next();
        }
        return accum;
    }

    exp fn reduce(self, func: Fn) {
        init = self.next();
        return self.fold(init, func);
    }

    exp fn take(self, n: Int) {
        // Hacky way to share state with new_fn
        current = [n];
        new_fn = fn() {
            if current[0] != 0 {
                current[0] -= 1;
                return self.next();
            } else {
                return StopIteration();
            }
        };
        return Iterator(new_fn);
    }

    exp fn nth(self, n: Int) {
        for _ in Range(0, n) {
            self.next();
        }
        return self.next();
    }

    exp fn to_array(self) > Array {
        arr = [];
        next = self.next();
        while next != StopIteration() {
            arr.push(next);
            next = self.next();
        }
        return arr;
    }

    exp fn chain(self, other: Iterator) > Iterator {
        new_fn = fn() {
            next = self.next();
            if next == StopIteration() {
                return other.next();
            }
            return next;
        };
        return Iterator(new_fn);
    }
}

decl Range {
    current: Int,
    end: Int

    [Iterator] {
        .next_
    }

    fn $new(self, start: Int, end: Int) {
        self.current = start - 1;
        self.end = end;
    }

    fn next_(self) {
        self.current += 1;
        if self.current >= self.end {
            return StopIteration();
        }
        return self.current;
    }
}