A smart sbt parser for the most used jmh parameters (copy to project dir and use jmh:runny)

JmhCompletion.scala

package akka.sbt.jmh.completion

import java.io.FileInputStream
import java.util.Collections

import pl.project13.scala.sbt.SbtJmh.JmhKeys.Jmh
import sbt._
import Keys._
import org.openjdk.jmh.runner.BenchmarkList
import org.openjdk.jmh.runner.BenchmarkListEntry
import org.openjdk.jmh.runner.Runner

import scala.reflect.internal.util.ScalaClassLoader.URLClassLoader

object JmhCompletion extends AutoPlugin {
  override def requires = pl.project13.scala.sbt.SbtJmh
  override def trigger = allRequirements

  val runny = inputKey[Unit]("run")

  override def projectSettings: Seq[Def.Setting[_]] = Seq(
    runny in Jmh := jmhRunTaskWithCompletion.evaluated
  )

  import complete.DefaultParsers._
  import complete.Parser

  class JmhParser(entries: Seq[BenchmarkListEntry]) {
    def singleUseArgs: Parser[String] =
      tokenDisplay("-h", "-h -- Display help, and exit.") |
      tokenDisplay("-lrf", "-lrf -- List machine-readable result formats, and exit.") |
      tokenDisplay("-lprof", "-lprof -- List profilers, and exit.")

    def simple(option: Option): Parser[Any] =
      tokenDisplay(option.name, option.description)

    def numericArg(option: Option): Parser[Any] =
      tokenDisplay(option.name ~ OptSpace ~ NatBasic, s"${option.name} <number> -- ${option.description}")

    def timeArg(option: Option): Parser[Any] =
      tokenDisplay(option.name ~ OptSpace ~ NatBasic, s"${option.name} <time> -- ${option.description}")

    def stringArg(option: Option): Parser[Any] =
      tokenDisplay(option.name ~ OptSpace ~ NotSpace, s"${option.name} <string> -- ${option.description}")

    def parser: Parser[JmhParserState] = nextParser(NothingYet)

    def nextParser(state: JmhParserState): Parser[JmhParserState] = {
      val nextParserOpt: scala.Option[Parser[JmhParserState]] =
        state match {
          case s@NothingYet => Some(
            onlyThatOptionsParser |
            withOptionsParser(s.options)
          )
          case o: JmhOptions => Some(withOptionsParser(o))
          case OnlyThatOption => None // single option already set
        }

      nextParserOpt match {
        case Some(next) => (token(Space) ~> next).flatMap(nextParser) | success(state)
        case None => success(state)
      }
    }

    case class Option(name: String, description: String, parserCons: Option => Parser[Any]) {
      val parser: Parser[Any] = parserCons(this)
    }
    val onlyThatOptions = Seq(
      Option("-h", "-h -- Display help, and exit.", simple),
      Option("-lrf", "-lrf -- List machine-readable result formats, and exit.", simple),
      Option("-lprof", "-lprof -- List profilers, and exit.", simple),
      Option("-l", "-l -- List the benchmarks, and exit.", simple),
      Option("-lp", "-lp -- List the benchmarks along with parameters, and exit.", simple)
    )
    val options = Seq(
      Option("-f", "How many times to fork a single benchmark", numericArg),
      Option("-i", "Number of measurement iterations to do", numericArg),
      Option("-wi", "Number of warmup iterations to do.", numericArg),
      Option("-w", "Minimum time to spend at each warmup iteration", timeArg),
      Option("-r", "Minimum time to spend at each measurement iteration", timeArg),
      Option("-jvm", "Use given JVM for runs", stringArg)
    )

    def onlyThatOptionsParser: Parser[JmhParserState] =
      onlyThatOptions.map(_.parser).reduce(_ | _).map(_ => OnlyThatOption)

    def withOptionsParser(options: JmhOptions): Parser[JmhParserState] =
      paramParser(options) |
      optionParser(options) |
      benchmarkParser(options)

    def optionParser(state: JmhOptions): Parser[JmhParserState] =
      options
        .filterNot(o => state.definedOptions(o.name))
        .map { o =>
          o.parser.map(_ => state.addOption(o.name))
        }
        .reduceOption(_ | _)
        .getOrElse(success(state))

    def matchedByFilter(entry: BenchmarkListEntry, filter: String): Boolean =
      entry.getUsername == filter // FIXME: improve

    def matchedByFilters(entry: BenchmarkListEntry, filters: Set[String]): Boolean =
      filters.exists(filter => matchedByFilter(entry, filter))

    val benchmarkFilterClass: Parser[Char] = charClass(c => c.isLetterOrDigit || c == '.')
    val benchmarkFilter: Parser[String] = identifier(benchmarkFilterClass, benchmarkFilterClass)
    def benchmarkParser(state: JmhOptions): Parser[JmhParserState] =
      if (entries.nonEmpty) {
        val classes =
          entries
            .filterNot(matchedByFilters(_, state.definedFilters))
            .map(_.getUsername)
            .toSet

        token(benchmarkFilter examples classes).map(state.addFilter)
      } else
        failure("No benchmarks found")

    def paramParser(state: JmhOptions): Parser[JmhParserState] = {
      val matchedEntries = entries.filter(matchedByFilters(_, state.definedFilters))
      import scala.collection.JavaConverters._
      val matchedParams = matchedEntries.flatMap(_.getParams.orElse(Collections.emptyMap()).asScala.toSeq)
      val matchedParamNames = matchedParams.map(_._1).toSet.filterNot(state.definedParams)

      token("-p" ~> Space ~> (ID examples matchedParamNames)).flatMap { paramName =>
        val paramValues = matchedParams.filter(_._1 == paramName).flatMap(_._2.toSet).toSet
        token("=") ~> selectValues(paramValues) ^^^ state.addParam(paramName)
      }
    }

    val paramValue: Parser[String] = charClass(c => !c.isWhitespace && c != ',').+.string

    def selectValues(potentialValues: Set[String]): Parser[Unit] =
      token(paramValue examples potentialValues).flatMap { value =>
        val remaining = potentialValues - value
        token(",") ~> selectValues(remaining) | success(())
      }

    sealed trait JmhParserState
    sealed trait StateWithOptions extends JmhParserState {
      def options: JmhOptions
    }
    case object NothingYet extends StateWithOptions {
      def options: JmhOptions = JmhOptions()
    }
    case object OnlyThatOption extends JmhParserState
    case class JmhOptions(
      definedOptions: Set[String] = Set.empty,
      definedFilters: Set[String] = Set.empty,
      definedParams: Set[String] = Set.empty
    ) extends StateWithOptions {
      override def options: JmhOptions = this

      def addOption(name: String): JmhParserState = copy(definedOptions = definedOptions + name)
      def addFilter(filter: String): JmhParserState = copy(definedFilters = definedFilters + filter)
      def addParam(param: String): JmhParserState = copy(definedParams = definedParams + param)
    }
  }

  def jmhParser: Def.Initialize[State => Parser[String]] =
    Def.setting {
      (state: State) =>
        val file = (classDirectory in Compile).value / "META-INF" / "BenchmarkList"
        val list: Seq[BenchmarkListEntry] =
          if (file.exists()) {
            import scala.collection.JavaConverters._
            BenchmarkList.readBenchmarkList(new FileInputStream(file)).asScala.toVector
          }
          else Nil

        matched(new JmhParser(list).parser)
    }

  def jmhRunTaskWithCompletion: Def.Initialize[InputTask[Unit]] =
    Def.inputTaskDyn {
      val args = jmhParser.parsed
      (run in Jmh).toTask(args)
    }
}