From 6861ef7707a56725769594aaa725518cbe65f628 Mon Sep 17 00:00:00 2001
From: Silvan Mosberger <silvan.mosberger@tweag.io>
Date: Wed, 1 Feb 2023 18:27:23 +0100
Subject: [PATCH] lib.lists.ifilter0: init

---
 lib/default.nix    |  2 +-
 lib/lists.nix      | 50 +++++++++++++++++++++++++++++++++++++++++++++-
 lib/tests/misc.nix | 27 +++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/lib/default.nix b/lib/default.nix
index 486d412fbb6fd..53ec802427db0 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -89,7 +89,7 @@ let
       recurseIntoAttrs dontRecurseIntoAttrs cartesianProduct cartesianProductOfSets
       mapCartesianProduct updateManyAttrsByPath;
     inherit (self.lists) singleton forEach foldr fold foldl foldl' imap0 imap1
-      concatMap flatten remove findSingle findFirst any all count
+      ifilter0 concatMap flatten remove findSingle findFirst any all count
       optional optionals toList range replicate partition zipListsWith zipLists
       reverseList listDfs toposort sort sortOn naturalSort compareLists take
       drop sublist last init crossLists unique allUnique intersectLists
diff --git a/lib/lists.nix b/lib/lists.nix
index 28fa277b22b1f..ca436d7a9c94b 100644
--- a/lib/lists.nix
+++ b/lib/lists.nix
@@ -4,7 +4,7 @@
 { lib }:
 let
   inherit (lib.strings) toInt;
-  inherit (lib.trivial) compare min id warn;
+  inherit (lib.trivial) compare min id warn pipe;
   inherit (lib.attrsets) mapAttrs;
 in
 rec {
@@ -333,6 +333,54 @@ rec {
   */
   imap1 = f: list: genList (n: f (n + 1) (elemAt list n)) (length list);
 
+  /**
+    Filter a list for elements that satisfy a predicate function.
+    The predicate function is called with both the index and value for each element.
+    It must return `true`/`false` to include/exclude a given element in the result.
+    This function is strict in the result of the predicate function for each element.
+    This function has O(n) complexity.
+
+    Also see [`builtins.filter`](https://nixos.org/manual/nix/stable/language/builtins.html#builtins-filter) (available as `lib.lists.filter`),
+    which can be used instead when the index isn't needed.
+
+    # Inputs
+
+    `ipred`
+
+    : The predicate function, it takes two arguments:
+      - 1. (int): the index of the element.
+      - 2. (a): the value of the element.
+
+      It must return `true`/`false` to include/exclude a given element from the result.
+
+    `list`
+
+    : The list to filter using the predicate.
+
+    # Type
+    ```
+    ifilter0 :: (int -> a -> bool) -> [a] -> [a]
+    ```
+
+    # Examples
+    :::{.example}
+    ## `lib.lists.ifilter0` usage example
+
+    ```nix
+    ifilter0 (i: v: i == 0 || v > 2) [ 1 2 3 ]
+    => [ 1 3 ]
+    ```
+    :::
+  */
+  ifilter0 =
+    ipred:
+    input:
+    map (idx: elemAt input idx) (
+      filter (idx: ipred idx (elemAt input idx)) (
+        genList (x: x) (length input)
+      )
+    );
+
   /**
     Map and concatenate the result.
 
diff --git a/lib/tests/misc.nix b/lib/tests/misc.nix
index cf4a185c14689..6774939023d20 100644
--- a/lib/tests/misc.nix
+++ b/lib/tests/misc.nix
@@ -63,8 +63,10 @@ let
     hasAttrByPath
     hasInfix
     id
+    ifilter0
     isStorePath
     lazyDerivation
+    length
     lists
     listToAttrs
     makeExtensible
@@ -651,6 +653,31 @@ runTests {
     expected = ["b" "c"];
   };
 
+  testIfilter0Example = {
+    expr = ifilter0 (i: v: i == 0 || v > 2) [ 1 2 3 ];
+    expected = [ 1 3 ];
+  };
+  testIfilter0Empty = {
+    expr = ifilter0 (i: v: abort "shouldn't be evaluated!") [ ];
+    expected = [ ];
+  };
+  testIfilter0IndexOnly = {
+    expr = length (ifilter0 (i: v: mod i 2 == 0) [ (throw "0") (throw "1") (throw "2") (throw "3")]);
+    expected = 2;
+  };
+  testIfilter0All = {
+    expr = ifilter0 (i: v: true) [ 10 11 12 13 14 15 ];
+    expected = [ 10 11 12 13 14 15 ];
+  };
+  testIfilter0First = {
+    expr = ifilter0 (i: v: i == 0) [ 10 11 12 13 14 15 ];
+    expected = [ 10 ];
+  };
+  testIfilter0Last = {
+    expr = ifilter0 (i: v: i == 5) [ 10 11 12 13 14 15 ];
+    expected = [ 15 ];
+  };
+
   testFold =
     let
       f = op: fold: fold op 0 (range 0 100);