Copy ImmutableLinkedHashMap without rehashing or tuple allocation

benchmarks/src/main/scala/play/api/libs/json/JsObjectBench.scala

@@ -46,9 +46,12 @@ class JsObjectBench {
   private var jsObject: JsObject      = _
   private var jsObjectClone: JsObject = _
   private var jsObjectCopy: JsObject  = _
+  private var fields: Seq[(String, JsValue)] = _
+  private val small: JsObject                = Json.obj("a" -> 1, "b" -> 2)
 
   @Setup def setup(): Unit = {
-    jsObject = JsObject(0.until(size).map(i => i.toString -> JsTrue))
+    fields = 0.until(size).map(i => i.toString -> JsTrue)
+    jsObject = JsObject(fields)
     jsObjectClone = JsObject(jsObject.underlying)
     jsObjectCopy = Json.parse(jsObject.toString).as[JsObject]
   }
@@ -57,4 +60,10 @@ class JsObjectBench {
   @Benchmark def oEqualsNe     = jsObject == jsObjectClone
   @Benchmark def oEqualsCopy   = jsObject == jsObjectCopy
   @Benchmark def oHashCodeWarm = jsObject.hashCode()
+
+  @Benchmark def oPlusNewKey      = jsObject + ("new" -> JsFalse)
+  @Benchmark def oPlusExistingKey = jsObject + ("0" -> JsFalse)
+  @Benchmark def oMinus           = jsObject - "0"
+  @Benchmark def oConcat          = jsObject ++ small
+  @Benchmark def oBuildByAdds     = fields.foldLeft(JsObject.empty)(_ + _)
 }

play-json/jvm/src/test/scala-3/play/api/libs/json/JsonMemoryFootprintSpec.scala

@@ -44,8 +44,21 @@ class JsonMemoryFootprintSpec extends AnyFreeSpec {
     "obj1 ++ obj0" in assertSize(obj("""{"1":true}""") ++ obj("""{}"""), 152, expectedJdk21 = Some(160))
 
     "obj1.value" in assertSize(obj("""{"1":true}""").tap(_.value), 152, expectedJdk21 = Some(160))
+
+    "built with + takes the same memory as built at once" in {
+      1.to(20).foreach { n =>
+        val fields   = 0.until(n).map(i => s"$i" -> JsTrue)
+        val stepwise = footprint(fields.foldLeft(JsObject.empty)(_ + _))
+        val oneGo    = footprint(JsObject(fields))
+        withClue(s"$n fields:")(assert(stepwise > 0 && stepwise === oneGo))
+      }
+    }
   }
 
+  // By-name: each layout is diffed against a second instance, so shared singletons don't count.
+  private def footprint(a: => JsValue): Long =
+    GraphLayout.parseInstance(a).subtract(GraphLayout.parseInstance(a)).totalSize()
+
   "malicious" - {
     // if we pack data into ~1KB of input, how much memory amplification can we achieve?
     def arr1KB(elem: String, targetSize: Int = 1000): String =

play-json/shared/src/main/scala-2.13+/play/api/libs/json/ImmutableLinkedHashMap.scala

@@ -27,7 +27,7 @@ private[json] class ImmutableLinkedHashMap[A, +B](underlying: JLinkedHashMap[A,
   }
 
   override def updated[V1 >: B](key: A, value: V1): Map[A, V1] = {
-    val c = shallowCopy[V1](size + 1)
+    val c = shallowCopy[V1](if (underlying.containsKey(key)) size else size + 1)
     c.put(key, value)
     new ImmutableLinkedHashMap(c)
   }
@@ -50,16 +50,28 @@ private[json] class ImmutableLinkedHashMap[A, +B](underlying: JLinkedHashMap[A,
   override def knownSize: Int = underlying.size()
   override def size: Int      = underlying.size()
 
+  // Sized for sizeHint entries, so the copy never rehashes, and copied from the entries
+  // directly rather than through the tuple iterator.
   private def shallowCopy[V1 >: B](sizeHint: Int = size): JLinkedHashMap[A, V1] = {
-    val c = new JLinkedHashMap[A, V1](sizeHint)
-    for ((k, v) <- this) c.put(k, v)
+    val c  = new JLinkedHashMap[A, V1](ImmutableLinkedHashMap.capacityFor(sizeHint))
+    val it = underlying.entrySet().iterator()
+
+    while (it.hasNext) {
+      val e = it.next()
+      c.put(e.getKey, e.getValue)
+    }
+
     c
   }
 }
 
 private[json] object ImmutableLinkedHashMap extends MapFactory[Map] {
   private object EmptyMap extends ImmutableLinkedHashMap[Any, Nothing](new JLinkedHashMap(0))
 
+  // Smallest initial capacity that holds `size` entries at the default 0.75 load factor
+  // (what LinkedHashMap.newLinkedHashMap computes on JDK 19+, not available at --release 17).
+  private[json] def capacityFor(size: Int): Int = math.ceil(size / 0.75).toInt
+
   override def empty[K, V]: Map[K, V] = EmptyMap.asInstanceOf[Map[K, V]]
 
   override def from[K, V](it: IterableOnce[(K, V)]): Map[K, V] = (newBuilder ++= it).result()

play-json/shared/src/main/scala-2.13-/play/api/libs/json/ImmutableLinkedHashMap.scala

@@ -27,7 +27,7 @@ private[json] class ImmutableLinkedHashMap[A, +B](underlying: JLinkedHashMap[A,
   override def get(key: A): Option[B] = Option(underlying.get(key))
 
   override def +[V1 >: B](kv: (A, V1)): Map[A, V1] = {
-    val c = shallowCopy[V1]()
+    val c = shallowCopy[V1](if (underlying.containsKey(kv._1)) size else size + 1)
     c.put(kv._1, kv._2)
     new ImmutableLinkedHashMap(c)
   }
@@ -59,16 +59,28 @@ private[json] class ImmutableLinkedHashMap[A, +B](underlying: JLinkedHashMap[A,
 
   override def empty = ImmutableLinkedHashMap.empty
 
-  private def shallowCopy[V1 >: B](): JLinkedHashMap[A, V1] = {
-    val c = new JLinkedHashMap[A, V1](underlying.size())
-    for ((k, v) <- this) c.put(k, v)
+  // Sized for sizeHint entries, so the copy never rehashes, and copied from the entries
+  // directly rather than through the tuple iterator.
+  private def shallowCopy[V1 >: B](sizeHint: Int = size): JLinkedHashMap[A, V1] = {
+    val c  = new JLinkedHashMap[A, V1](ImmutableLinkedHashMap.capacityFor(sizeHint))
+    val it = underlying.entrySet().iterator()
+
+    while (it.hasNext) {
+      val e = it.next()
+      c.put(e.getKey, e.getValue)
+    }
+
     c
   }
 }
 
 private[json] object ImmutableLinkedHashMap extends ImmutableMapFactory[ImmutableLinkedHashMap] {
   private object EmptyMap extends ImmutableLinkedHashMap[Any, Nothing](new JLinkedHashMap(0))
 
+  // Smallest initial capacity that holds `size` entries at the default 0.75 load factor
+  // (what LinkedHashMap.newLinkedHashMap computes on JDK 19+, not available at --release 17).
+  private[json] def capacityFor(size: Int): Int = math.ceil(size / 0.75).toInt
+
   override def empty[A, B]: ImmutableLinkedHashMap[A, B] = EmptyMap.asInstanceOf[ImmutableLinkedHashMap[A, B]]
 
   override def newBuilder[A, B]: mutable.Builder[(A, B), ImmutableLinkedHashMap[A, B]] = {

play-json/shared/src/test/scala/play/api/libs/json/JsObjectSpec.scala

@@ -250,6 +250,42 @@ class JsObjectSpec extends AnyWordSpec with Matchers {
     }
   }
 
+  "JsObject.+" should {
+    "append a new field and keep the position of an overwritten one" in {
+      val obj = Json.obj("a" -> 1, "b" -> 2, "c" -> 3)
+
+      Json.stringify(obj + ("d" -> JsNumber(4))).mustEqual("""{"a":1,"b":2,"c":3,"d":4}""")
+      Json.stringify(obj + ("b" -> JsNumber(20))).mustEqual("""{"a":1,"b":20,"c":3}""")
+      Json.stringify(obj).mustEqual("""{"a":1,"b":2,"c":3}""")
+    }
+
+    "build the same object as JsObject.apply, at every size" in {
+      // crosses the LinkedHashMap resize thresholds
+      0.to(40).foreach { n =>
+        val fields   = 0.until(n).map(i => s"f$i" -> JsNumber(i))
+        val stepwise = fields.foldLeft(Json.obj())(_ + _)
+        val oneGo    = JsObject(fields)
+
+        stepwise.mustEqual(oneGo)
+        stepwise.hashCode.mustEqual(oneGo.hashCode)
+        Json.stringify(stepwise).mustEqual(Json.stringify(oneGo))
+
+        val removed = fields.indices.filter(_ % 3 == 0).foldLeft(stepwise)((o, i) => o - s"f$i")
+        val kept    = fields.zipWithIndex.collect { case (f, i) if i % 3 != 0 => f }
+
+        Json.stringify(removed).mustEqual(Json.stringify(JsObject(kept)))
+      }
+    }
+  }
+
+  "JsObject.++" should {
+    "keep the left fields in order, overwrite in place and append new ones" in {
+      val merged = Json.obj("a" -> 1, "b" -> 2) ++ Json.obj("b" -> 20, "c" -> 3)
+
+      Json.stringify(merged).mustEqual("""{"a":1,"b":20,"c":3}""")
+    }
+  }
+
   "JsObject" should {
     // see https://github.com/playframework/play-json/issues/390
     "accept null fields when calling Json.obj" in {