ProtobufBuffer.java

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.
 
package org.wpilib.util.protobuf;
 
import java.io.IOException;
import java.nio.ByteBuffer;
import us.hebi.quickbuf.ProtoMessage;
import us.hebi.quickbuf.ProtoSink;
import us.hebi.quickbuf.ProtoSource;

public final class ProtobufBuffer<T, MessageType extends ProtoMessage<?>> {
  private ProtobufBuffer(Protobuf<T, MessageType> proto) {
    m_buf = ByteBuffer.allocateDirect(1024);
    m_sink = ProtoSink.newDirectSink();
    m_sink.setOutput(m_buf);
    m_source = ProtoSource.newDirectSource();
    m_msg = proto.createMessage();
    m_proto = proto;
  }

  public static <T, MessageType extends ProtoMessage<?>> ProtobufBuffer<T, MessageType> create(
      Protobuf<T, MessageType> proto) {
    return new ProtobufBuffer<>(proto);
  }

  public Protobuf<T, MessageType> getProto() {
    return m_proto;
  }

  public String getTypeString() {
    return m_proto.getTypeString();
  }

  public ByteBuffer write(T value) throws IOException {
    m_msg.clearQuick();
    m_proto.pack(m_msg, value);
    int size = m_msg.getSerializedSize();
    if (size > m_buf.capacity()) {
      m_buf = ByteBuffer.allocateDirect(size * 2);
      m_sink.setOutput(m_buf);
    }
    m_sink.reset();
    m_msg.writeTo(m_sink);
    m_buf.position(m_sink.getTotalBytesWritten());
    return m_buf;
  }

  public T read(byte[] buf, int start, int len) throws IOException {
    m_msg.clearQuick();
    m_source.setInput(buf, start, len);
    m_msg.mergeFrom(m_source);
    return m_proto.unpack(m_msg);
  }

  public T read(byte[] buf) throws IOException {
    return read(buf, 0, buf.length);
  }

  public T read(ByteBuffer buf) throws IOException {
    m_msg.clearQuick();
    m_source.setInput(buf);
    m_msg.mergeFrom(m_source);
    return m_proto.unpack(m_msg);
  }

  public void readInto(T out, byte[] buf, int start, int len) throws IOException {
    m_msg.clearQuick();
    m_source.setInput(buf, start, len);
    m_msg.mergeFrom(m_source);
    m_proto.unpackInto(out, m_msg);
  }

  public void readInto(T out, byte[] buf) throws IOException {
    readInto(out, buf, 0, buf.length);
  }

  public void readInto(T out, ByteBuffer buf) throws IOException {
    m_msg.clearQuick();
    m_source.setInput(buf);
    m_msg.mergeFrom(m_source);
    m_proto.unpackInto(out, m_msg);
  }
 
  private ByteBuffer m_buf;
  private final ProtoSink m_sink;
  private final ProtoSource m_source;
  private final MessageType m_msg;
  private final Protobuf<T, MessageType> m_proto;
}