Skip to content

Add iterate and iterateIndexed methods for Collections in the stdlib #1113

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add iterate and iterateIndexed methods for Collections in the stdlib #1113

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e3b761
Add `iterate` and `iterateIndexed` methods for `Collections` in the s…
stackoverflow Jul 1, 2025
e3d5187
Address review comments.
stackoverflow Jul 1, 2025
d353b85
Remove iterate functions and add `folded` function to short-circuit t…
stackoverflow Jul 14, 2025
a3ec194
Change folded function to not throw an exception
stackoverflow Jul 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 76 additions & 0 deletions pkl-core/src/main/java/org/pkl/core/stdlib/base/ListNodes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,10 @@
import org.pkl.core.ast.internal.IsInstanceOfNodeGen;
import org.pkl.core.ast.lambda.*;
import org.pkl.core.ast.type.TypeNode;
import org.pkl.core.ast.type.TypeNode.BooleanTypeNode;
import org.pkl.core.ast.type.TypeNode.PairTypeNode;
import org.pkl.core.ast.type.TypeNode.UInt8TypeAliasTypeNode;
import org.pkl.core.ast.type.TypeNode.UnknownTypeNode;
import org.pkl.core.ast.type.VmTypeMismatchException;
import org.pkl.core.runtime.*;
import org.pkl.core.stdlib.*;
Expand Down Expand Up @@ -785,6 +788,79 @@ protected Object eval(VmList self, Object initial, VmFunction function) {
}
}

public abstract static class iterate extends ExternalMethod2Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Child @LateInit private TypeNode typeNode;

private TypeNode getTypeNode() {
if (typeNode == null) {
CompilerDirectives.transferToInterpreterAndInvalidate();
var section = VmUtils.unavailableSourceSection();
var booleanTypeNode = new BooleanTypeNode(section);
var unknownTypeNode = new UnknownTypeNode(section);
typeNode = new PairTypeNode(section, booleanTypeNode, unknownTypeNode);
}
return typeNode;
}

@Specialization
protected Object eval(VirtualFrame frame, VmList self, Object initial, VmFunction function) {
var iter = self.iterator();
var typeNode = getTypeNode();
var result = initial;
var loop = 0;

while (iter.hasNext()) {
var elem = iter.next();
var pairUnchecked = applyLambdaNode.execute(function, result, elem);
var pair = (VmPair) typeNode.executeEagerly(frame, pairUnchecked);
result = pair.getSecond();
loop++;
if (!(Boolean) pair.getFirst()) break;
}

LoopNode.reportLoopCount(this, loop);
return result;
}
}

public abstract static class iterateIndexed extends ExternalMethod2Node {
@Child private ApplyVmFunction3Node applyLambdaNode = ApplyVmFunction3NodeGen.create();
@Child @LateInit private TypeNode typeNode;

private TypeNode getTypeNode() {
if (typeNode == null) {
CompilerDirectives.transferToInterpreterAndInvalidate();
var section = VmUtils.unavailableSourceSection();
var booleanTypeNode = new BooleanTypeNode(section);
var unknownTypeNode = new UnknownTypeNode(section);
typeNode = new PairTypeNode(section, booleanTypeNode, unknownTypeNode);
}
return typeNode;
}

@Specialization
protected Object eval(VirtualFrame frame, VmList self, Object initial, VmFunction function) {
var iter = self.iterator();
var typeNode = getTypeNode();
var result = initial;
long index = 0;
var loop = 0;

while (iter.hasNext()) {
var elem = iter.next();
var pairUnchecked = applyLambdaNode.execute(function, index++, result, elem);
var pair = (VmPair) typeNode.executeEagerly(frame, pairUnchecked);
result = pair.getSecond();
loop++;
if (!(Boolean) pair.getFirst()) break;
}

LoopNode.reportLoopCount(this, loop);
return result;
}
}

public abstract static class reduce extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();

Expand Down
79 changes: 79 additions & 0 deletions pkl-core/src/main/java/org/pkl/core/stdlib/base/SetNodes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.oracle.truffle.api.CompilerDirectives;
import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
import com.oracle.truffle.api.dsl.Specialization;
import com.oracle.truffle.api.frame.VirtualFrame;
import com.oracle.truffle.api.nodes.LoopNode;
import org.pkl.core.ast.expression.binary.GreaterThanNode;
import org.pkl.core.ast.expression.binary.GreaterThanNodeGen;
Expand All @@ -26,11 +27,16 @@
import org.pkl.core.ast.internal.IsInstanceOfNode;
import org.pkl.core.ast.internal.IsInstanceOfNodeGen;
import org.pkl.core.ast.lambda.*;
import org.pkl.core.ast.type.TypeNode;
import org.pkl.core.ast.type.TypeNode.BooleanTypeNode;
import org.pkl.core.ast.type.TypeNode.PairTypeNode;
import org.pkl.core.ast.type.TypeNode.UnknownTypeNode;
import org.pkl.core.runtime.*;
import org.pkl.core.stdlib.*;
import org.pkl.core.stdlib.base.CollectionNodes.CompareByNode;
import org.pkl.core.stdlib.base.CollectionNodes.CompareNode;
import org.pkl.core.stdlib.base.CollectionNodes.CompareWithNode;
import org.pkl.core.util.LateInit;

// duplication between ListNodes and SetNodes is "intentional"
// (sharing nodes between VmCollection subtypes results in
Expand Down Expand Up @@ -561,6 +567,79 @@ protected Object eval(VmSet self, Object initial, VmFunction function) {
}
}

public abstract static class iterate extends ExternalMethod2Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Child @LateInit private TypeNode typeNode;

private TypeNode getTypeNode() {
if (typeNode == null) {
CompilerDirectives.transferToInterpreterAndInvalidate();
var section = VmUtils.unavailableSourceSection();
var booleanTypeNode = new BooleanTypeNode(section);
var unknownTypeNode = new UnknownTypeNode(section);
typeNode = new PairTypeNode(section, booleanTypeNode, unknownTypeNode);
}
return typeNode;
}

@Specialization
protected Object eval(VirtualFrame frame, VmSet self, Object initial, VmFunction function) {
var iter = self.iterator();
var typeNode = getTypeNode();
var result = initial;
var loop = 0;

while (iter.hasNext()) {
var elem = iter.next();
var pairUnchecked = applyLambdaNode.execute(function, result, elem);
var pair = (VmPair) typeNode.executeEagerly(frame, pairUnchecked);
result = pair.getSecond();
loop++;
if (!(Boolean) pair.getFirst()) break;
}

LoopNode.reportLoopCount(this, loop);
return result;
}
}

public abstract static class iterateIndexed extends ExternalMethod2Node {
@Child private ApplyVmFunction3Node applyLambdaNode = ApplyVmFunction3NodeGen.create();
@Child @LateInit private TypeNode typeNode;

private TypeNode getTypeNode() {
if (typeNode == null) {
CompilerDirectives.transferToInterpreterAndInvalidate();
var section = VmUtils.unavailableSourceSection();
var booleanTypeNode = new BooleanTypeNode(section);
var unknownTypeNode = new UnknownTypeNode(section);
typeNode = new PairTypeNode(section, booleanTypeNode, unknownTypeNode);
}
return typeNode;
}

@Specialization
protected Object eval(VirtualFrame frame, VmSet self, Object initial, VmFunction function) {
var iter = self.iterator();
var typeNode = getTypeNode();
var result = initial;
long index = 0;
var loop = 0;

while (iter.hasNext()) {
var elem = iter.next();
var pairUnchecked = applyLambdaNode.execute(function, index++, result, elem);
var pair = (VmPair) typeNode.executeEagerly(frame, pairUnchecked);
result = pair.getSecond();
loop++;
if (!(Boolean) pair.getFirst()) break;
}

LoopNode.reportLoopCount(this, loop);
return result;
}
}

public abstract static class reduce extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();

Expand Down
14 changes: 14 additions & 0 deletions pkl-core/src/test/files/LanguageSnippetTests/input/api/list.pkl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,20 @@ examples {
List().foldBack(0, (x, acc) -> x + acc)
}

["iterate()"] {
list1.iterate(0, (acc, x) -> Pair(x < 2, x + acc))
list1.iterate(0, (acc, x) -> Pair(true, x + acc))
module.catch(() -> list1.iterate(0, (acc, x) -> x + acc))
module.catch(() -> list1.iterate(0, (acc, x) -> Pair(x, x + acc)))
}

["iterateIndexed()"] {
list1.iterateIndexed(0, (i, acc, x) -> Pair(i < 1, x + acc))
list1.iterateIndexed(0, (i, acc, x) -> Pair(true, x + acc))
module.catch(() -> list1.iterateIndexed(0, (i, acc, x) -> x + acc))
module.catch(() -> list1.iterateIndexed(0, (i, acc, x) -> Pair(i, x + acc)))
}

["reduce()"] {
list1.reduce((x, y) -> x + y)
List(1).reduce((x, y) -> x + y)
Expand Down
12 changes: 12 additions & 0 deletions pkl-core/src/test/files/LanguageSnippetTests/input/api/set.pkl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,18 @@ examples {
Set().foldBack(0, (x, acc) -> x + acc)
}

["iterate()"] {
set1.iterate(0, (acc, x) -> Pair(x < 2, x + acc))
set1.iterate(0, (acc, x) -> Pair(true, x + acc))
module.catch(() -> set1.iterate(0, (acc, x) -> Pair(1, x + acc)))
}

["iterateIndexed()"] {
set1.iterateIndexed(0, (i, acc, x) -> Pair(i < 1, x + acc))
set1.iterateIndexed(0, (i, acc, x) -> Pair(true, x + acc))
module.catch(() -> set1.iterateIndexed(0, (i, acc, x) -> Pair(i, x + acc)))
}

["reduce()"] {
set1.reduce((x, y) -> x + y)
Set(1).reduce((x, y) -> x + y)
Expand Down
12 changes: 12 additions & 0 deletions pkl-core/src/test/files/LanguageSnippetTests/output/api/list.pcf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,18 @@ examples {
1
0
}
["iterate()"] {
3
6
"Expected value of type `Pair`, but got type `Int`. Value: 1"
"Expected value of type `Boolean`, but got type `Int`. Value: 1"
}
["iterateIndexed()"] {
3
6
"Expected value of type `Pair`, but got type `Int`. Value: 1"
"Expected value of type `Boolean`, but got type `Int`. Value: 0"
}
["reduce()"] {
6
1
Expand Down
10 changes: 10 additions & 0 deletions pkl-core/src/test/files/LanguageSnippetTests/output/api/set.pcf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,16 @@ examples {
1
0
}
["iterate()"] {
3
6
"Expected value of type `Boolean`, but got type `Int`. Value: 1"
}
["iterateIndexed()"] {
3
6
"Expected value of type `Boolean`, but got type `Int`. Value: 0"
}
["reduce()"] {
6
1
Expand Down
18 changes: 18 additions & 0 deletions stdlib/base.pkl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2571,6 +2571,18 @@ abstract external class Collection<out Element> extends Any {
/// The first parameter of [operator] is the zero-based index of the current element.
abstract function foldIndexed<Result>(initial: Result, operator: (Int, Result, Element) -> Result): Result

/// Iterate over elements of this collection in order, accumulating a result.
/// The looping function returns a [Pair] where the first value is a [Boolean] indicating if the iteration
/// should keep going ([true]) or stop, and the second value is the accumulated result.
/// This function is similar to [fold], but it can short-circuit.
abstract function iterate<Result>(initial: Result, loopFunction: (Result, Element) -> Pair<Boolean, Result>): Result
Copy link
Member

@bioball bioball Jul 11, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there other libraries that also call this iterate?

I also wonder if we should use a clojure-style "return a reduced value" to short-circuit?

In addition to fold, we also have foldBack and reduce. It'd be strange if you can short-circuit here, but not in the other two methods.

Copy link
Contributor Author

@stackoverflow stackoverflow Jul 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, that's a good point. I can change the type signature of the folding function to return Element|Reduced<Element> and just change the original functions, instead of creating new ones.

Edit: scratch that. I went with the Clojure approach of having a short-circuiting function that you can call from within any folding/reducing function.


/// Iterate over the indexes and elements of this collection in order, accumulating a result.
/// The looping function returns a [Pair] where the first value is a [Boolean] indicating if the iteration
/// should keep going ([true]) or stop, and the second value is the accumulated result.
/// This function is similar to [foldIndexed], but it can short-circuit.
abstract function iterateIndexed<Result>(initial: Result, loopFunction: (Int, Result, Element) -> Pair<Boolean, Result>): Result

/// Folds this collection in iteration order using [operator], starting with the first element.
///
/// Throws if this collection is empty.
Expand Down Expand Up @@ -3050,6 +3062,9 @@ external class List<out Element> extends Collection<Element> {
external function foldBack<Result>(initial: Result, operator: (Element, Result) -> Result): Result
external function foldIndexed<Result>(initial: Result, operator: (Int, Result, Element) -> Result): Result

external function iterate<Result>(initial: Result, loopFunction: (Element, Result) -> Pair<Boolean, Result>): Result
external function iterateIndexed<Result>(initial: Result, loopFunction: (Int, Result, Element) -> Pair<Boolean, Result>): Result

external function reduce<Result>(operator: (Element|Result, Element) -> Result): Result
external function reduceOrNull<Result>(operator: (Element|Result, Element) -> Result): Result?

Expand Down Expand Up @@ -3191,6 +3206,9 @@ external class Set<out Element> extends Collection<Element> {
external function foldBack<Result>(initial: Result, operator: (Element, Result) -> Result): Result
external function foldIndexed<Result>(initial: Result, operator: (Int, Result, Element) -> Result): Result

external function iterate<Result>(initial: Result, loopFunction: (Element, Result) -> Pair<Boolean, Result>): Result
external function iterateIndexed<Result>(initial: Result, loopFunction: (Int, Result, Element) -> Pair<Boolean, Result>): Result

external function reduce<Result>(operator: (Element|Result, Element) -> Result): Result
external function reduceOrNull<Result>(operator: (Element|Result, Element) -> Result): Result?

Expand Down