diff --git a/src/main/java/com/annimon/ownlang/lib/modules/functions/functional_stream.java b/src/main/java/com/annimon/ownlang/lib/modules/functions/functional_stream.java index ea5fd1e..aac8408 100644 --- a/src/main/java/com/annimon/ownlang/lib/modules/functions/functional_stream.java +++ b/src/main/java/com/annimon/ownlang/lib/modules/functions/functional_stream.java @@ -29,7 +29,7 @@ public final class functional_stream implements Function { private final ArrayValue container; public StreamValue(ArrayValue container) { - super(12); + super(13); this.container = container; init(); } @@ -43,6 +43,7 @@ public final class functional_stream implements Function { set("dropWhile", wrapIntermediate(new functional_dropwhile())); set("skip", this::skip); set("limit", this::limit); + set("custom", this::custom); set("reduce", wrapTerminal(new functional_reduce())); set("forEach", wrapTerminal(new functional_foreach())); @@ -82,6 +83,19 @@ public final class functional_stream implements Function { return new StreamValue(new ArrayValue(result)); } + private Value custom(Value... args) { + Arguments.check(1, args.length); + if (args[0].type() != Types.FUNCTION) { + throw new TypeException("Function expected in first argument"); + } + final Function f = ((FunctionValue) args[0]).getValue(); + final Value result = f.execute(container); + if (result.type() == Types.ARRAY) { + return new StreamValue((ArrayValue) result); + } + return result; + } + private FunctionValue wrapIntermediate(Function f) { return wrap(f, true); } diff --git a/src/test/resources/modules/functional/stream.own b/src/test/resources/modules/functional/stream.own index 48647c2..19fa66d 100644 --- a/src/test/resources/modules/functional/stream.own +++ b/src/test/resources/modules/functional/stream.own @@ -1,3 +1,4 @@ +use "std" use "functional" def testStream() { @@ -37,3 +38,17 @@ def testDropWhile() { assertEquals([5, 6, 7, 8], stream(data).dropWhile(def(x) = x % 2 == 0).toArray()) assertEquals(data, stream(data).dropWhile(def(x) = x % 2 != 0).toArray()) } + +def testCustom() { + data = [2,4,6,5] + assertEquals([5,6,4,2], stream(data).custom(::reverse).toArray()) +} + +def reverse(container) { + size = length(container) + result = newarray(size) + for i : range(size) { + result[size - i - 1] = container[i] + } + return result +} \ No newline at end of file