001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.util.protobuf;
006
007import java.io.IOException;
008import java.nio.ByteBuffer;
009import us.hebi.quickbuf.ProtoMessage;
010import us.hebi.quickbuf.ProtoSink;
011import us.hebi.quickbuf.ProtoSource;
012
013/**
014 * Reusable buffer for serialization/deserialization to/from a protobuf.
015 *
016 * @param <T> object type
017 * @param <MessageType> protobuf message type
018 */
019public final class ProtobufBuffer<T, MessageType extends ProtoMessage<?>> {
020  private ProtobufBuffer(Protobuf<T, MessageType> proto) {
021    m_buf = ByteBuffer.allocateDirect(1024);
022    m_sink = ProtoSink.newDirectSink();
023    m_sink.setOutput(m_buf);
024    m_source = ProtoSource.newDirectSource();
025    m_msg = proto.createMessage();
026    m_proto = proto;
027  }
028
029  public static <T, MessageType extends ProtoMessage<?>> ProtobufBuffer<T, MessageType> create(
030      Protobuf<T, MessageType> proto) {
031    return new ProtobufBuffer<>(proto);
032  }
033
034  /**
035   * Gets the protobuf object of the stored type.
036   *
037   * @return protobuf object
038   */
039  public Protobuf<T, MessageType> getProto() {
040    return m_proto;
041  }
042
043  /**
044   * Gets the type string.
045   *
046   * @return type string
047   */
048  public String getTypeString() {
049    return m_proto.getTypeString();
050  }
051
052  /**
053   * Serializes a value to a ByteBuffer. The returned ByteBuffer is a direct byte buffer with the
054   * position set to the end of the serialized data.
055   *
056   * @param value value
057   * @return byte buffer
058   * @throws IOException if serialization failed
059   */
060  public ByteBuffer write(T value) throws IOException {
061    m_msg.clearQuick();
062    m_proto.pack(m_msg, value);
063    int size = m_msg.getSerializedSize();
064    if (size > m_buf.capacity()) {
065      m_buf = ByteBuffer.allocateDirect(size * 2);
066      m_sink.setOutput(m_buf);
067    }
068    m_sink.reset();
069    m_msg.writeTo(m_sink);
070    m_buf.position(m_sink.getTotalBytesWritten());
071    return m_buf;
072  }
073
074  /**
075   * Deserializes a value from a byte array, creating a new object.
076   *
077   * @param buf byte array
078   * @param start starting location within byte array
079   * @param len length of serialized data
080   * @return new object
081   * @throws IOException if deserialization failed
082   */
083  public T read(byte[] buf, int start, int len) throws IOException {
084    m_msg.clearQuick();
085    m_source.setInput(buf, start, len);
086    m_msg.mergeFrom(m_source);
087    return m_proto.unpack(m_msg);
088  }
089
090  /**
091   * Deserializes a value from a byte array, creating a new object.
092   *
093   * @param buf byte array
094   * @return new object
095   * @throws IOException if deserialization failed
096   */
097  public T read(byte[] buf) throws IOException {
098    return read(buf, 0, buf.length);
099  }
100
101  /**
102   * Deserializes a value from a ByteBuffer, creating a new object.
103   *
104   * @param buf byte buffer
105   * @return new object
106   * @throws IOException if deserialization failed
107   */
108  public T read(ByteBuffer buf) throws IOException {
109    m_msg.clearQuick();
110    m_source.setInput(buf);
111    m_msg.mergeFrom(m_source);
112    return m_proto.unpack(m_msg);
113  }
114
115  /**
116   * Deserializes a value from a byte array into a mutable object.
117   *
118   * @param out object (will be updated with deserialized contents)
119   * @param buf byte array
120   * @param start starting location within byte array
121   * @param len length of serialized data
122   * @throws IOException if deserialization failed
123   * @throws UnsupportedOperationException if the object is immutable
124   */
125  public void readInto(T out, byte[] buf, int start, int len) throws IOException {
126    m_msg.clearQuick();
127    m_source.setInput(buf, start, len);
128    m_msg.mergeFrom(m_source);
129    m_proto.unpackInto(out, m_msg);
130  }
131
132  /**
133   * Deserializes a value from a byte array into a mutable object.
134   *
135   * @param out object (will be updated with deserialized contents)
136   * @param buf byte array
137   * @throws IOException if deserialization failed
138   * @throws UnsupportedOperationException if the object is immutable
139   */
140  public void readInto(T out, byte[] buf) throws IOException {
141    readInto(out, buf, 0, buf.length);
142  }
143
144  /**
145   * Deserializes a value from a ByteBuffer into a mutable object.
146   *
147   * @param out object (will be updated with deserialized contents)
148   * @param buf byte buffer
149   * @throws IOException if deserialization failed
150   * @throws UnsupportedOperationException if the object is immutable
151   */
152  public void readInto(T out, ByteBuffer buf) throws IOException {
153    m_msg.clearQuick();
154    m_source.setInput(buf);
155    m_msg.mergeFrom(m_source);
156    m_proto.unpackInto(out, m_msg);
157  }
158
159  private ByteBuffer m_buf;
160  private final ProtoSink m_sink;
161  private final ProtoSource m_source;
162  private final MessageType m_msg;
163  private final Protobuf<T, MessageType> m_proto;
164}