[kairos] add 5-way combine() overload

Flag: EXEMPT unused
Test: atest kairos-tests
Change-Id: Ia3ce305702419f5a498c3a8341879b5aadae6e55

2 files changed, 73 insertions, 0 deletions

diff --git a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/TState.kt b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/TState.kt
index a4c695657f8d..46b0f075af22 100644
--- a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/TState.kt
+++ b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/TState.kt
@@ -312,6 +312,57 @@ fun <A, B, C, D, Z> combine(
     )
 }
 
+/**
+ * Returns a [TState] whose value is generated with [transform] by combining the current values of
+ * each given [TState].
+ *
+ * @see TState.combineWith
+ */
+@ExperimentalFrpApi
+fun <A, B, C, D, E, Z> combine(
+    stateA: TState<A>,
+    stateB: TState<B>,
+    stateC: TState<C>,
+    stateD: TState<D>,
+    stateE: TState<E>,
+    transform: suspend FrpScope.(A, B, C, D, E) -> Z,
+): TState<Z> {
+    val operatorName = "combine"
+    val name = operatorName
+    return TStateInit(
+        init(name) {
+            coroutineScope {
+                val dl1: Deferred<TStateImpl<A>> = async {
+                    stateA.init.connect(evalScope = this@init)
+                }
+                val dl2: Deferred<TStateImpl<B>> = async {
+                    stateB.init.connect(evalScope = this@init)
+                }
+                val dl3: Deferred<TStateImpl<C>> = async {
+                    stateC.init.connect(evalScope = this@init)
+                }
+                val dl4: Deferred<TStateImpl<D>> = async {
+                    stateD.init.connect(evalScope = this@init)
+                }
+                val dl5: Deferred<TStateImpl<E>> = async {
+                    stateE.init.connect(evalScope = this@init)
+                }
+                zipStates(
+                    name,
+                    operatorName,
+                    dl1.await(),
+                    dl2.await(),
+                    dl3.await(),
+                    dl4.await(),
+                    dl5.await(),
+                ) { a, b, c, d, e ->
+                    NoScope.runInFrpScope { transform(a, b, c, d, e) }
+                }
+            }
+        }
+    )
+}
+
 /** Returns a [TState] by applying [transform] to the value held by the original [TState]. */
 @ExperimentalFrpApi
 fun <A, B> TState<A>.flatMap(transform: suspend FrpScope.(A) -> TState<B>): TState<B> {
diff --git a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/TStateImpl.kt b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/TStateImpl.kt
index 5cec05c8ef2d..c68b4c366776 100644
--- a/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/TStateImpl.kt
+++ b/packages/SystemUI/utils/kairos/src/com/android/systemui/kairos/internal/TStateImpl.kt
@@ -314,6 +314,28 @@ internal fun <A, B, C, D, Z> zipStates(
         @Suppress("UNCHECKED_CAST") transform(a as A, b as B, c as C, d as D)
     }
 
+internal fun <A, B, C, D, E, Z> zipStates(
+    name: String?,
+    operatorName: String,
+    l1: TStateImpl<A>,
+    l2: TStateImpl<B>,
+    l3: TStateImpl<C>,
+    l4: TStateImpl<D>,
+    l5: TStateImpl<E>,
+    transform: suspend EvalScope.(A, B, C, D, E) -> Z,
+): TStateImpl<Z> =
+    zipStates(null, operatorName, mapOf(0 to l1, 1 to l2, 2 to l3, 3 to l4, 4 to l5)).map(
+        name,
+        operatorName,
+    ) {
+        val a = it.getValue(0)
+        val b = it.getValue(1)
+        val c = it.getValue(2)
+        val d = it.getValue(3)
+        val e = it.getValue(4)
+        @Suppress("UNCHECKED_CAST") transform(a as A, b as B, c as C, d as D, e as E)
+    }
+
 internal fun <K : Any, A> zipStates(
     name: String?,
     operatorName: String,