Optimize ByColor flatMap/flatten

[lenguyenthanh]

benchmarks code:

package benchmarks

import org.openjdk.jmh.annotations.*
import org.openjdk.jmh.infra.Blackhole
import java.util.concurrent.TimeUnit
import chess.ByColor

@State(Scope.Thread)
@BenchmarkMode(Array(Mode.Throughput))
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Measurement(iterations = 15, timeUnit = TimeUnit.SECONDS, time = 3)
@Warmup(iterations = 15, timeUnit = TimeUnit.SECONDS, time = 3)
@Fork(value = 3)
@Threads(value = 1)
class ByColorBench:

  // the unit of CPU work per iteration
  private[this] val Work: Long = 10

  case class Game(pgn: String)

  var byColor: ByColor[Game] = _
  var byColorO: ByColor[Option[Game]] = _
  def f(s: String): Option[String] =
    if s.size % 3 == 0 then None
    else Some(s)

  @Setup
  def setup() =
    byColor = ByColor(Game("foo"), Game("bar"))
    byColorO = ByColor(Some(Game("foo")), Some(Game("bar")))

  @Benchmark
  def byColorFlatMap =
    Blackhole.consumeCPU(Work)
    byColor.flatMap(g => f(g.pgn).map(Game(_)))

  @Benchmark
  def allFlatMap =
    Blackhole.consumeCPU(Work)
    byColor.all.flatMap(g => f(g.pgn).map(Game(_)))

  @Benchmark
  def byColorFlatten =
    Blackhole.consumeCPU(Work)
    byColorO.flatten

  @Benchmark
  def allFlatten =
    Blackhole.consumeCPU(Work)
    byColorO.all.flatten

results:

[info] Benchmark                            Mode  Cnt       Score     Error   Units
[info] ByColorBench.allFlatMap             thrpt   45   99880.913 ± 516.563  ops/ms
[info] ByColorBench.allFlatten             thrpt   45   61322.261 ± 178.359  ops/ms
[info] ByColorBench.byColorFlatMap         thrpt   45  117163.184 ± 418.543  ops/ms
[info] ByColorBench.byColorFlatten         thrpt   45   93482.812 ± 413.525  ops/ms
[info] Benchmark result is saved to jmh-result.json