Add mapNonNullIndexed API to List and Set (#1063)

This commit is contained in:
Josh B
2025-05-08 06:57:28 -07:00
committed by GitHub
parent 919d63e51a
commit 948a20ad0c
7 changed files with 78 additions and 6 deletions

View File

@@ -1,5 +1,5 @@
/*
* Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
* Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -554,6 +554,22 @@ public final class ListNodes {
}
}
public abstract static class mapIndexed extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Specialization
protected VmList eval(VmList self, VmFunction function) {
var builder = self.builder();
long index = 0;
for (var elem : self) {
builder.add(applyLambdaNode.execute(function, index++, elem));
}
LoopNode.reportLoopCount(this, self.getLength());
return builder.build();
}
}
public abstract static class mapNonNull extends ExternalMethod1Node {
@Child private ApplyVmFunction1Node applyLambdaNode = ApplyVmFunction1Node.create();
@@ -570,7 +586,7 @@ public final class ListNodes {
}
}
public abstract static class mapIndexed extends ExternalMethod1Node {
public abstract static class mapNonNullIndexed extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Specialization
@@ -579,7 +595,9 @@ public final class ListNodes {
long index = 0;
for (var elem : self) {
builder.add(applyLambdaNode.execute(function, index++, elem));
var newValue = applyLambdaNode.execute(function, index++, elem);
if (newValue instanceof VmNull) continue;
builder.add(newValue);
}
LoopNode.reportLoopCount(this, self.getLength());
return builder.build();

View File

@@ -1,5 +1,5 @@
/*
* Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
* Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -335,6 +335,22 @@ public final class SetNodes {
}
}
public abstract static class mapIndexed extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Specialization
protected VmSet eval(VmSet self, VmFunction function) {
var builder = self.builder();
long index = 0;
for (var elem : self) {
builder.add(applyLambdaNode.execute(function, index++, elem));
}
LoopNode.reportLoopCount(this, self.getLength());
return builder.build();
}
}
public abstract static class mapNonNull extends ExternalMethod1Node {
@Child private ApplyVmFunction1Node applyLambdaNode = ApplyVmFunction1Node.create();
@@ -351,7 +367,7 @@ public final class SetNodes {
}
}
public abstract static class mapIndexed extends ExternalMethod1Node {
public abstract static class mapNonNullIndexed extends ExternalMethod1Node {
@Child private ApplyVmFunction2Node applyLambdaNode = ApplyVmFunction2NodeGen.create();
@Specialization
@@ -360,7 +376,9 @@ public final class SetNodes {
long index = 0;
for (var elem : self) {
builder.add(applyLambdaNode.execute(function, index++, elem));
var newValue = applyLambdaNode.execute(function, index++, elem);
if (newValue instanceof VmNull) continue;
builder.add(newValue);
}
LoopNode.reportLoopCount(this, self.getLength());
return builder.build();

View File

@@ -332,6 +332,11 @@ examples {
List(1, 2, 3).mapIndexed((i, n) -> i.isOdd && n.isEven)
}
["mapNonNullIndexed()"] {
List(1, 2, 3, 4).mapNonNullIndexed((i, n) -> if (n.isOdd) null else n * i)
List(1, 2, 3, 4, null).mapNonNullIndexed((i, n) -> if (n?.isOdd ?? true) null else n * i)
}
["flatMapIndexed()"] {
List(1, 2, 3).flatMapIndexed((i, n) -> List(n * i))
List(1, 2, 3).flatMapIndexed((i, n) -> List(i.isOdd && n.isEven))

View File

@@ -280,6 +280,11 @@ examples {
Set(1, 2, 3, 3, 2, 1).mapIndexed((i, n) -> i.isOdd && n.isEven)
}
["mapNonNullIndexed()"] {
Set(1, 2, 3, 4, 4, 3, 2, 1).mapNonNullIndexed((i, n) -> if (n.isOdd) null else n * i)
Set(1, 2, 3, 4, null, 4, 3, 2, 1).mapNonNullIndexed((i, n) -> if (n?.isOdd ?? true) null else n * i)
}
["flatMapIndexed()"] {
Set(1, 2, 3, 3, 2, 1).flatMapIndexed((i, n) -> Set(n * i))
Set(1, 2, 3, 3, 2, 1).flatMapIndexed((i, n) -> Set(i.isOdd && n.isEven))

View File

@@ -274,6 +274,10 @@ examples {
List(0, 2, 6)
List(false, true, false)
}
["mapNonNullIndexed()"] {
List(2, 12)
List(2, 12)
}
["flatMapIndexed()"] {
List(0, 2, 6)
List(false, true, false)

View File

@@ -227,6 +227,10 @@ examples {
Set(0, 2, 6)
Set(false, true)
}
["mapNonNullIndexed()"] {
Set(2, 12)
Set(2, 12)
}
["flatMapIndexed()"] {
Set(0, 2, 6)
Set(false, true)

View File

@@ -2441,6 +2441,18 @@ abstract external class Collection<out Element> extends Any {
@AlsoKnownAs { names { "filterMap" } }
abstract function mapNonNull<Result>(transform: (Element) -> Result): Collection<Result(this != null)>
/// Transforms this collection by applying [transform] to each element and removing resulting [null] elements.
///
/// [transform] takes two arguments: the index of the element, and the element itself.
///
/// Facts:
/// ```
/// List(1, 2, 3, 4).mapNonNullIndexed((i, n) -> if (n.isOdd) null else n * i) == List(2, 12)
/// List(1, 2, 3, 4, null).mapNonNullIndexed((i, n) -> if (n?.isOdd ?? true) null else n * i) == List(2, 12)
/// ```
@Since { version = "0.29.0" }
abstract function mapNonNullIndexed<Result>(transform: (Int, Element) -> Result): Collection<Result(this != null)>
/// Applies a collection-generating [transform] to each element in this collection
/// and concatenates the resulting collections.
///
@@ -2874,6 +2886,9 @@ external class List<out Element> extends Collection<Element> {
external function filterIndexed(predicate: (Int, Element) -> Boolean): List<Element>
external function mapIndexed<Result>(transform: (Int, Element) -> Result): List<Result>
@Since { version = "0.29.0" }
external function mapNonNullIndexed<Result>(transform: (Int, Element) -> Result): List<Result(this != null)>
external function flatMapIndexed<Result>(transform: (Int, Element) -> Collection<Result>): List<Result>
external function filterIsInstance<Type>(clazz: Class<Type>): List<Type>
@@ -3082,6 +3097,9 @@ external class Set<out Element> extends Collection<Element> {
external function filterIndexed(predicate: (Int, Element) -> Boolean): Set<Element>
external function mapIndexed<Result>(transform: (Int, Element) -> Result): Set<Result>
@Since { version = "0.29.0" }
external function mapNonNullIndexed<Result>(transform: (Int, Element) -> Result): Set<Result(this != null)>
external function flatMapIndexed<Result>(transform: (Int, Element) -> Collection<Result>): Set<Result>
external function filterIsInstance<Type>(clazz: Class<Type>): Set<Type>