Ammonite script to merge JMH outputs for comparison

merge.sc

#!/usr/bin/env amm

import $ivy.`com.lihaoyi::ammonite-ops:2.4.0`

import ammonite.ops._

import scala.collection.immutable.ArraySeq
import scala.util.control.NoStackTrace

object Abort extends Throwable("unable to process inputs") with NoStackTrace

def err(msg: String): Unit = System.err.println(msg)

def terminate(msg: String): Nothing = {
  err(msg)
  throw Abort
}

def expect(cond: Boolean, msg: String): Unit = {
  if (!cond) terminate(msg)
}

case class Segment(index: Int, text: String)

case class OffsetSegment(additionalPadding: Int, segment: Segment)

case class Segments(before: IndexedSeq[Segment], after: IndexedSeq[Segment]) {
  require(before.length == after.length, "before and after must have the same number of segments")

  def withOffsets: OffsetSegments = {
    val before = this.before
    val after = this.after
    val len = before.length
    val beforeBuilder = ArraySeq.newBuilder[OffsetSegment]
    val afterBuilder = ArraySeq.newBuilder[OffsetSegment]
    beforeBuilder.sizeHint(len)
    afterBuilder.sizeHint(len)
    
    var beforeTotalPadding = 0
    var afterTotalPadding = 0
    var i = 0
    while (i < len) {
      // our goal is to pad the one that starts sooner to line it up
      val bSegment = before(i)
      val aSegment = after(i)
      val bEffectiveIndex = bSegment.index + beforeTotalPadding
      val aEffectiveIndex = aSegment.index + afterTotalPadding
      val paddingDiff = bEffectiveIndex - aEffectiveIndex // realistically nowhere near max or min `Int`
      if (paddingDiff < 0) {
        // after has more padding
        val bPadding = -paddingDiff
        beforeTotalPadding += bPadding
        beforeBuilder += OffsetSegment(bPadding, bSegment)
        afterBuilder += OffsetSegment(0, aSegment)
      } else if (paddingDiff > 0) {
        // before has more padding
        afterTotalPadding += paddingDiff
        beforeBuilder += OffsetSegment(0, bSegment)
        afterBuilder += OffsetSegment(paddingDiff, aSegment)
      } else /* paddingDiff == 0 */ {
        // same padding
        beforeBuilder += OffsetSegment(0, bSegment)
        afterBuilder += OffsetSegment(0, aSegment)
      }
      i += 1
    }
    
    OffsetSegments(beforeBuilder.result(), afterBuilder.result())
  }
}

case class OffsetSegments(before: IndexedSeq[OffsetSegment], after: IndexedSeq[OffsetSegment]) {
  require(before.length == after.length, "before and after must have the same number of segments")
  
  private def alignWithSegments(line: String, segments: IndexedSeq[OffsetSegment]): String = {
    val b = new java.lang.StringBuilder
    val len = segments.length
    
    var lastTextIndex = 0
    var i = 0
    while (i < len) {
      val segment = segments(i)
      val padding = segment.additionalPadding
      var index = segment.segment.index
      // most headings are aligned right, so we need to find the start of text
      while (index > 0 && line.charAt(index) != ' ') { index -= 1 }
      expect(index >= lastTextIndex, "no spaces since start of last text—is this JMH output?") // TODO: improve this check
      b.append(line, lastTextIndex, index)
      lastTextIndex = index
      b.append(" " * segment.additionalPadding)
      i += 1
    }
    // append remaining text with no spaces after it
    b.append(line, lastTextIndex, line.length)
    b.toString()
  }
  
  def alignAsBefore(line: String): String = alignWithSegments(line, before)
  def alignAsAfter(line: String): String = alignWithSegments(line, after)
}

def findCommonPrefix(a: String, b: String): String = {
  var i = 0
  val limit = math.min(a.length, b.length)
  while (i < limit && a.charAt(i) == b.charAt(i)) { i += 1 }
  a.substring(0, i)
}

def findCommonPrefix(strings: Seq[String]): String = { 
  assert(strings.nonEmpty, "cannot find common prefix of zero strings")
  strings.reduce(findCommonPrefix)
}

def cleanOutput(path: Path): IndexedSeq[String] = {
  val lines = read.lines(path)
  expect(lines.nonEmpty, s"${path.relativeTo(pwd)} is empty")
  val prefix = findCommonPrefix(lines)
  lines.map(_.stripPrefix(prefix))
}

def segments(line: String): IndexedSeq[Segment] = {
  val b = ArraySeq.newBuilder[Segment]
  
  var lastSegmentStart = 0 // doesn't matter, will be overwritten
  var i = 0
  var scanningText = false
  val len = line.length
  while (i < len) {
    if (!scanningText) {
      if (line.charAt(i) == ' ') {
        // still not scanning text, nothing to do
      } else {
        // found text. set `lastSegmentStart` to keep track of start of text
        lastSegmentStart = i
        scanningText = true
      }
    } else /* scanningText */ {
      if (line.charAt(i) == ' ') {
        // reached the end of the text
        val text = line.substring(lastSegmentStart, i)
        b += Segment(lastSegmentStart, text)
        scanningText = false
      } else {
        // still in the middle of text, nothing to do
      }
    }
    i += 1
  }
  if (scanningText) {
    // add text with no spaces after it
    val text = line.substring(lastSegmentStart, i)
    b += Segment(lastSegmentStart, text)
  }

  b.result()
}

def interleaveAsDiff(before: IndexedSeq[String], after: IndexedSeq[String]): IndexedSeq[String] = {
  assert(before.length == after.length, "cannot interleave if before and after are different lengths")
  
  val b = ArraySeq.newBuilder[String]
  b += s" ${before.head}" // header
  
  val len = before.length
  var i = 1
  while (i < len) {
    b += s"-${before(i)}"
    b += s"+${after(i)}"
    i += 1
  }
  
  b.result()
}

@main
def main(beforePath: Path = pwd/"before.txt",
         afterPath: Path = pwd/"after.txt",
         mergedPath: Path = pwd/"merged.txt"): Unit = {
  def relative(path: Path): String = path.relativeTo(pwd).toString
  
  def checkInput(path: Path): Unit = {
    expect(exists(path), s"No such file: ${relative(path)}")
    expect(path.isFile, s"${relative(path)} is not a normal file (probably a directory)")
  }
  checkInput(beforePath)
  checkInput(afterPath)
  
  val before = cleanOutput(beforePath)
  val after = cleanOutput(afterPath)
  expect(before.nonEmpty, s"${relative(beforePath)} is empty")
  expect(after.nonEmpty, s"${relative(afterPath)} is empty")
  expect(before.length == after.length,
         s"${relative(beforePath)} and ${relative(afterPath)} must have the same number of lines")

  val beforeHeading = before.head
  val afterHeading = after.head
  expect(beforeHeading.startsWith("Benchmark"), s"${relative(beforePath)} is missing JMH column headings")
  expect(afterHeading.startsWith("Benchmark"), s"${relative(afterPath)} is missing JMH column headings")
  val offsetSegments = Segments(segments(beforeHeading), segments(afterHeading)).withOffsets
  
  expect(before.length >= 2, "must have at least one benchmark result in each file after the JMH column headings")
  val beforeAligned = before.map(offsetSegments.alignAsBefore)
  val afterAligned = after.map(offsetSegments.alignAsAfter)

  val diff = interleaveAsDiff(beforeAligned, afterAligned)
  write.over(mergedPath, diff.mkString("", "\n", "\n"))
}