array.heron

language heron:std:0.1;

module heron:std.array:0.1
{
    function unit(x)
        = [x];

    function map(xs, f)
        = array(xs.count, i => f(xs[i]));

    function mapWithIndex(xs, f)
        = array(xs.count, i => f(xs[i], i));

    function mapIndex(xs, f)
        = array(xs.count, f);

    function min(x, y)
        = x <= y ? x : y;

    function max(x, y)
        = x >= y ? x : y;

    function shorter(xs, ys)
        = xs.count <= ys.count ? xs : ys;

    function longer(xs, ys)
        = xs.count >= ys.count ? xs : ys;

    function empty(xs)
        = xs.count == 0;

    function selectByIndex(xs, indices)
        = indices.map(i => xs.at(i));

    function indices(xs)
        = 0 .. xs.count;

    function zip(xs, ys, f)
        = xs.count <= ys.count
            ? xs.mapWithIndex((x, i) => f(x, ys[i]))
            : ys.mapWithIndex((y, i) => f(xs[i], y));

    function all(xs, p)
        = xs.reduce(true, (prev, x) => prev && p(x));

    function any(xs, p)
        = xs.reduce(false, (prev, x) => prev || p(x));

    //function count(xs, p)
    //    = xs.reduce(0, (prev, x) => p(x) ? prev + 1 : prev);

    function eq(xs, ys)
        = xs.count == ys.count;

    function filter(xs, p) {
        var ys = xs.mutable;
        var i = 0;
        for (var x in xs)
            if (p(x))
                ys[i++] = x;
        return ys.immutable.take(i);
    }

    function repeat(x, n)
        = (0 .. n).map(i => x);

    function prefixScan(xs, op) {
        if (xs.empty) return xs;
        var ys = xs[0].repeat(xs.count).mutable;
        for (var i in 1 .. ys.count)
            ys[i] = op(xs[i], ys[i-1]);
        return ys.immutable;
    }

    function adjacentDifferences(xs)
        = xs.indices.map(i => i > 0 ? xs[i] - xs[i-1] : xs[i]);

    function slice(xs, from, to)
        = (from .. to).map(i => xs.at(i));

    function stride(xs, n)
        = (0 .. xs.count / n).map(i => xs[i * n]);

    function stride(xs, from, n)
        = (0 .. xs.count / n).map(i => xs[from + i * n]);

    function strides(xs, n)
        = (0 .. n).map(i => xs.stride(i, n));

    function slices(xs, n)
        = (0 .. n).map(i => xs.slice(i * n, (i+1) * n));

    function take(xs, n)
        = xs.slice(0, n);

    function take(xs, i, n)
        = xs.skip(i).take(n);

    function skip(xs, n)
        = xs.slice(n, xs.count - n);

    function drop(xs, n)
        = xs.take(xs.count - n);

    function last(xs, n)
        = xs.skip(xs.count-n);

    function reverse(xs, n)
        = xs.indices.map(i => xs[xs.count-1-i]);

    function gen(cnt, f)
        = (0 .. cnt).map(f);

    function concat(xs, ys)
        = gen(xs.count + ys.count, i => i < xs.count ? xs[i] : ys[i - xs.count]);

    function cut(xs, from, n)
        = gen(xs.count - n, i => i < from ? xs[i] : xs[i + n]);

    function splice(xs, from, ys)
        = xs.take(from).concat(ys).concat(xs.skip(from));

    function sum(xs)
        = xs.reduce(0.0, op+);

    function product(xs)
        = xs.reduce(1.0, op*);

    function average(xs)
        = xs.sum / xs.count.float;

    function min(xs)
        = xs.reduce(xs[0], min);

    function max(xs)
        = xs.reduce(xs[0], max);

    function swapElements(xs, i, j) {
        var tmp = xs[i];
        xs[i] = xs[j];
        xs[j] = tmp;
        return xs;
    }

    function partition(a, lo, hi) {
        var p = a[lo];
        var i = lo - 1;
        var j = hi + 1;
        while (true) {
            do { i++; } while (a[i] < p);
            do { j--; } while (a[j] > p);
            if (i >= j) return j;
            swapElements(a, i, j);
        }
    }

    function qsort(a, lo, hi) {
        if (lo < hi) {
            var p = partition(a, lo, hi);
            qsort(a, lo, p);
            qsort(a, p+1, hi);
        }
        return a;
    }

    function sort(xs)
        = xs.mutable.qsort(0, xs.count-1).immutable;

    function median(xs) {
        var ys = xs.sort;
        return (ys.count - 1) % 2 == 0
            ? ys[(ys.count - 1) / 2]
            : ys[(ys.count - 2) / 2] + ys[ys.count / 2] / 2;
    }

    function inRange(xs, n)
        = n >= 0 && n < xs.count;

    function last(xs)
        = xs[xs.count - 1];

    function first(xs)
        = xs[0];

    function tail(xs)
        = xs.skip(1);

    function reduce(xs, acc, f) {
        for (var x in xs)
            acc = f(acc, x);
        return acc;
    }

    function flatten(xs)
        = xs.reduce([], concat);

    function flatMap(xs, f)
        = xs.map(f).flatten;

    // TODO: rewrite using flatMap, but that is slow for now. I will need to write an optimizer    
    function cartesianProduct(xs, ys, f) {
        var r = [].mutable;
        for (var x in xs)
            for (var y in ys) 
                r.push(f(x, y));
        return r.immutable;
    }

    function setAll(xs, x)
        = xs.map(_ => x);
}