perf: reduce TOML and repeated join allocations

sjsonnet/src/sjsonnet/TomlRenderer.scala

@@ -1,6 +1,6 @@
 package sjsonnet
 
-import upickle.core.{ArrVisitor, CharBuilder, ObjVisitor, SimpleVisitor, Visitor}
+import upickle.core.{ArrVisitor, ObjVisitor, SimpleVisitor, Visitor}
 
 import java.io.StringWriter
 
@@ -20,7 +20,7 @@ class TomlRenderer(
     override def visitString(s: CharSequence, index: Int): StringWriter = {
       if (s == null) visitNull(index)
       else {
-        out.write(TomlRenderer.escapeKey(s.toString))
+        TomlRenderer.writeEscapedKey(out, s)
         out
       }
     }
@@ -29,8 +29,7 @@ class TomlRenderer(
   private var depth = 0
 
   private def flush = {
-    if (depth == 0) out.write("\n")
-    out.flush()
+    if (depth == 0) out.write('\n')
     out
   }
 
@@ -50,15 +49,7 @@ class TomlRenderer(
     if (s == null) {
       visitNull(index)
     } else {
-      val charBuilder = new CharBuilder()
-      upickle.core.RenderUtils.escapeChar(
-        null,
-        charBuilder,
-        s,
-        escapeUnicode = true,
-        wrapQuotes = true
-      )
-      out.write(charBuilder.makeString())
+      BaseRenderer.escape(out, s, unicode = true)
       flush
     }
   }
@@ -83,9 +74,13 @@ class TomlRenderer(
       private var addComma = false
 
       depth += 1
-      out.write("[" + separator)
+      out.write('[')
+      out.write(separator)
       def subVisitor: Visitor[StringWriter, StringWriter] = {
-        if (addComma) out.write("," + separator)
+        if (addComma) {
+          out.write(',')
+          out.write(separator)
+        }
         out.write(newElementIndent)
         TomlRenderer.this
       }
@@ -131,11 +126,35 @@ class TomlRenderer(
 }
 
 object TomlRenderer {
-  private val bareAllowed = Platform.getPatternFromCache("[A-Za-z0-9_-]+")
-  def escapeKey(key: String): String = if (bareAllowed.matcher(key).matches()) key
+  @inline private def isBareKeyChar(c: Char): Boolean =
+    (c >= 'A' && c <= 'Z') ||
+    (c >= 'a' && c <= 'z') ||
+    (c >= '0' && c <= '9') ||
+    c == '_' ||
+    c == '-'
+
+  private def isBareKey(key: CharSequence): Boolean = {
+    val len = key.length
+    if (len == 0) false
+    else {
+      var i = 0
+      while (i < len) {
+        if (!isBareKeyChar(key.charAt(i))) return false
+        i += 1
+      }
+      true
+    }
+  }
+
+  def writeEscapedKey(out: StringWriter, key: CharSequence): Unit = {
+    if (isBareKey(key)) out.write(key.toString)
+    else BaseRenderer.escape(out, key, unicode = true)
+  }
+
+  def escapeKey(key: String): String = if (isBareKey(key)) key
   else {
     val out = new StringWriter()
-    BaseRenderer.escape(out, key, unicode = true)
+    writeEscapedKey(out, key)
     out.toString
   }
 }

sjsonnet/src/sjsonnet/Val.scala

@@ -582,6 +582,8 @@ object Val {
       if (!isConcatView && !_reversed && (arr ne null)) arr.asInstanceOf[Array[Eval]]
       else null
 
+    private[sjsonnet] def constantEval: Eval = null
+
     /**
      * If both this and other are ConcatViews sharing the same left array, return the shared prefix
      * length. Otherwise return 0. Used by compare/equal to skip identical prefix elements entirely,
@@ -835,12 +837,46 @@ object Val {
     }
   }
 
-  private final class RepeatedArr(pos0: Position, private[this] val source: Arr, count: Int)
+  private final class ConstArr(pos0: Position, size: Int, private val elem: Eval)
+      extends Arr(pos0, null) {
+    _length = size
+
+    @inline private def checkIndex(i: Int): Unit =
+      if (i < 0 || i >= _length) throw new ArrayIndexOutOfBoundsException(i)
+
+    override def value(i: Int): Val = {
+      checkIndex(i)
+      elem.value
+    }
+
+    override def eval(i: Int): Eval = {
+      checkIndex(i)
+      elem
+    }
+
+    override private[sjsonnet] def constantEval: Eval = elem
+
+    override def asLazyArray: Array[Eval] = {
+      val materialized = arr
+      if (materialized != null) materialized.asInstanceOf[Array[Eval]]
+      else {
+        val result = new Array[Eval](_length)
+        java.util.Arrays.fill(result.asInstanceOf[Array[AnyRef]], elem.asInstanceOf[AnyRef])
+        arr = result
+        result
+      }
+    }
+
+    override def reversed(newPos: Position): Arr =
+      new ConstArr(newPos, _length, elem)
+  }
+
+  private final class RepeatedArr(pos0: Position, private val source: Arr, count: Int)
       extends Arr(pos0, null) {
     // Keep std.repeat(array, n) as an indexed view. The common consumers either index a subset or
     // materialize later anyway; eagerly copying here multiplies thunks and stresses young-gen GC.
-    private[this] val sourceLen = source.length
-    private[this] val totalLen = sourceLen.toLong * count.toLong
+    private val sourceLen = source.length
+    private val totalLen = sourceLen.toLong * count.toLong
     if (totalLen > Int.MaxValue) throw new IllegalArgumentException("array too large")
     _length = totalLen.toInt
 
@@ -850,6 +886,9 @@ object Val {
     override def eval(i: Int): Eval =
       source.eval(i % sourceLen)
 
+    override private[sjsonnet] def constantEval: Eval =
+      if (sourceLen == 1) source.eval(0) else source.constantEval
+
     override def asLazyArray: Array[Eval] = {
       val sourceArr = source.asLazyArray
       val sourceLen = this.sourceLen
@@ -1278,6 +1317,14 @@ object Val {
   object Arr {
     def apply(pos: Position, arr: Array[? <: Eval]): Arr = new Arr(pos, arr)
 
+    def constant(pos: Position, size: Int, elem: Eval): Arr =
+      if (size == 0) Arr(pos, EMPTY_EVAL_ARRAY)
+      else if (size < LAZY_VIEW_THRESHOLD) {
+        val result = new Array[Eval](size)
+        java.util.Arrays.fill(result.asInstanceOf[Array[AnyRef]], elem.asInstanceOf[AnyRef])
+        Arr(pos, result)
+      } else new ConstArr(pos, size, elem)
+
     def repeated(pos: Position, source: Arr, count: Int): Arr =
       if (count == 0 || source.length == 0) Arr(pos, EMPTY_EVAL_ARRAY)
       else new RepeatedArr(pos, source, count)

sjsonnet/src/sjsonnet/stdlib/ArrayModule.scala

@@ -919,10 +919,7 @@ object ArrayModule extends AbstractFunctionModule {
       val body = func.bodyExpr
       if (func.params.names.length == 1 && body != null && body.isInstanceOf[Val.Literal]) {
         // Function body is a constant (e.g. `function(_) 'x'`).
-        // Keep the eager shared-value array: it is smaller and faster than a lazy view here.
-        val a = new Array[Eval](sz)
-        java.util.Arrays.fill(a.asInstanceOf[Array[AnyRef]], body.asInstanceOf[Val])
-        Val.Arr(pos, a)
+        Val.Arr.constant(pos, sz, body.asInstanceOf[Val])
       } else {
         Val.Arr.makeArray(pos, sz, func, pos, pos.noOffset, ev)
       }

sjsonnet/src/sjsonnet/stdlib/ManifestModule.scala

@@ -153,7 +153,7 @@ object ManifestModule extends AbstractFunctionModule {
    * that are used for indentation.
    */
   private object ManifestTomlEx extends Val.Builtin2("manifestTomlEx", "value", "indent") {
-    private def isTableArray(v: Val) = v.value match {
+    private def isTableArray(v: Val) = v match {
       case s: Val.Arr =>
         if (s.length == 0) false
         else {
@@ -168,68 +168,114 @@ object ManifestModule extends AbstractFunctionModule {
       case _ => false
     }
 
-    private def isSection(v: Val) = v.value.isInstanceOf[Val.Obj] || isTableArray(v.value)
+    private def isSection(v: Val) = v.isInstanceOf[Val.Obj] || isTableArray(v)
+
+    private def sortedKeyIndex(keys: Array[String], key: String): Int = {
+      var low = 0
+      var high = keys.length - 1
+      while (low <= high) {
+        val mid = (low + high) >>> 1
+        val comparison = Util.CodepointStringOrdering.compare(keys(mid), key)
+        if (comparison < 0) low = mid + 1
+        else if (comparison > 0) high = mid - 1
+        else return mid
+      }
+      -1
+    }
 
     private def renderTableInternal(
         out: StringWriter,
         v: Val.Obj,
         cumulatedIndent: String,
         indent: String,
-        path: Seq[String],
-        indexedPath: Seq[String])(implicit ev: EvalScope): StringWriter = {
-      val (sections, nonSections) =
-        v.visibleKeyNames.partition(k => isSection(v.value(k, v.pos)(ev)))
-      for (k <- nonSections.sorted(Util.CodepointStringOrdering)) {
-        out.write(cumulatedIndent)
-        out.write(TomlRenderer.escapeKey(k))
-        out.write(" = ")
-        Materializer.apply0(v.value(k, v.pos)(ev), new TomlRenderer(out, cumulatedIndent, indent))(
-          ev
-        )
+        path: mutable.ArrayBuffer[String])(implicit ev: EvalScope): StringWriter = {
+      val keys = v.sortedVisibleKeyNames
+      if (keys.length == 0) {
+        out.write('\n')
+        return out
+      }
+
+      val sectionFlags = new Array[Boolean](keys.length)
+      val visibleKeys = v.visibleKeyNames
+      var visibleKeyIdx = 0
+      while (visibleKeyIdx < visibleKeys.length) {
+        val k = visibleKeys(visibleKeyIdx)
+        val sortedIdx = sortedKeyIndex(keys, k)
+        sectionFlags(sortedIdx) = isSection(v.value(k, v.pos)(ev))
+        visibleKeyIdx += 1
+      }
+
+      val renderer = new TomlRenderer(out, cumulatedIndent, indent)
+      var keyIdx = 0
+      while (keyIdx < keys.length) {
+        if (!sectionFlags(keyIdx)) {
+          val k = keys(keyIdx)
+          out.write(cumulatedIndent)
+          TomlRenderer.writeEscapedKey(out, k)
+          out.write(" = ")
+          Materializer.apply0(v.value(k, v.pos)(ev), renderer)(ev)
+        }
+        keyIdx += 1
       }
       out.write('\n')
 
-      for (k <- sections.sorted(Util.CodepointStringOrdering)) {
-        val v0 = v.value(k, v.pos, v)(ev)
-        if (isTableArray(v0)) {
-          for (i <- 0 until v0.asArr.length) {
-            out.write(cumulatedIndent)
-            renderTableArrayHeader(out, path :+ k)
-            out.write('\n')
-            renderTableInternal(
-              out,
-              v0.asArr.value(i).asObj,
-              cumulatedIndent + indent,
-              indent,
-              path :+ k,
-              indexedPath ++ Seq(k, i.toString)
-            )
+      keyIdx = 0
+      while (keyIdx < keys.length) {
+        if (sectionFlags(keyIdx)) {
+          val k = keys(keyIdx)
+          val v0 = v.value(k, v.pos, v)(ev)
+          val childIndent = cumulatedIndent + indent
+          path += k
+          v0 match {
+            case arr: Val.Arr =>
+              var i = 0
+              while (i < arr.length) {
+                out.write(cumulatedIndent)
+                renderTableArrayHeader(out, path)
+                out.write('\n')
+                renderTableInternal(
+                  out,
+                  arr.value(i).asObj,
+                  childIndent,
+                  indent,
+                  path
+                )
+                i += 1
+              }
+            case obj: Val.Obj =>
+              out.write(cumulatedIndent)
+              renderTableHeader(out, path)
+              out.write('\n')
+              renderTableInternal(
+                out,
+                obj,
+                childIndent,
+                indent,
+                path
+              )
+            case _ =>
+              ()
           }
-        } else {
-          out.write(cumulatedIndent)
-          renderTableHeader(out, path :+ k)
-          out.write('\n')
-          renderTableInternal(
-            out,
-            v0.asObj,
-            cumulatedIndent + indent,
-            indent,
-            path :+ k,
-            indexedPath :+ k
-          )
+          path.remove(path.length - 1)
         }
+        keyIdx += 1
       }
       out
     }
 
-    private def renderTableHeader(out: StringWriter, path: Seq[String]) = {
+    private def renderTableHeader(out: StringWriter, path: mutable.ArrayBuffer[String]) = {
       out.write('[')
-      out.write(path.map(TomlRenderer.escapeKey).mkString("."))
+      var i = 0
+      while (i < path.length) {
+        if (i != 0) out.write('.')
+        TomlRenderer.writeEscapedKey(out, path(i))
+        i += 1
+      }
       out.write(']')
       out
     }
 
-    private def renderTableArrayHeader(out: StringWriter, path: Seq[String]) = {
+    private def renderTableArrayHeader(out: StringWriter, path: mutable.ArrayBuffer[String]) = {
       out.write('[')
       renderTableHeader(out, path)
       out.write(']')
@@ -243,8 +289,7 @@ object ManifestModule extends AbstractFunctionModule {
         v.value.asObj,
         "",
         indent.value.asString,
-        Seq.empty[String],
-        Seq.empty[String]
+        new mutable.ArrayBuffer[String](8)
       )(ev)
       Val.Str(pos, out.toString.strip)
     }