/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.jmh;

import org.openjdk.jmh.results.RunResult;
import org.openjdk.jmh.results.format.ResultFormatType;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.ChainedOptionsBuilder;
import org.openjdk.jmh.runner.options.OptionsBuilder;
import org.openjdk.jmh.runner.options.VerboseMode;
import org.openjdk.jmh.runner.options.WarmupMode;

import java.time.LocalDateTime;
import java.util.Collection;
import java.util.function.Consumer;

import static java.lang.String.format;
import static java.time.format.DateTimeFormatter.ISO_DATE_TIME;
import static java.util.Objects.requireNonNull;

public final class Benchmarks
{
    private Benchmarks() {}

    public static BenchmarkBuilder benchmark(Class<?> benchmarkClass)
    {
        ChainedOptionsBuilder optionsBuilder = new OptionsBuilder()
                .verbosity(VerboseMode.NORMAL)
                .include("^\\Q" + benchmarkClass.getName() + ".\\E")
                .resultFormat(ResultFormatType.JSON)
                .result(format("%s/%s-result-%s.json", System.getProperty("java.io.tmpdir"), benchmarkClass.getSimpleName(), ISO_DATE_TIME.format(LocalDateTime.now())));
        return new BenchmarkBuilder(optionsBuilder);
    }

    public static BenchmarkBuilder benchmark(Class<?> benchmarkClass, WarmupMode warmupMode)
    {
        return benchmark(benchmarkClass)
                .withOptions(optionsBuilder -> optionsBuilder.warmupMode(warmupMode));
    }

    public static class BenchmarkBuilder
    {
        private final ChainedOptionsBuilder optionsBuilder;

        private BenchmarkBuilder(ChainedOptionsBuilder optionsBuilder)
        {
            this.optionsBuilder = requireNonNull(optionsBuilder, "optionsBuilder is null");
        }

        public BenchmarkBuilder withOptions(Consumer<ChainedOptionsBuilder> optionsConsumer)
        {
            optionsConsumer.accept(optionsBuilder);
            return this;
        }

        public Collection<RunResult> run()
                throws RunnerException
        {
            return new Runner(optionsBuilder.build()).run();
        }
    }
}

@@ -578,16 +578,63 @@
 o.jmh;%0A%0A
+import org.intellij.lang.annotations.Language;%0A
 import o
@@ -1531,77 +1531,8 @@
 AL)%0A
-                .include(%22%5E%5C%5CQ%22 + benchmarkClass.getName() + %22.%5C%5CE%22)%0A
     
@@ -1796,24 +1796,40 @@
 tionsBuilder
+, benchmarkClass
 );%0A    %7D%0A%0A  
@@ -2166,16 +2166,63 @@
 Builder;
+%0A        private final Class%3C?%3E benchmarkClass;
 %0A%0A      
@@ -2284,16 +2284,41 @@
 sBuilder
+, Class%3C?%3E benchmarkClass
 )%0A      
@@ -2413,16 +2413,108 @@
 null%22);%0A
+            this.benchmarkClass = requireNonNull(benchmarkClass, %22benchmarkClass is null%22);%0A
         
@@ -2726,84 +2726,607 @@
 lic 
-Collection%3CRunResult%3E run()%0A                throws RunnerException%0A        %7B
+BenchmarkBuilder includeAll()%0A        %7B%0A            optionsBuilder.include(%22%5E%5C%5CQ%22 + benchmarkClass.getName() + %22.%5C%5CE%22);%0A            return this;%0A        %7D%0A%0A        public BenchmarkBuilder includeMethod(@Language(%22RegExp%22) String benchmarkMethod)%0A        %7B%0A            optionsBuilder.include(%22%5E%5C%5CQ%22 + benchmarkClass.getName() + %22.%5C%5CE(%22 + benchmarkMethod + %22)$%22);%0A            return this;%0A        %7D%0A%0A        public Collection%3CRunResult%3E run()%0A                throws RunnerException%0A        %7B%0A            if (optionsBuilder.build().getIncludes().isEmpty()) %7B%0A                includeAll();%0A            %7D
 %0A   

Add option to benchmark only selected methods