perf: inline fields, shared keys, and concrete Members for objComp/mapWithKey (#758)

stephenamar-db · web-flow · commit ef04cb48135b · 2026-04-11T17:41:52.000-07:00
## Summary - **Shared allKeyNames/visibleKeyNames** (Val.Obj): Objects from the same `MemberList` expression now lazily share key-name arrays via a `_sourceMemberList` reference on `Val.Obj`. The first access computes and caches on the MemberList AST node; subsequent objects reuse the shared arrays. - **Inline field arrays for visitObjComp**: Apply the same `singleKey -> inlineKeys -> builder` tiered storage from `visitMemberList` to object comprehensions, avoiding `LinkedHashMap` for comprehensions producing 1-8 fields. - **Concrete ObjCompMember/MapWithKeyMember classes**: Replace anonymous `Val.Obj.Member` inner classes in `visitObjComp` and `std.mapWithKey` with concrete final classes for monomorphic JIT dispatch. - **Inline field arrays for std.mapWithKey**: Use `singleFieldKey`/`inlineFieldKeys` constructor paths for mapped objects with up to 8 keys, sharing the source object's `allKeyNames` array directly. - **Remove fieldLookups/fieldCacheHits counters**: These DebugStats counters in `Val.Obj.value()` were bloating the method past the JIT inlining threshold, causing a ~5% regression on bench.02. ## Benchmark results (vs pre-optimization baseline at 8c33fd9) | Benchmark | Baseline (ms/op) | After (ms/op) | Change | |---|---|---|---| | bench.02 | 76.368 | 73.352 | -3.9% | | bench.07 | 9.948 | 5.254 | **-47.2%** | | comparison | 6.941 | 4.204 | **-39.4%** | | large_string_template | 5.591 | 3.806 | **-31.9%** | | large_string_join | 1.476 | 1.048 | **-29.0%** | | base64 | 1.150 | 0.820 | **-28.7%** | | setDiff | 1.373 | 0.999 | **-27.2%** | | setUnion | 1.969 | 1.461 | **-25.8%** | | base64_byte_array | 2.925 | 2.198 | **-24.9%** | | realistic1 | 4.592 | 3.534 | **-23.0%** | | bench.04 | 0.486 | 0.375 | **-22.8%** | | foldl | 0.277 | 0.215 | **-22.4%** | | setInter | 1.072 | 0.838 | **-21.8%** | | substr | 0.212 | 0.174 | **-17.9%** | | realistic2 | 115.444 | 98.109 | **-15.0%** | All other benchmarks improved or within noise. No regressions. ## Test plan - [x] `./mill 'sjsonnet.jvm[3.3.7]'.test` — all 3 test suites pass - [x] `./mill __.checkFormat` — formatting clean - [x] `./mill bench.runRegressions` — no regressions
diff --git a/sjsonnet/src/sjsonnet/DebugStats.scala b/sjsonnet/src/sjsonnet/DebugStats.scala
@@ -29,8 +29,6 @@ final class DebugStats {
   var addSuperChainWalks: Long = 0
   var maxSuperChainDepth: Int = 0
   var valueCacheOverflows: Long = 0
-  var fieldLookups: Long = 0
-  var fieldCacheHits: Long = 0
 
   // -- Parse / import --
   var filesParsed: Long = 0
@@ -59,8 +57,6 @@ final class DebugStats {
     sb.append(formatLine("add_super_chain_walks", addSuperChainWalks))
     sb.append(formatLine("max_super_chain_depth", maxSuperChainDepth))
     sb.append(formatLine("value_cache_overflows", valueCacheOverflows))
-    sb.append(formatLine("field_lookups", fieldLookups))
-    sb.append(formatLine("field_cache_hits", fieldCacheHits))
     sb.append(formatLine("files_parsed", filesParsed))
     sb.append(formatLine("import_calls", importCalls))
     sb.append(formatLine("import_cache_hits", importCacheHits))
diff --git a/sjsonnet/src/sjsonnet/Evaluator.scala b/sjsonnet/src/sjsonnet/Evaluator.scala
@@ -1748,9 +1748,10 @@ class Evaluator(
         inlineFieldKeys = finalKeys,
         inlineFieldMembers = finalMembers
       )
-      // Cache sorted field order on MemberList (shared across all objects from same expression)
+      // Cache sorted field order on MemberList (shared across all objects from same expression).
+      // Only safe when all field names are fixed (not computed at runtime).
       var sortedOrder = e._cachedSortedOrder
-      if (sortedOrder == null && sup == null) {
+      if (sortedOrder == null && sup == null && e.allFieldNamesFixed) {
         sortedOrder = Materializer.computeSortedInlineOrder(finalKeys, finalMembers)
         e._cachedSortedOrder = sortedOrder
       }
@@ -1767,40 +1768,119 @@ class Evaluator(
         sup
       )
     }
+    // Only share key-name caches when all field names are fixed (compile-time constants).
+    // MemberLists with computed (Dyn) field names produce different keys per evaluation,
+    // so sharing allKeyNames/visibleKeyNames across objects would be incorrect.
+    if (sup == null && e.allFieldNamesFixed) factory.cachedObj._sourceMemberList = e
     factory.cachedObj
   }
 
   def visitObjComp(e: ObjBody.ObjComp, sup: Val.Obj)(implicit scope: ValScope): Val.Obj = {
     val binds = e.preLocals ++ e.postLocals
     val compScope: ValScope = scope // .clearSuper
-    val builder = new java.util.LinkedHashMap[String, Val.Obj.Member]
     val compScopes = visitComp(e.first :: e.rest, Array(compScope))
     if (debugStats != null) debugStats.objectCompIterations += compScopes.length
+
+    var builder: java.util.LinkedHashMap[String, Val.Obj.Member] = null
+    var singleKey: String = null
+    var singleMember: Val.Obj.Member = null
+    var inlineKeys: Array[String] = null
+    var inlineMembers: Array[Val.Obj.Member] = null
+    var fieldCount = 0
+    val maxInlineFields = 8
+
     for (s <- compScopes) {
       visitExpr(e.key)(s) match {
         case Val.Str(_, k) =>
-          val previousValue = builder.put(
-            k,
-            new Val.Obj.Member(e.plus, Visibility.Normal, deprecatedSkipAsserts = true) {
-              def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
-                checkStackDepth(e.value.pos, "object comprehension")
-                try {
-                  lazy val newScope: ValScope = s.extend(newBindings, self, sup)
-                  lazy val newBindings = visitBindings(binds, newScope)
-                  visitExpr(e.value)(newScope)
-                } finally decrementStackDepth()
+          val member = new ObjCompMember(
+            e.plus,
+            this,
+            binds,
+            s,
+            e.value
+          )
+          if (fieldCount == 0) {
+            singleKey = k
+            singleMember = member
+          } else if (fieldCount == 1) {
+            inlineKeys = new Array[String](math.min(compScopes.length, maxInlineFields))
+            inlineMembers = new Array[Val.Obj.Member](inlineKeys.length)
+            inlineKeys(0) = singleKey
+            inlineMembers(0) = singleMember
+            if (singleKey.equals(k))
+              Error.fail(s"Duplicate key $k in evaluated object comprehension.", e.pos)
+            inlineKeys(1) = k
+            inlineMembers(1) = member
+            singleKey = null
+            singleMember = null
+          } else if (fieldCount <= maxInlineFields && inlineKeys != null) {
+            var di = 0
+            while (di < fieldCount) {
+              if (inlineKeys(di).equals(k))
+                Error.fail(s"Duplicate key $k in evaluated object comprehension.", e.pos)
+              di += 1
+            }
+            if (fieldCount < inlineKeys.length) {
+              inlineKeys(fieldCount) = k
+              inlineMembers(fieldCount) = member
+            } else {
+              builder = Util.preSizedJavaLinkedHashMap[String, Val.Obj.Member](compScopes.length)
+              var mi = 0
+              while (mi < fieldCount) {
+                builder.put(inlineKeys(mi), inlineMembers(mi))
+                mi += 1
               }
+              inlineKeys = null
+              inlineMembers = null
+              builder.put(k, member)
             }
-          )
-          if (previousValue != null) {
-            Error.fail(s"Duplicate key $k in evaluated object comprehension.", e.pos)
+          } else {
+            val previousValue = builder.put(k, member)
+            if (previousValue != null)
+              Error.fail(s"Duplicate key $k in evaluated object comprehension.", e.pos)
           }
+          fieldCount += 1
         case Val.Null(_) => // do nothing
         case x           => fieldNameTypeError(x, e.pos)
       }
     }
     if (debugStats != null) debugStats.objectsCreated += 1
-    new Val.Obj(e.pos, builder, false, null, sup)
+    if (fieldCount == 1 && singleKey != null) {
+      new Val.Obj(
+        e.pos,
+        null,
+        false,
+        null,
+        sup,
+        singleFieldKey = singleKey,
+        singleFieldMember = singleMember
+      )
+    } else if (inlineKeys != null && fieldCount >= 2) {
+      val finalKeys =
+        if (fieldCount == inlineKeys.length) inlineKeys
+        else java.util.Arrays.copyOf(inlineKeys, fieldCount)
+      val finalMembers =
+        if (fieldCount == inlineMembers.length) inlineMembers
+        else java.util.Arrays.copyOf(inlineMembers, fieldCount)
+      new Val.Obj(
+        e.pos,
+        null,
+        false,
+        null,
+        sup,
+        inlineFieldKeys = finalKeys,
+        inlineFieldMembers = finalMembers
+      )
+    } else {
+      new Val.Obj(
+        e.pos,
+        if (builder != null) builder
+        else Util.preSizedJavaLinkedHashMap[String, Val.Obj.Member](0),
+        false,
+        null,
+        sup
+      )
+    }
   }
 
   @tailrec
@@ -2098,6 +2178,27 @@ private[sjsonnet] final class ObjectScopeFactory(
   }
 }
 
+/**
+ * Concrete [[Val.Obj.Member]] for object comprehension fields. Replaces the anonymous inner class
+ * in `visitObjComp`, giving the JIT a monomorphic type at `invoke` call sites.
+ */
+private[sjsonnet] final class ObjCompMember(
+    plus0: Boolean,
+    private val evaluator: Evaluator,
+    private val binds: Array[Expr.Bind],
+    private val compScope: ValScope,
+    private val valueExpr: Expr)
+    extends Val.Obj.Member(plus0, Visibility.Normal, deprecatedSkipAsserts = true) {
+  def invoke(self: Val.Obj, sup: Val.Obj, fs: FileScope, ev: EvalScope): Val = {
+    evaluator.checkStackDepth(valueExpr.pos, "object comprehension")
+    try {
+      lazy val newScope: ValScope = compScope.extend(newBindings, self, sup)
+      lazy val newBindings = evaluator.visitBindings(binds, newScope)
+      evaluator.visitExpr(valueExpr)(newScope)
+    } finally evaluator.decrementStackDepth()
+  }
+}
+
 /**
  * Concrete [[Val.Obj.Member]] for expression-body fields (no argSpec). Replaces the anonymous inner
  * class in `visitMemberList`, giving the JIT a single monomorphic type at `invoke` call sites for
diff --git a/sjsonnet/src/sjsonnet/Expr.scala b/sjsonnet/src/sjsonnet/Expr.scala
@@ -416,6 +416,18 @@ object Expr {
        */
       @volatile var _noSelfRef: java.lang.Boolean = null
 
+      /**
+       * Cached key-name arrays shared across all Val.Obj instances from this MemberList (when super
+       * == null). Computed from the first instance and reused by subsequent ones, avoiding per-object
+       * allKeyNames/visibleKeyNames allocation.
+       */
+      var _cachedAllKeyNames: Array[String] = null
+      var _cachedVisibleKeyNames: Array[String] = null
+
+      /** True if all fields have compile-time fixed names (no computed/Dyn field names). */
+      lazy val allFieldNamesFixed: Boolean =
+        fields.forall(_.fieldName.isInstanceOf[FieldName.Fixed])
+
       override def toString: String =
         s"MemberList($pos, ${arrStr(binds)}, ${arrStr(fields)}, ${arrStr(asserts)})"
     }
diff --git a/sjsonnet/src/sjsonnet/Val.scala b/sjsonnet/src/sjsonnet/Val.scala
@@ -686,6 +686,13 @@ object Val {
      */
     private[sjsonnet] var _skipFieldCache: Boolean = false
 
+    /**
+     * Reference to the source MemberList expression, set for objects with super == null. Enables
+     * lazy sharing of allKeyNames/visibleKeyNames: the first access computes and caches on the
+     * MemberList; subsequent objects from the same expression find the cached result.
+     */
+    private[sjsonnet] var _sourceMemberList: Expr.ObjBody.MemberList = null
+
     /**
      * Store a computed field value in the object's inline cache, preserving memoization semantics
      * when bypassing `value()` during direct iteration. This ensures that subsequent accesses via
@@ -987,79 +994,87 @@ object Val {
     }
 
     lazy val allKeyNames: Array[String] = {
-      if (inlineFieldKeys != null && `super` == null) inlineFieldKeys.clone()
+      val ml = _sourceMemberList
+      val cached = if (ml != null) ml._cachedAllKeyNames else null
+      if (cached != null) cached
       else {
-        val m = if (static || `super` != null) getAllKeys else getValue0
-        m.keySet().toArray(new Array[String](m.size()))
+        val result =
+          if (inlineFieldKeys != null && `super` == null) inlineFieldKeys.clone()
+          else {
+            val m = if (static || `super` != null) getAllKeys else getValue0
+            m.keySet().toArray(new Array[String](m.size()))
+          }
+        if (ml != null) ml._cachedAllKeyNames = result
+        result
       }
     }
 
     lazy val visibleKeyNames: Array[String] = {
-      if (static) {
-        allKeyNames
-      } else if (inlineFieldKeys != null && `super` == null) {
-        // Inline multi-field fast path: check if all visible (common case)
-        val keys = inlineFieldKeys
-        val members = inlineFieldMembers
-        val n = keys.length
-        var allVisible = true
-        var i = 0
-        while (allVisible && i < n) {
-          if (members(i).visibility == Visibility.Hidden) allVisible = false
-          i += 1
-        }
-        if (allVisible) keys.clone()
-        else {
+      val ml = _sourceMemberList
+      val cached = if (ml != null) ml._cachedVisibleKeyNames else null
+      if (cached != null) cached
+      else {
+        val result = if (static) {
+          allKeyNames
+        } else if (inlineFieldKeys != null && `super` == null) {
+          // Inline multi-field fast path: check if all visible (common case)
+          val keys = inlineFieldKeys
+          val members = inlineFieldMembers
+          val n = keys.length
+          var allVisible = true
+          var i = 0
+          while (allVisible && i < n) {
+            if (members(i).visibility == Visibility.Hidden) allVisible = false
+            i += 1
+          }
+          if (allVisible) keys.clone()
+          else {
+            val buf = new mutable.ArrayBuilder.ofRef[String]
+            buf.sizeHint(n)
+            var j = 0
+            while (j < n) {
+              if (members(j).visibility != Visibility.Hidden) buf += keys(j)
+              j += 1
+            }
+            buf.result()
+          }
+        } else {
           val buf = new mutable.ArrayBuilder.ofRef[String]
-          buf.sizeHint(n)
-          var j = 0
-          while (j < n) {
-            if (members(j).visibility != Visibility.Hidden) buf += keys(j)
-            j += 1
+          if (`super` == null) {
+            val v0 = getValue0
+            // This size hint is based on an optimistic assumption that most fields are visible,
+            // avoiding re-sizing or trimming the buffer in the common case:
+            buf.sizeHint(v0.size())
+            v0.forEach((k, m) => if (m.visibility != Visibility.Hidden) buf += k)
+          } else {
+            getAllKeys.forEach((k, b) => if (b == java.lang.Boolean.FALSE) buf += k)
           }
           buf.result()
         }
-      } else {
-        val buf = new mutable.ArrayBuilder.ofRef[String]
-        if (`super` == null) {
-          val v0 = getValue0
-          // This size hint is based on an optimistic assumption that most fields are visible,
-          // avoiding re-sizing or trimming the buffer in the common case:
-          buf.sizeHint(v0.size())
-          v0.forEach((k, m) => if (m.visibility != Visibility.Hidden) buf += k)
-        } else {
-          getAllKeys.forEach((k, b) => if (b == java.lang.Boolean.FALSE) buf += k)
-        }
-        buf.result()
+        if (ml != null) ml._cachedVisibleKeyNames = result
+        result
       }
     }
 
     def value(k: String, pos: Position, self: Obj = this)(implicit evaluator: EvalScope): Val = {
-      val ds = evaluator.debugStats
-      if (ds != null) ds.fieldLookups += 1
       if (static) {
         valueCache.get(k) match {
           case null => Error.fail("Field does not exist: " + k, pos)
-          case x    =>
-            if (ds != null) ds.fieldCacheHits += 1
-            x
+          case x    => x
         }
       } else {
         if ((self eq this) && excludedKeys != null && excludedKeys.contains(k)) {
           Error.fail("Field does not exist: " + k, pos)
         }
         val cacheKey: Any = if (self eq this) k else (k, self)
         if (ck1 != null && ck1 == cacheKey) {
-          if (ds != null) ds.fieldCacheHits += 1
           return cv1
         }
         if (ck2 != null && ck2 == cacheKey) {
-          if (ds != null) ds.fieldCacheHits += 1
           return cv2
         }
         val cachedValue = if (valueCache != null) valueCache.get(cacheKey) else null
         if (cachedValue != null) {
-          if (ds != null) ds.fieldCacheHits += 1
           cachedValue
         } else {
           valueRaw(k, self, pos, this, cacheKey) match {
diff --git a/sjsonnet/src/sjsonnet/stdlib/ObjectModule.scala b/sjsonnet/src/sjsonnet/stdlib/ObjectModule.scala
