From 22ff3844495608b00216dff5ec1c03210fc3cb55 Mon Sep 17 00:00:00 2001
From: Laimonas Turauskas <laimiux@gmail.com>
Date: Thu, 7 Dec 2023 15:24:49 -0800
Subject: [PATCH] Handle duplicate child key errors more gracefully.

---
 .../formula/internal/ChildrenManager.kt       | 63 ++++++++++++++++---
 1 file changed, 55 insertions(+), 8 deletions(-)

diff --git a/formula/src/main/java/com/instacart/formula/internal/ChildrenManager.kt b/formula/src/main/java/com/instacart/formula/internal/ChildrenManager.kt
index ee9d20e4..551c5e5c 100644
--- a/formula/src/main/java/com/instacart/formula/internal/ChildrenManager.kt
+++ b/formula/src/main/java/com/instacart/formula/internal/ChildrenManager.kt
@@ -2,6 +2,7 @@ package com.instacart.formula.internal
 
 import com.instacart.formula.IFormula
 import com.instacart.formula.Inspector
+import java.lang.IllegalStateException
 
 /**
  * Keeps track of child formula managers.
@@ -11,6 +12,7 @@ internal class ChildrenManager(
     private val inspector: Inspector?,
 ) {
     private var children: SingleRequestMap<Any, FormulaManager<*, *>>? = null
+    private var indexes: MutableMap<Any, Int>? = null
     private var pendingRemoval: MutableList<FormulaManager<*, *>>? = null
 
     /**
@@ -19,6 +21,8 @@ internal class ChildrenManager(
      * in post evaluation, which will call [terminateChildren] function.
      */
     fun prepareForPostEvaluation() {
+        indexes?.clear()
+
         children?.clearUnrequested {
             pendingRemoval = pendingRemoval ?: mutableListOf()
             it.markAsTerminated()
@@ -51,6 +55,29 @@ internal class ChildrenManager(
         formula: IFormula<ChildInput, ChildOutput>,
         input: ChildInput,
     ): FormulaManager<ChildInput, ChildOutput> {
+        val childHolder = childFormulaHolder(key, formula, input)
+        return if (childHolder.requested) {
+            if (key is IndexedKey) {
+                // This should never happen, but added as safety
+                throw IllegalStateException("Key already indexed (and still duplicate).")
+            }
+
+            // TODO: How to handle duplicate keys
+            val index = nextIndex(key)
+            val indexedKey = IndexedKey(key, index)
+            findOrInitChild(indexedKey, formula, input)
+        } else {
+            childHolder.requestAccess {
+                "There already is a child with same key: $key. Override [Formula.key] function."
+            }
+        }
+    }
+
+    private fun <ChildInput, ChildOutput>  childFormulaHolder(
+        key: Any,
+        formula: IFormula<ChildInput, ChildOutput>,
+        input: ChildInput,
+    ): SingleRequestHolder<FormulaManager<ChildInput, ChildOutput>> {
         @Suppress("UNCHECKED_CAST")
         val children = children ?: run {
             val initialized: SingleRequestMap<Any, FormulaManager<*, *>> = LinkedHashMap()
@@ -58,13 +85,33 @@ internal class ChildrenManager(
             initialized
         }
 
-        return children
-            .findOrInit(key) {
-                val implementation = formula.implementation()
-                FormulaManagerImpl(delegate, implementation, input, loggingType = formula::class, inspector = inspector)
-            }
-            .requestAccess {
-                "There already is a child with same key: $key. Override [Formula.key] function."
-            } as FormulaManager<ChildInput, ChildOutput>
+        val childFormulaHolder = children.findOrInit(key) {
+            val implementation = formula.implementation()
+            FormulaManagerImpl(
+                delegate,
+                implementation,
+                input,
+                loggingType = formula::class,
+                inspector = inspector
+            )
+        }
+        @Suppress("UNCHECKED_CAST")
+        return childFormulaHolder as SingleRequestHolder<FormulaManager<ChildInput, ChildOutput>>
+    }
+
+    /**
+     * Function which returns next index for a given key. It will
+     * mutate the [indexes] map.
+     */
+    private fun nextIndex(key: Any): Int {
+        val indexes = indexes ?: run {
+            val initialized = mutableMapOf<Any, Int>()
+            this.indexes = initialized
+            initialized
+        }
+
+        val index = indexes.getOrElse(key) { 0 } + 1
+        indexes[key] = index
+        return index
     }
 }
\ No newline at end of file