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