1 /*
2  * Copyright (C) 2022 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.modules.utils;
18 
19 import android.annotation.NonNull;
20 
21 import java.io.BufferedOutputStream;
22 import java.io.Closeable;
23 import java.io.DataOutput;
24 import java.io.DataOutputStream;
25 import java.io.Flushable;
26 import java.io.IOException;
27 import java.io.OutputStream;
28 import java.util.HashMap;
29 import java.util.Objects;
30 
31 /**
32  * Optimized implementation of {@link DataOutput} which buffers data in memory
33  * before flushing to the underlying {@link OutputStream}.
34  * <p>
35  * Benchmarks have demonstrated this class is 2x more efficient than using a
36  * {@link DataOutputStream} with a {@link BufferedOutputStream}.
37  */
38 public class FastDataOutput implements DataOutput, Flushable, Closeable {
39     protected static final int MAX_UNSIGNED_SHORT = 65_535;
40 
41     protected static final int DEFAULT_BUFFER_SIZE = 32_768;
42 
43     protected final byte[] mBuffer;
44     protected final int mBufferCap;
45 
46     private OutputStream mOut;
47     protected int mBufferPos;
48 
49     /**
50      * Values that have been "interned" by {@link #writeInternedUTF(String)}.
51      */
52     private final HashMap<String, Integer> mStringRefs = new HashMap<>();
53 
FastDataOutput(@onNull OutputStream out, int bufferSize)54     public FastDataOutput(@NonNull OutputStream out, int bufferSize) {
55         if (bufferSize < 8) {
56             throw new IllegalArgumentException();
57         }
58 
59         mBuffer = newByteArray(bufferSize);
60         mBufferCap = mBuffer.length;
61 
62         setOutput(out);
63     }
64 
65     /**
66      * Obtain a {@link FastDataOutput} configured with the given
67      * {@link OutputStream} and which encodes large code-points using 3-byte
68      * sequences.
69      * <p>
70      * This <em>is</em> compatible with the {@link DataOutput} API contract,
71      * which specifies that large code-points must be encoded with 3-byte
72      * sequences.
73      */
obtain(@onNull OutputStream out)74     public static FastDataOutput obtain(@NonNull OutputStream out) {
75         return new FastDataOutput(out, DEFAULT_BUFFER_SIZE);
76     }
77 
78     /**
79      * Release a {@link FastDataOutput} to potentially be recycled. You must not
80      * interact with the object after releasing it.
81      */
release()82     public void release() {
83         if (mBufferPos > 0) {
84             throw new IllegalStateException("Lingering data, call flush() before releasing.");
85         }
86 
87         mOut = null;
88         mBufferPos = 0;
89         mStringRefs.clear();
90     }
91 
newByteArray(int bufferSize)92     public byte[] newByteArray(int bufferSize) {
93         return new byte[bufferSize];
94     }
95 
96     /**
97      * Re-initializes the object for the new output.
98      */
setOutput(@onNull OutputStream out)99     protected void setOutput(@NonNull OutputStream out) {
100         if (mOut != null) {
101             throw new IllegalStateException("setOutput() called before calling release()");
102         }
103 
104         mOut = Objects.requireNonNull(out);
105         mBufferPos = 0;
106         mStringRefs.clear();
107     }
108 
drain()109     protected void drain() throws IOException {
110         if (mBufferPos > 0) {
111             mOut.write(mBuffer, 0, mBufferPos);
112             mBufferPos = 0;
113         }
114     }
115 
116     @Override
flush()117     public void flush() throws IOException {
118         drain();
119         mOut.flush();
120     }
121 
122     @Override
close()123     public void close() throws IOException {
124         mOut.close();
125         release();
126     }
127 
128     @Override
write(int b)129     public void write(int b) throws IOException {
130         writeByte(b);
131     }
132 
133     @Override
write(byte[] b)134     public void write(byte[] b) throws IOException {
135         write(b, 0, b.length);
136     }
137 
138     @Override
write(byte[] b, int off, int len)139     public void write(byte[] b, int off, int len) throws IOException {
140         if (mBufferCap < len) {
141             drain();
142             mOut.write(b, off, len);
143         } else {
144             if (mBufferCap - mBufferPos < len) drain();
145             System.arraycopy(b, off, mBuffer, mBufferPos, len);
146             mBufferPos += len;
147         }
148     }
149 
150     @Override
writeUTF(String s)151     public void writeUTF(String s) throws IOException {
152         final int len = (int) ModifiedUtf8.countBytes(s, false);
153         if (len > MAX_UNSIGNED_SHORT) {
154             throw new IOException("Modified UTF-8 length too large: " + len);
155         }
156 
157         // Attempt to write directly to buffer space if there's enough room,
158         // otherwise fall back to chunking into place
159         if (mBufferCap >= 2 + len) {
160             if (mBufferCap - mBufferPos < 2 + len) drain();
161             writeShort(len);
162             ModifiedUtf8.encode(mBuffer, mBufferPos, s);
163             mBufferPos += len;
164         } else {
165             final byte[] tmp = newByteArray(len + 1);
166             ModifiedUtf8.encode(tmp, 0, s);
167             writeShort(len);
168             write(tmp, 0, len);
169         }
170     }
171 
172     /**
173      * Write a {@link String} value with the additional signal that the given
174      * value is a candidate for being canonicalized, similar to
175      * {@link String#intern()}.
176      * <p>
177      * Canonicalization is implemented by writing each unique string value once
178      * the first time it appears, and then writing a lightweight {@code short}
179      * reference when that string is written again in the future.
180      *
181      * @see FastDataInput#readInternedUTF()
182      */
writeInternedUTF(@onNull String s)183     public void writeInternedUTF(@NonNull String s) throws IOException {
184         Integer ref = mStringRefs.get(s);
185         if (ref != null) {
186             writeShort(ref);
187         } else {
188             writeShort(MAX_UNSIGNED_SHORT);
189             writeUTF(s);
190 
191             // We can only safely intern when we have remaining values; if we're
192             // full we at least sent the string value above
193             ref = mStringRefs.size();
194             if (ref < MAX_UNSIGNED_SHORT) {
195                 mStringRefs.put(s, ref);
196             }
197         }
198     }
199 
200     @Override
writeBoolean(boolean v)201     public void writeBoolean(boolean v) throws IOException {
202         writeByte(v ? 1 : 0);
203     }
204 
205     @Override
writeByte(int v)206     public void writeByte(int v) throws IOException {
207         if (mBufferCap - mBufferPos < 1) drain();
208         mBuffer[mBufferPos++] = (byte) ((v >>  0) & 0xff);
209     }
210 
211     @Override
writeShort(int v)212     public void writeShort(int v) throws IOException {
213         if (mBufferCap - mBufferPos < 2) drain();
214         mBuffer[mBufferPos++] = (byte) ((v >>  8) & 0xff);
215         mBuffer[mBufferPos++] = (byte) ((v >>  0) & 0xff);
216     }
217 
218     @Override
writeChar(int v)219     public void writeChar(int v) throws IOException {
220         writeShort((short) v);
221     }
222 
223     @Override
writeInt(int v)224     public void writeInt(int v) throws IOException {
225         if (mBufferCap - mBufferPos < 4) drain();
226         mBuffer[mBufferPos++] = (byte) ((v >> 24) & 0xff);
227         mBuffer[mBufferPos++] = (byte) ((v >> 16) & 0xff);
228         mBuffer[mBufferPos++] = (byte) ((v >>  8) & 0xff);
229         mBuffer[mBufferPos++] = (byte) ((v >>  0) & 0xff);
230     }
231 
232     @Override
writeLong(long v)233     public void writeLong(long v) throws IOException {
234         if (mBufferCap - mBufferPos < 8) drain();
235         int i = (int) (v >> 32);
236         mBuffer[mBufferPos++] = (byte) ((i >> 24) & 0xff);
237         mBuffer[mBufferPos++] = (byte) ((i >> 16) & 0xff);
238         mBuffer[mBufferPos++] = (byte) ((i >>  8) & 0xff);
239         mBuffer[mBufferPos++] = (byte) ((i >>  0) & 0xff);
240         i = (int) v;
241         mBuffer[mBufferPos++] = (byte) ((i >> 24) & 0xff);
242         mBuffer[mBufferPos++] = (byte) ((i >> 16) & 0xff);
243         mBuffer[mBufferPos++] = (byte) ((i >>  8) & 0xff);
244         mBuffer[mBufferPos++] = (byte) ((i >>  0) & 0xff);
245     }
246 
247     @Override
writeFloat(float v)248     public void writeFloat(float v) throws IOException {
249         writeInt(Float.floatToIntBits(v));
250     }
251 
252     @Override
writeDouble(double v)253     public void writeDouble(double v) throws IOException {
254         writeLong(Double.doubleToLongBits(v));
255     }
256 
257     @Override
writeBytes(String s)258     public void writeBytes(String s) throws IOException {
259         // Callers should use writeUTF()
260         throw new UnsupportedOperationException();
261     }
262 
263     @Override
writeChars(String s)264     public void writeChars(String s) throws IOException {
265         // Callers should use writeUTF()
266         throw new UnsupportedOperationException();
267     }
268 }
269