Add mapNonNullIndexed API to List and Set (#1063)

pkl-core/src/main/java/org/pkl/core/stdlib/base/ListNodes.java

@@ -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();

pkl-core/src/main/java/org/pkl/core/stdlib/base/SetNodes.java

@@ -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();

pkl-core/src/test/files/LanguageSnippetTests/input/api/list.pkl

@@ -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))

pkl-core/src/test/files/LanguageSnippetTests/input/api/set.pkl

@@ -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))

pkl-core/src/test/files/LanguageSnippetTests/output/api/list.pcf

@@ -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)

pkl-core/src/test/files/LanguageSnippetTests/output/api/set.pcf

@@ -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)

stdlib/base.pkl

@@ -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>