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