1 /*
<lambda>null2 * Copyright (C) 2024 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 package com.android.systemui.kairos
18
19 import com.android.systemui.kairos.internal.DerivedMapCheap
20 import com.android.systemui.kairos.internal.Init
21 import com.android.systemui.kairos.internal.InitScope
22 import com.android.systemui.kairos.internal.Network
23 import com.android.systemui.kairos.internal.NoScope
24 import com.android.systemui.kairos.internal.Schedulable
25 import com.android.systemui.kairos.internal.TFlowImpl
26 import com.android.systemui.kairos.internal.TStateImpl
27 import com.android.systemui.kairos.internal.TStateSource
28 import com.android.systemui.kairos.internal.activated
29 import com.android.systemui.kairos.internal.cached
30 import com.android.systemui.kairos.internal.constInit
31 import com.android.systemui.kairos.internal.constS
32 import com.android.systemui.kairos.internal.filterNode
33 import com.android.systemui.kairos.internal.flatMap
34 import com.android.systemui.kairos.internal.init
35 import com.android.systemui.kairos.internal.map
36 import com.android.systemui.kairos.internal.mapCheap
37 import com.android.systemui.kairos.internal.mapImpl
38 import com.android.systemui.kairos.internal.util.hashString
39 import com.android.systemui.kairos.internal.zipStates
40 import kotlin.reflect.KProperty
41 import kotlinx.coroutines.CompletableDeferred
42 import kotlinx.coroutines.Deferred
43 import kotlinx.coroutines.async
44 import kotlinx.coroutines.coroutineScope
45
46 /**
47 * A time-varying value with discrete changes. Essentially, a combination of a [Transactional] that
48 * holds a value, and a [TFlow] that emits when the value changes.
49 */
50 @ExperimentalFrpApi sealed class TState<out A>
51
52 /** A [TState] that never changes. */
53 @ExperimentalFrpApi
54 fun <A> tStateOf(value: A): TState<A> {
55 val operatorName = "tStateOf"
56 val name = "$operatorName($value)"
57 return TStateInit(constInit(name, constS(name, operatorName, value)))
58 }
59
60 /** TODO */
<lambda>null61 @ExperimentalFrpApi fun <A> Lazy<TState<A>>.defer(): TState<A> = deferInline { value }
62
63 /** TODO */
64 @ExperimentalFrpApi
<lambda>null65 fun <A> FrpDeferredValue<TState<A>>.defer(): TState<A> = deferInline { unwrapped.await() }
66
67 /** TODO */
68 @ExperimentalFrpApi
deferTStatenull69 fun <A> deferTState(block: suspend FrpScope.() -> TState<A>): TState<A> = deferInline {
70 NoScope.runInFrpScope(block)
71 }
72
73 /**
74 * Returns a [TState] containing the results of applying [transform] to the value held by the
75 * original [TState].
76 */
77 @ExperimentalFrpApi
mapnull78 fun <A, B> TState<A>.map(transform: suspend FrpScope.(A) -> B): TState<B> {
79 val operatorName = "map"
80 val name = operatorName
81 return TStateInit(
82 init(name) {
83 init.connect(evalScope = this).map(name, operatorName) {
84 NoScope.runInFrpScope { transform(it) }
85 }
86 }
87 )
88 }
89
90 /**
91 * Returns a [TState] that transforms the value held inside this [TState] by applying it to the
92 * [transform].
93 *
94 * Note that unlike [map], the result is not cached. This means that not only should [transform] be
95 * fast and pure, it should be *monomorphic* (1-to-1). Failure to do this means that [stateChanges]
96 * for the returned [TState] will operate unexpectedly, emitting at rates that do not reflect an
97 * observable change to the returned [TState].
98 */
99 @ExperimentalFrpApi
mapCheapUnsafenull100 fun <A, B> TState<A>.mapCheapUnsafe(transform: suspend FrpScope.(A) -> B): TState<B> {
101 val operatorName = "map"
102 val name = operatorName
103 return TStateInit(
104 init(name) {
105 init.connect(evalScope = this).mapCheap(name, operatorName) {
106 NoScope.runInFrpScope { transform(it) }
107 }
108 }
109 )
110 }
111
112 /**
113 * Returns a [TState] by combining the values held inside the given [TState]s by applying them to
114 * the given function [transform].
115 */
116 @ExperimentalFrpApi
combineWithnull117 fun <A, B, C> TState<A>.combineWith(
118 other: TState<B>,
119 transform: suspend FrpScope.(A, B) -> C,
120 ): TState<C> = combine(this, other, transform)
121
122 /**
123 * Splits a [TState] of pairs into a pair of [TFlows][TState], where each returned [TState] holds
124 * hald of the original.
125 *
126 * Shorthand for:
127 * ```kotlin
128 * val lefts = map { it.first }
129 * val rights = map { it.second }
130 * return Pair(lefts, rights)
131 * ```
132 */
133 @ExperimentalFrpApi
134 fun <A, B> TState<Pair<A, B>>.unzip(): Pair<TState<A>, TState<B>> {
135 val left = map { it.first }
136 val right = map { it.second }
137 return left to right
138 }
139
140 /**
141 * Returns a [TState] by combining the values held inside the given [TStates][TState] into a [List].
142 *
143 * @see TState.combineWith
144 */
145 @ExperimentalFrpApi
combinenull146 fun <A> Iterable<TState<A>>.combine(): TState<List<A>> {
147 val operatorName = "combine"
148 val name = operatorName
149 return TStateInit(
150 init(name) {
151 zipStates(name, operatorName, states = map { it.init.connect(evalScope = this) })
152 }
153 )
154 }
155
156 /**
157 * Returns a [TState] by combining the values held inside the given [TStates][TState] into a [Map].
158 *
159 * @see TState.combineWith
160 */
161 @ExperimentalFrpApi
combinenull162 fun <K : Any, A> Map<K, TState<A>>.combine(): TState<Map<K, A>> {
163 val operatorName = "combine"
164 val name = operatorName
165 return TStateInit(
166 init(name) {
167 zipStates(
168 name,
169 operatorName,
170 states = mapValues { it.value.init.connect(evalScope = this) },
171 )
172 }
173 )
174 }
175
176 /**
177 * Returns a [TState] whose value is generated with [transform] by combining the current values of
178 * each given [TState].
179 *
180 * @see TState.combineWith
181 */
182 @ExperimentalFrpApi
combinenull183 fun <A, B> Iterable<TState<A>>.combine(transform: suspend FrpScope.(List<A>) -> B): TState<B> =
184 combine().map(transform)
185
186 /**
187 * Returns a [TState] by combining the values held inside the given [TState]s into a [List].
188 *
189 * @see TState.combineWith
190 */
191 @ExperimentalFrpApi
192 fun <A> combine(vararg states: TState<A>): TState<List<A>> = states.asIterable().combine()
193
194 /**
195 * Returns a [TState] whose value is generated with [transform] by combining the current values of
196 * each given [TState].
197 *
198 * @see TState.combineWith
199 */
200 @ExperimentalFrpApi
201 fun <A, B> combine(
202 vararg states: TState<A>,
203 transform: suspend FrpScope.(List<A>) -> B,
204 ): TState<B> = states.asIterable().combine(transform)
205
206 /**
207 * Returns a [TState] whose value is generated with [transform] by combining the current values of
208 * each given [TState].
209 *
210 * @see TState.combineWith
211 */
212 @ExperimentalFrpApi
213 fun <A, B, Z> combine(
214 stateA: TState<A>,
215 stateB: TState<B>,
216 transform: suspend FrpScope.(A, B) -> Z,
217 ): TState<Z> {
218 val operatorName = "combine"
219 val name = operatorName
220 return TStateInit(
221 init(name) {
222 coroutineScope {
223 val dl1: Deferred<TStateImpl<A>> = async {
224 stateA.init.connect(evalScope = this@init)
225 }
226 val dl2: Deferred<TStateImpl<B>> = async {
227 stateB.init.connect(evalScope = this@init)
228 }
229 zipStates(name, operatorName, dl1.await(), dl2.await()) { a, b ->
230 NoScope.runInFrpScope { transform(a, b) }
231 }
232 }
233 }
234 )
235 }
236
237 /**
238 * Returns a [TState] whose value is generated with [transform] by combining the current values of
239 * each given [TState].
240 *
241 * @see TState.combineWith
242 */
243 @ExperimentalFrpApi
combinenull244 fun <A, B, C, Z> combine(
245 stateA: TState<A>,
246 stateB: TState<B>,
247 stateC: TState<C>,
248 transform: suspend FrpScope.(A, B, C) -> Z,
249 ): TState<Z> {
250 val operatorName = "combine"
251 val name = operatorName
252 return TStateInit(
253 init(name) {
254 coroutineScope {
255 val dl1: Deferred<TStateImpl<A>> = async {
256 stateA.init.connect(evalScope = this@init)
257 }
258 val dl2: Deferred<TStateImpl<B>> = async {
259 stateB.init.connect(evalScope = this@init)
260 }
261 val dl3: Deferred<TStateImpl<C>> = async {
262 stateC.init.connect(evalScope = this@init)
263 }
264 zipStates(name, operatorName, dl1.await(), dl2.await(), dl3.await()) { a, b, c ->
265 NoScope.runInFrpScope { transform(a, b, c) }
266 }
267 }
268 }
269 )
270 }
271
272 /**
273 * Returns a [TState] whose value is generated with [transform] by combining the current values of
274 * each given [TState].
275 *
276 * @see TState.combineWith
277 */
278 @ExperimentalFrpApi
combinenull279 fun <A, B, C, D, Z> combine(
280 stateA: TState<A>,
281 stateB: TState<B>,
282 stateC: TState<C>,
283 stateD: TState<D>,
284 transform: suspend FrpScope.(A, B, C, D) -> Z,
285 ): TState<Z> {
286 val operatorName = "combine"
287 val name = operatorName
288 return TStateInit(
289 init(name) {
290 coroutineScope {
291 val dl1: Deferred<TStateImpl<A>> = async {
292 stateA.init.connect(evalScope = this@init)
293 }
294 val dl2: Deferred<TStateImpl<B>> = async {
295 stateB.init.connect(evalScope = this@init)
296 }
297 val dl3: Deferred<TStateImpl<C>> = async {
298 stateC.init.connect(evalScope = this@init)
299 }
300 val dl4: Deferred<TStateImpl<D>> = async {
301 stateD.init.connect(evalScope = this@init)
302 }
303 zipStates(name, operatorName, dl1.await(), dl2.await(), dl3.await(), dl4.await()) {
304 a,
305 b,
306 c,
307 d ->
308 NoScope.runInFrpScope { transform(a, b, c, d) }
309 }
310 }
311 }
312 )
313 }
314
315 /**
316 * Returns a [TState] whose value is generated with [transform] by combining the current values of
317 * each given [TState].
318 *
319 * @see TState.combineWith
320 */
321 @ExperimentalFrpApi
combinenull322 fun <A, B, C, D, E, Z> combine(
323 stateA: TState<A>,
324 stateB: TState<B>,
325 stateC: TState<C>,
326 stateD: TState<D>,
327 stateE: TState<E>,
328 transform: suspend FrpScope.(A, B, C, D, E) -> Z,
329 ): TState<Z> {
330 val operatorName = "combine"
331 val name = operatorName
332 return TStateInit(
333 init(name) {
334 coroutineScope {
335 val dl1: Deferred<TStateImpl<A>> = async {
336 stateA.init.connect(evalScope = this@init)
337 }
338 val dl2: Deferred<TStateImpl<B>> = async {
339 stateB.init.connect(evalScope = this@init)
340 }
341 val dl3: Deferred<TStateImpl<C>> = async {
342 stateC.init.connect(evalScope = this@init)
343 }
344 val dl4: Deferred<TStateImpl<D>> = async {
345 stateD.init.connect(evalScope = this@init)
346 }
347 val dl5: Deferred<TStateImpl<E>> = async {
348 stateE.init.connect(evalScope = this@init)
349 }
350 zipStates(
351 name,
352 operatorName,
353 dl1.await(),
354 dl2.await(),
355 dl3.await(),
356 dl4.await(),
357 dl5.await(),
358 ) { a, b, c, d, e ->
359 NoScope.runInFrpScope { transform(a, b, c, d, e) }
360 }
361 }
362 }
363 )
364 }
365
366 /** Returns a [TState] by applying [transform] to the value held by the original [TState]. */
367 @ExperimentalFrpApi
flatMapnull368 fun <A, B> TState<A>.flatMap(transform: suspend FrpScope.(A) -> TState<B>): TState<B> {
369 val operatorName = "flatMap"
370 val name = operatorName
371 return TStateInit(
372 init(name) {
373 init.connect(this).flatMap(name, operatorName) { a ->
374 NoScope.runInFrpScope { transform(a) }.init.connect(this)
375 }
376 }
377 )
378 }
379
380 /** Shorthand for `flatMap { it }` */
flattennull381 @ExperimentalFrpApi fun <A> TState<TState<A>>.flatten() = flatMap { it }
382
383 /**
384 * Returns a [TStateSelector] that can be used to efficiently check if the input [TState] is
385 * currently holding a specific value.
386 *
387 * An example:
388 * ```
389 * val lInt: TState<Int> = ...
390 * val intSelector: TStateSelector<Int> = lInt.selector()
391 * // Tracks if lInt is holding 1
392 * val isOne: TState<Boolean> = intSelector.whenSelected(1)
393 * ```
394 *
395 * This is semantically equivalent to `val isOne = lInt.map { i -> i == 1 }`, but is significantly
396 * more efficient; specifically, using [TState.map] in this way incurs a `O(n)` performance hit,
397 * where `n` is the number of different [TState.map] operations used to track a specific value.
398 * [selector] internally uses a [HashMap] to lookup the appropriate downstream [TState] to update,
399 * and so operates in `O(1)`.
400 *
401 * Note that the result [TStateSelector] should be cached and re-used to gain the performance
402 * benefit.
403 *
404 * @see groupByKey
405 */
406 @ExperimentalFrpApi
selectornull407 fun <A> TState<A>.selector(numDistinctValues: Int? = null): TStateSelector<A> =
408 TStateSelector(
409 this,
410 stateChanges
411 .map { new -> mapOf(new to true, sampleDeferred().get() to false) }
412 .groupByKey(numDistinctValues),
413 )
414
415 /**
416 * Tracks the currently selected value of type [A] from an upstream [TState].
417 *
418 * @see selector
419 */
420 @ExperimentalFrpApi
421 class TStateSelector<in A>
422 internal constructor(
423 private val upstream: TState<A>,
424 private val groupedChanges: GroupedTFlow<A, Boolean>,
425 ) {
426 /**
427 * Returns a [TState] that tracks whether the upstream [TState] is currently holding the given
428 * [value].
429 *
430 * @see selector
431 */
432 @ExperimentalFrpApi
whenSelectednull433 fun whenSelected(value: A): TState<Boolean> {
434 val operatorName = "TStateSelector#whenSelected"
435 val name = "$operatorName[$value]"
436 return TStateInit(
437 init(name) {
438 DerivedMapCheap(
439 name,
440 operatorName,
441 upstream = upstream.init.connect(evalScope = this),
442 changes = groupedChanges.impl.eventsForKey(value),
443 ) {
444 it == value
445 }
446 }
447 )
448 }
449
getnull450 @ExperimentalFrpApi operator fun get(value: A): TState<Boolean> = whenSelected(value)
451 }
452
453 /** TODO */
454 @ExperimentalFrpApi
455 class MutableTState<T>
456 internal constructor(internal val network: Network, initialValue: Deferred<T>) : TState<T>() {
457
458 private val input: CoalescingMutableTFlow<Deferred<T>, Deferred<T>?> =
459 CoalescingMutableTFlow(
460 coalesce = { _, new -> new },
461 network = network,
462 getInitialValue = { null },
463 )
464
465 internal val tState = run {
466 val changes = input.impl
467 val name = null
468 val operatorName = "MutableTState"
469 lateinit var state: TStateSource<T>
470 val calm: TFlowImpl<T> =
471 filterNode({ mapImpl(upstream = { changes.activated() }) { it!!.await() } }) { new ->
472 new != state.getCurrentWithEpoch(evalScope = this).first
473 }
474 .cached()
475 state = TStateSource(name, operatorName, initialValue, calm)
476 @Suppress("DeferredResultUnused")
477 network.transaction {
478 calm.activate(evalScope = this, downstream = Schedulable.S(state))?.let {
479 (connection, needsEval) ->
480 state.upstreamConnection = connection
481 if (needsEval) {
482 schedule(state)
483 }
484 }
485 }
486 TStateInit(constInit(name, state))
487 }
488
489 /** TODO */
490 @ExperimentalFrpApi fun setValue(value: T) = input.emit(CompletableDeferred(value))
491
492 @ExperimentalFrpApi
493 fun setValueDeferred(value: FrpDeferredValue<T>) = input.emit(value.unwrapped)
494 }
495
496 /** A forward-reference to a [TState], allowing for recursive definitions. */
497 @ExperimentalFrpApi
498 class TStateLoop<A> : TState<A>() {
499
500 private val name: String? = null
501
502 private val deferred = CompletableDeferred<TState<A>>()
503
504 internal val init: Init<TStateImpl<A>> =
<lambda>null505 init(name) { deferred.await().init.connect(evalScope = this) }
506
507 /** The [TState] this [TStateLoop] will forward to. */
508 @ExperimentalFrpApi
509 var loopback: TState<A>? = null
510 set(value) {
<lambda>null511 value?.let {
512 check(deferred.complete(value)) { "TStateLoop.loopback has already been set." }
513 field = value
514 }
515 }
516
517 @ExperimentalFrpApi
getValuenull518 operator fun getValue(thisRef: Any?, property: KProperty<*>): TState<A> = this
519
520 @ExperimentalFrpApi
521 operator fun setValue(thisRef: Any?, property: KProperty<*>, value: TState<A>) {
522 loopback = value
523 }
524
toStringnull525 override fun toString(): String = "${this::class.simpleName}@$hashString"
526 }
527
528 internal class TStateInit<A> internal constructor(internal val init: Init<TStateImpl<A>>) :
529 TState<A>() {
530 override fun toString(): String = "${this::class.simpleName}@$hashString"
531 }
532
533 internal val <A> TState<A>.init: Init<TStateImpl<A>>
534 get() =
535 when (this) {
536 is TStateInit -> init
537 is TStateLoop -> init
538 is MutableTState -> tState.init
539 }
540
deferInlinenull541 private inline fun <A> deferInline(
542 crossinline block: suspend InitScope.() -> TState<A>
543 ): TState<A> = TStateInit(init(name = null) { block().init.connect(evalScope = this) })
544