better helper function for mutable maps with default puts

Author: aSemy

modules/kxs-ts-gen-core/src/commonMain/kotlin/dev.adamko.kxstsgen/util/mapWithDefaultDelegate.kt

@@ -1,5 +1,6 @@
 package dev.adamko.kxstsgen.util
 
+import kotlin.jvm.JvmName
 import kotlin.properties.ReadWriteProperty
 import kotlin.reflect.KProperty
 
@@ -9,14 +10,12 @@ class MutableMapWithDefaultPut<K, V>(
   private val defaultValue: (key: K) -> V,
 ) : ReadWriteProperty<Any?, MutableMap<K, V>> {
 
-  private var map: MutableMap<K, V> = with(initial.toMutableMap()) {
-    withDefault { key -> getOrPut(key) { defaultValue(key) } }
-  }
+  private var map: MutableMap<K, V> = initial.toMutableMap().withDefaultPut(defaultValue)
 
   override fun getValue(thisRef: Any?, property: KProperty<*>): MutableMap<K, V> = map
 
   override fun setValue(thisRef: Any?, property: KProperty<*>, value: MutableMap<K, V>) {
-    this.map = value
+    this.map = value.withDefaultPut(defaultValue)
   }
 }
 
@@ -33,6 +32,20 @@ class MapWithDefaultPut<K, V>(
   override fun getValue(thisRef: Any?, property: KProperty<*>): Map<K, V> = map
 
   override fun setValue(thisRef: Any?, property: KProperty<*>, value: Map<K, V>) {
-    this.map = value
+    this.map = value.toMutableMap().withDefaultPut(defaultValue)
   }
 }
+
+
+@JvmName("mapWithDefaultPut")
+fun <K, V> Map<K, V>.withDefaultPut(defaultValue: (key: K) -> V): Map<K, V> =
+  with(this.toMutableMap()) {
+    withDefault { key -> getOrPut(key) { defaultValue(key) } }
+  }
+
+
+@JvmName("mutableMapWithDefaultPut")
+fun <K, V> MutableMap<K, V>.withDefaultPut(defaultValue: (key: K) -> V): MutableMap<K, V> =
+  with(this) {
+    withDefault { key -> getOrPut(key) { defaultValue(key) } }
+  }