diff --git a/doc/site/modules/core/sequence.markdown b/doc/site/modules/core/sequence.markdown index 968b88bf..e1bb5c9d 100644 --- a/doc/site/modules/core/sequence.markdown +++ b/doc/site/modules/core/sequence.markdown @@ -125,6 +125,23 @@ It is a runtime error to call this on an empty sequence. Similar to above, but uses `seed` for the initial value of the accumulator. If the sequence is empty, returns `seed`. +### **skip**(count) + +Creates a new sequence that skips the first `count` elements of the original +sequence. + +The returned sequence is *lazy*. The first `count` elements are only skipped +once you start to iterate the returned sequence. Changes to the original +sequence will be reflected in the filtered sequence. + +### **take**(count) + +Creates a new sequence that iterates only the first `count` elements of the +original sequence. + +The returned sequence is *lazy*. Changes to the original sequence will be +reflected in the filtered sequence. + ### **toList** Creates a [list][] containing all the elements in the sequence. diff --git a/src/vm/wren_core.wren b/src/vm/wren_core.wren index 3d756964..36660f54 100644 --- a/src/vm/wren_core.wren +++ b/src/vm/wren_core.wren @@ -56,6 +56,10 @@ class Sequence { map(transformation) { MapSequence.new(this, transformation) } + skip(count) { SkipSequence.new(this, count) } + + take(count) { TakeSequence.new(this, count) } + where(predicate) { WhereSequence.new(this, predicate) } reduce(acc, f) { @@ -112,6 +116,43 @@ class MapSequence is Sequence { iteratorValue(iterator) { _fn.call(_sequence.iteratorValue(iterator)) } } +class SkipSequence is Sequence { + construct new(sequence, count) { + _sequence = sequence + _count = count + } + + iterate(iterator) { + if (iterator) { + return _sequence.iterate(iterator) + } else { + iterator = _sequence.iterate(iterator) + var count = _count + while (count > 0 && iterator) { + iterator = _sequence.iterate(iterator) + count = count - 1 + } + return iterator + } + } + + iteratorValue(iterator) { _sequence.iteratorValue(iterator) } +} + +class TakeSequence is Sequence { + construct new(sequence, count) { + _sequence = sequence + _count = count + } + + iterate(iterator) { + if (!iterator) _taken = 1 else _taken = _taken + 1 + return _taken > _count ? null : _sequence.iterate(iterator) + } + + iteratorValue(iterator) { _sequence.iteratorValue(iterator) } +} + class WhereSequence is Sequence { construct new(sequence, fn) { _sequence = sequence diff --git a/src/vm/wren_core.wren.inc b/src/vm/wren_core.wren.inc index 97d59b63..7548dca4 100644 --- a/src/vm/wren_core.wren.inc +++ b/src/vm/wren_core.wren.inc @@ -58,6 +58,10 @@ static const char* coreModuleSource = "\n" " map(transformation) { MapSequence.new(this, transformation) }\n" "\n" +" skip(count) { SkipSequence.new(this, count) }\n" +"\n" +" take(count) { TakeSequence.new(this, count) }\n" +"\n" " where(predicate) { WhereSequence.new(this, predicate) }\n" "\n" " reduce(acc, f) {\n" @@ -114,6 +118,43 @@ static const char* coreModuleSource = " iteratorValue(iterator) { _fn.call(_sequence.iteratorValue(iterator)) }\n" "}\n" "\n" +"class SkipSequence is Sequence {\n" +" construct new(sequence, count) {\n" +" _sequence = sequence\n" +" _count = count\n" +" }\n" +"\n" +" iterate(iterator) {\n" +" if (iterator) {\n" +" return _sequence.iterate(iterator)\n" +" } else {\n" +" iterator = _sequence.iterate(iterator)\n" +" var count = _count\n" +" while (count > 0 && iterator) {\n" +" iterator = _sequence.iterate(iterator)\n" +" count = count - 1\n" +" }\n" +" return iterator\n" +" }\n" +" }\n" +"\n" +" iteratorValue(iterator) { _sequence.iteratorValue(iterator) }\n" +"}\n" +"\n" +"class TakeSequence is Sequence {\n" +" construct new(sequence, count) {\n" +" _sequence = sequence\n" +" _count = count\n" +" }\n" +"\n" +" iterate(iterator) {\n" +" if (!iterator) _taken = 1 else _taken = _taken + 1\n" +" return _taken > _count ? null : _sequence.iterate(iterator)\n" +" }\n" +"\n" +" iteratorValue(iterator) { _sequence.iteratorValue(iterator) }\n" +"}\n" +"\n" "class WhereSequence is Sequence {\n" " construct new(sequence, fn) {\n" " _sequence = sequence\n" diff --git a/test/core/sequence/skip.wren b/test/core/sequence/skip.wren new file mode 100644 index 00000000..b54a51fc --- /dev/null +++ b/test/core/sequence/skip.wren @@ -0,0 +1,28 @@ +class TestSequence is Sequence { + construct new() {} + + iterate(iterator) { + if (iterator == null) return 1 + if (iterator == 3) return false + return iterator + 1 + } + + iteratorValue(iterator) { iterator } +} + +var test = TestSequence.new().skip(0) + +System.print(test is Sequence) // expect: true +System.print(test) // expect: instance of SkipSequence + +// Skipping 0 changes nothing +System.print(test.toList) // expect: [1, 2, 3] + +// Skipping 1 works +System.print(test.skip(1).toList) // expect: [2, 3] + +// Skipping more than length of sequence produces empty list +System.print(test.skip(4).isEmpty) // expect: true + +// Skipping less than 0 changes nothing +System.print(test.skip(-10).toList) // expect: [1, 2, 3] diff --git a/test/core/sequence/take.wren b/test/core/sequence/take.wren new file mode 100644 index 00000000..81c188e8 --- /dev/null +++ b/test/core/sequence/take.wren @@ -0,0 +1,28 @@ +class TestSequence is Sequence { + construct new() {} + + iterate(iterator) { + if (iterator == null) return 1 + if (iterator == 3) return false + return iterator + 1 + } + + iteratorValue(iterator) { iterator } +} + +var test = TestSequence.new().take(3) + +System.print(test is Sequence) // expect: true +System.print(test) // expect: instance of TakeSequence + +// Taking 0 produces empty list +System.print(test.take(0).isEmpty) // expect: true + +// Taking 1 works +System.print(test.take(1).toList) // expect: [1] + +// Taking more than length of sequence produces whole sequence +System.print(test.take(4).toList) // expect: [1, 2, 3] + +// Taking less than 0 produces empty list +System.print(test.take(-10).isEmpty) // expect: true