ProtoOutputStream.java

package de.schegge.rosinante.io;

import java.io.ByteArrayOutputStream;
import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.function.Function;

public class ProtoOutputStream extends FilterOutputStream {

    public ProtoOutputStream(OutputStream out) {
        super(out);
    }

    private void writeWireTypeAndFieldNumber(WireType wireType, int fieldNumber) throws IOException {
        writeVariableByteInt(fieldNumber << 3 | wireType.ordinal() & 0b00000111);
    }

    protected void writeVariableByteInt(int number) throws IOException {
        int continuationBytes = (31 - Integer.numberOfLeadingZeros(number)) / 7;
        for (int i = 0; i < continuationBytes; ++i) {
            write(((byte) ((number & 0x7F) | 0x80)));
            number >>>= 7;
        }
        write((byte) number);
    }

    protected void writeVariableByteLong(long number) throws IOException {
        int continuationBytes = (63 - Long.numberOfLeadingZeros(number)) / 7;
        for (int i = 0; i < continuationBytes; ++i) {
            write(((byte) ((number & 0x7F) | 0x80)));
            number >>>= 7;
        }
        write((byte) number);
    }

    private <T> void writeRepeatedVarInt(int fieldNumber, List<T> values, Function<T, Integer> convert) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ProtoOutputStream protoOutputStream = new ProtoOutputStream(outputStream);
        for (T value : values) {
            protoOutputStream.writeVariableByteInt(convert.apply(value));
        }
        byte[] bytes = outputStream.toByteArray();
        writeWireTypeAndFieldNumber(WireType.LEN, fieldNumber);
        writeVariableByteInt(bytes.length);
        write(bytes);
    }

    private void writeRepeatedVarLong(int fieldNumber, List<Long> values, Function<Long, Long> convert) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ProtoOutputStream protoOutputStream = new ProtoOutputStream(outputStream);
        for (Long value : values) {
            protoOutputStream.writeVariableByteLong(convert.apply(value));
        }
        byte[] bytes = outputStream.toByteArray();
        writeWireTypeAndFieldNumber(WireType.LEN, fieldNumber);
        writeVariableByteInt(bytes.length);
        write(bytes);
    }

    public void writeBoolean(int fieldNumber, boolean value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.VARINT, fieldNumber);
        writeVariableByteInt(value ? 1 : 0);
    }

    public void writeInteger(int fieldNumber, int value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.VARINT, fieldNumber);
        writeVariableByteInt(value);
    }

    public void writeZigZagInteger(int fieldNumber, int value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.VARINT, fieldNumber);
        writeVariableByteInt((value >> 31) ^ (value << 1));
    }

    public void writeLong(int fieldNumber, long value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.VARINT, fieldNumber);
        writeVariableByteLong(value);
    }

    public void writeZigZagLong(int fieldNumber, long value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.VARINT, fieldNumber);
        writeVariableByteLong((value >> 63) ^ (value << 1));
    }

    public void writeString(int fieldNumber, String test) throws IOException {
        writeBytes(fieldNumber, test.getBytes(StandardCharsets.UTF_8));
    }

    public void writeBytes(int fieldNumber, byte[] bytes) throws IOException {
        writeWireTypeAndFieldNumber(WireType.LEN, fieldNumber);
        writeVariableByteInt(bytes.length);
        write(bytes);
    }

    public void writeBoolean(int fieldNumber, List<Boolean> values) throws IOException {
        writeRepeatedVarInt(fieldNumber, values, x -> x ? 1 : 0);
    }

    public void writeInteger(int fieldNumber, List<Integer> values) throws IOException {
        writeRepeatedVarInt(fieldNumber, values, x -> x);
    }

    public void writeZigZagInteger(int fieldNumber, List<Integer> values) throws IOException {
        writeRepeatedVarInt(fieldNumber, values, value -> (value >> 31) ^ (value << 1));
    }

    public void writeZigZagLong(int fieldNumber, List<Long> values) throws IOException {
        writeRepeatedVarLong(fieldNumber, values, value -> (value >> 63) ^ (value << 1));
    }

    public void writeLong(int fieldNumber, List<Long> values) throws IOException {
        writeRepeatedVarLong(fieldNumber, values, x -> x);
    }

    public void writeString(int fieldNumber, List<String> texts) throws IOException {
        for (String text : texts) {
            writeString(fieldNumber, text);
        }
    }

    public void writeBytes(int fieldNumber, List<byte[]> bytess) throws IOException {
        for (byte[] bytes : bytess) {
            writeBytes(fieldNumber, bytes);
        }
    }

    private void writeLittleEndian32(int value) throws IOException {
        write(value & 0xFF);
        write((value >> 8) & 0xFF);
        write((value >> 16) & 0xFF);
        write((value >> 24) & 0xFF);
    }

    private void writeLittleEndian64(long value) throws IOException {
        for (int i = 0; i < 8; i++) {
            write((int) (value & 0xFF));
            value >>= 8;
        }
    }

    public void writeFloat(int fieldNumber, float value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.I32, fieldNumber);
        writeLittleEndian32(Float.floatToRawIntBits(value));
    }

    public void writeDouble(int fieldNumber, double value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.I64, fieldNumber);
        writeLittleEndian64(Double.doubleToRawLongBits(value));
    }

    public void writeFixed32(int fieldNumber, int value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.I32, fieldNumber);
        writeLittleEndian32(value);
    }

    public void writeFixed64(int fieldNumber, long value) throws IOException {
        writeWireTypeAndFieldNumber(WireType.I64, fieldNumber);
        writeLittleEndian64(value);
    }

    public void writeSFixed32(int fieldNumber, int value) throws IOException {
        writeFixed32(fieldNumber, value);
    }

    public void writeSFixed64(int fieldNumber, long value) throws IOException {
        writeFixed64(fieldNumber, value);
    }

    private <T> void writeRepeatedFixed32(int fieldNumber, List<T> values, java.util.function.Function<T, Integer> convert) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ProtoOutputStream protoOutputStream = new ProtoOutputStream(outputStream);
        for (T value : values) {
            protoOutputStream.writeLittleEndian32(convert.apply(value));
        }
        byte[] bytes = outputStream.toByteArray();
        writeWireTypeAndFieldNumber(WireType.LEN, fieldNumber);
        writeVariableByteInt(bytes.length);
        write(bytes);
    }

    private <T> void writeRepeatedFixed64(int fieldNumber, List<T> values, java.util.function.Function<T, Long> convert) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        ProtoOutputStream protoOutputStream = new ProtoOutputStream(outputStream);
        for (T value : values) {
            protoOutputStream.writeLittleEndian64(convert.apply(value));
        }
        byte[] bytes = outputStream.toByteArray();
        writeWireTypeAndFieldNumber(WireType.LEN, fieldNumber);
        writeVariableByteInt(bytes.length);
        write(bytes);
    }

    public void writeFloat(int fieldNumber, List<Float> values) throws IOException {
        writeRepeatedFixed32(fieldNumber, values, v -> Float.floatToRawIntBits(v));
    }

    public void writeDouble(int fieldNumber, List<Double> values) throws IOException {
        writeRepeatedFixed64(fieldNumber, values, v -> Double.doubleToRawLongBits(v));
    }

    public void writeFixed32(int fieldNumber, List<Integer> values) throws IOException {
        writeRepeatedFixed32(fieldNumber, values, v -> v);
    }

    public void writeFixed64(int fieldNumber, List<Long> values) throws IOException {
        writeRepeatedFixed64(fieldNumber, values, v -> v);
    }

    public void writeSFixed32(int fieldNumber, List<Integer> values) throws IOException {
        writeFixed32(fieldNumber, values);
    }

    public void writeSFixed64(int fieldNumber, List<Long> values) throws IOException {
        writeFixed64(fieldNumber, values);
    }
}