content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.beam.sdk.jmh.util;
  
  import java.io.IOException;
  import java.io.OutputStream;
  import java.util.Random;
  import org.apache.beam.sdk.util.ByteStringOutputStream;
  import org.openjdk.jmh.annotations.Benchmark;
  import org.openjdk.jmh.annotations.Level;
  import org.openjdk.jmh.annotations.OperationsPerInvocation;
  import org.openjdk.jmh.annotations.Scope;
  import org.openjdk.jmh.annotations.Setup;
  import org.openjdk.jmh.annotations.State;
  import org.openjdk.jmh.annotations.TearDown;
  import org.openjdk.jmh.infra.Blackhole;
  
  /** Benchmarks for {@link org.apache.beam.sdk.util.VarInt} and variants. */
  @OperationsPerInvocation(VarIntBenchmark.VALUES_PER_INVOCATION)
  public class VarIntBenchmark {
    static final int VALUES_PER_INVOCATION = 2048;
    private static final Random RNG = new Random(314159);
  
    /** Output to {@link Blackhole}. Do nothing, assume nothing. */
    @State(Scope.Benchmark)
    public static class BlackholeOutput {
      OutputStream stream;
  
      @Setup
      public void setup(Blackhole bh) {
        stream =
            new OutputStream() {
              @Override
              public void write(int b) {
                bh.consume(b);
              }
  
              @Override
              public void write(byte[] b) throws IOException {
                bh.consume(b);
              }
  
              @Override
              public void write(byte[] b, int off, int len) throws IOException {
                bh.consume(b);
              }
            };
      }
    }
  
    /** Output to {@link ByteStringOutputStream}. */
    @State(Scope.Thread)
    public static class ByteStringOutput {
      final ByteStringOutputStream stream = new ByteStringOutputStream();
  
      // Unfortunately, this needs to be cleaned up after use to avoid OOMs.
      // It's not generally recommended to use Level.Invocation, but there's no way around it.
      @TearDown(Level.Invocation)
      public void tearDown(Blackhole bh) {
        bh.consume(stream.toByteStringAndReset());
      }
    }
  
    /** Input from randomly generated bytes. */
    @State(Scope.Benchmark)
    public static class Bytes {
      long[] values = new long[VALUES_PER_INVOCATION];
  
      @Setup
      public void setup() {
        values = new long[VALUES_PER_INVOCATION];
        byte[] bytes = new byte[VALUES_PER_INVOCATION];
        RNG.nextBytes(bytes);
  
        for (int i = 0; i < VALUES_PER_INVOCATION; i++) {
          values[i] = (long) (bytes[i] & 0x7F);
        }
      }
    }
  
    /** Input from randomly generated longs. */
    @State(Scope.Benchmark)
    public static class Longs {
      long[] values = new long[VALUES_PER_INVOCATION];
  
      @Setup
      public void setup() {
        values = new long[VALUES_PER_INVOCATION];
  
        for (int i = 0; i < VALUES_PER_INVOCATION; i++) {
          // This gaussian random is used to determine the encoded output size of the sample.
          // The distribution value is tweaked to favor small integers, positive more so than
          // negative.
          double g = RNG.nextGaussian();
          double s = 3;
          g = 10 * Math.min(Math.abs(g < 0 ? g + s : g / (s / 2)), s) / s;
  
          // Construct a bitmask to keep up to numBits of the input.
          // Find the lowest bit to set in the 7 bit segment below numBits.
          int numBits = 7 * (int) g;
          long mask = ~(~0x7fL << numBits);
          long low = 1L << numBits;
  
          values[i] = (RNG.nextLong() & mask) | low;
        }
      }
    }
  
    // Used in Beam 2.52.0
    static void encodeDoLoop(long v, OutputStream stream) throws IOException {
      do {
        // Encode next 7 bits + terminator bit
        long bits = v & 0x7F;
        v >>>= 7;
        byte b = (byte) (bits | ((v != 0) ? 0x80 : 0));
        stream.write(b);
      } while (v != 0);
    }
  
    // A tweak of the above, replacing a compare with a few bitwise operations.
    static void encodeDoLoopTwiddle(long v, OutputStream stream) throws IOException {
      do {
        // Encode next 7 bits + terminator bit
        long bits = v & 0x7F;
        v >>>= 7;
        long cont = (-v >> 63) & 0x80;
        byte b = (byte) (bits | cont);
        stream.write(b);
      } while (v != 0);
    }
  
    // Use a mask check to do less work for 1 byte output.
    static void encodeLoop(long v, OutputStream stream) throws IOException {
      while ((v & ~0x7FL) != 0) {
        stream.write((byte) (v | 0x80));
        v >>>= 7;
      }
      stream.write((byte) v);
    }
  
    // As above, but unrolled.
    static void encodeUnrolled(long v, OutputStream stream) throws IOException {
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      if ((v & ~0x7F) == 0) {
        stream.write((byte) v);
        return;
      }
      stream.write((byte) (v | 0x80));
      v >>>= 7;
      stream.write((byte) (v));
    }
  
    @Benchmark
    public void encodeDoLoopBlackhole(Longs input, BlackholeOutput output) throws IOException {
      for (long l : input.values) {
        encodeDoLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeDoLoopByteString(Longs input, ByteStringOutput output) throws IOException {
      for (long l : input.values) {
        encodeDoLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeDoLoopTwiddleBlackhole(Longs input, BlackholeOutput output) throws IOException {
      for (long l : input.values) {
        encodeDoLoopTwiddle(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeDoLoopTwiddleByteString(Longs input, ByteStringOutput output)
        throws IOException {
      for (long l : input.values) {
        encodeDoLoopTwiddle(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeLoopBlackhole(Longs input, BlackholeOutput output) throws IOException {
      for (long l : input.values) {
        encodeLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeLoopByteString(Longs input, ByteStringOutput output) throws IOException {
      for (long l : input.values) {
        encodeLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeUnrolledBlackhole(Longs input, BlackholeOutput output) throws IOException {
      for (long l : input.values) {
        encodeUnrolled(l, output.stream);
      }
    }
  
    @Benchmark
    public void encodeUnrolledByteString(Longs input, ByteStringOutput output) throws IOException {
      for (long l : input.values) {
        encodeUnrolled(l, output.stream);
      }
    }
  
    @Benchmark
    public void singleByteEncodeDoLoopByteString(Bytes input, ByteStringOutput output)
        throws IOException {
      for (long l : input.values) {
        encodeDoLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void singleByteEncodeDoLoopTwiddleByteString(Bytes input, ByteStringOutput output)
        throws IOException {
      for (long l : input.values) {
        encodeDoLoopTwiddle(l, output.stream);
      }
    }
  
    @Benchmark
    public void singleByteEncodeLoopByteString(Bytes input, ByteStringOutput output)
        throws IOException {
      for (long l : input.values) {
        encodeLoop(l, output.stream);
      }
    }
  
    @Benchmark
    public void singleByteEncodeUnrolledByteString(Bytes input, ByteStringOutput output)
        throws IOException {
      for (long l : input.values) {
        encodeUnrolled(l, output.stream);
      }
    }
  }
  