From ee14cbd3d94880f7e02d77849efa97ac56d8cbfb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sami=20Kalliom=C3=A4ki?= <sakal@google.com>
Date: Wed, 16 Nov 2022 10:33:11 +0000
Subject: [PATCH] Switch to manual for loop to support select statements.

---
 docs/collections_doc.md     |  4 +++-
 lib/collections.bzl         | 13 +++++++------
 tests/collections_tests.bzl | 13 ++++++++++---
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/docs/collections_doc.md b/docs/collections_doc.md
index e29feb0..0c74499 100755
--- a/docs/collections_doc.md
+++ b/docs/collections_doc.md
@@ -81,7 +81,9 @@ A new list with all unique elements from `iterable`.
 collections.flatten(<a href="#collections.flatten-iterable">iterable</a>)
 </pre>
 
-Collapses the first dimension of the iterable and returns the result as a list.
+Flattens an iterable to the sum of those iterables.
+
+This behaves like collapsing the first dimension of the iterable.
 
 **PARAMETERS**
 
diff --git a/lib/collections.bzl b/lib/collections.bzl
index 2e777e6..615d1f1 100644
--- a/lib/collections.bzl
+++ b/lib/collections.bzl
@@ -66,7 +66,9 @@ def _uniq(iterable):
     return list(unique_elements.keys())
 
 def _flatten(iterable):
-    """Collapses the first dimension of the iterable and returns the result as a list.
+    """Flattens an iterable to the sum of those iterables.
+
+    This behaves like collapsing the first dimension of the iterable.
 
     Args:
       iterable: An iterable of iterables to be collapsed to a list.
@@ -74,11 +76,10 @@ def _flatten(iterable):
     Returns:
       A new list with the collapsed elements from `iterable`.
     """
-    return [
-        element
-        for subiterable in iterable
-        for element in subiterable
-    ]
+    result = []
+    for element in iterable:
+      result += element
+    return result
 
 collections = struct(
     after_each = _after_each,
diff --git a/tests/collections_tests.bzl b/tests/collections_tests.bzl
index c74bdb8..1c911e4 100644
--- a/tests/collections_tests.bzl
+++ b/tests/collections_tests.bzl
@@ -105,9 +105,16 @@ uniq_test = unittest.make(_uniq_test)
 
 def _flatten_test(ctx):
     env = unittest.begin(ctx)
-    asserts.equals(env, collections.flatten([[0, 1], [2, 3]]), [0, 1, 2, 3])
-    asserts.equals(env, collections.flatten([]), [])
-    asserts.equals(env, collections.flatten([[], [0, 1]]), [0, 1])
+    asserts.equals(env, [0, 1, 2, 3], collections.flatten([[0, 1], [2, 3]]))
+    asserts.equals(env, [], collections.flatten([]), [])
+    asserts.equals(env, [0, 1], collections.flatten([[], [0, 1]]))
+    asserts.equals(env,
+        '[] + select({"//conditions:default": [0, 1]}) + select({"//conditions:default": [2, 3]})',
+        str(collections.flatten([select({
+            "//conditions:default": [0, 1],
+        }), select({
+            "//conditions:default": [2, 3],
+        })])))
 
     return unittest.end(env)