1/*
2 * LZMA2OutputStream
3 *
4 * Authors: Lasse Collin <lasse.collin@tukaani.org>
5 *          Igor Pavlov <http://7-zip.org/>
6 *
7 * This file has been put into the public domain.
8 * You can do whatever you want with this file.
9 */
10
11package org.tukaani.xz;
12
13import java.io.DataOutputStream;
14import java.io.IOException;
15import org.tukaani.xz.lz.LZEncoder;
16import org.tukaani.xz.rangecoder.RangeEncoder;
17import org.tukaani.xz.lzma.LZMAEncoder;
18
19class LZMA2OutputStream extends FinishableOutputStream {
20    static final int COMPRESSED_SIZE_MAX = 64 << 10;
21
22    private FinishableOutputStream out;
23    private final DataOutputStream outData;
24
25    private final LZEncoder lz;
26    private final RangeEncoder rc;
27    private final LZMAEncoder lzma;
28
29    private final int props; // Cannot change props on the fly for now.
30    private boolean dictResetNeeded = true;
31    private boolean stateResetNeeded = true;
32    private boolean propsNeeded = true;
33
34    private int pendingSize = 0;
35    private boolean finished = false;
36    private IOException exception = null;
37
38    private final byte[] tempBuf = new byte[1];
39
40    private static int getExtraSizeBefore(int dictSize) {
41        return COMPRESSED_SIZE_MAX > dictSize
42               ? COMPRESSED_SIZE_MAX - dictSize : 0;
43    }
44
45    static int getMemoryUsage(LZMA2Options options) {
46        // 64 KiB buffer for the range encoder + a little extra + LZMAEncoder
47        int dictSize = options.getDictSize();
48        int extraSizeBefore = getExtraSizeBefore(dictSize);
49        return 70 + LZMAEncoder.getMemoryUsage(options.getMode(),
50                                               dictSize, extraSizeBefore,
51                                               options.getMatchFinder());
52    }
53
54    LZMA2OutputStream(FinishableOutputStream out, LZMA2Options options) {
55        if (out == null)
56            throw new NullPointerException();
57
58        this.out = out;
59        outData = new DataOutputStream(out);
60        rc = new RangeEncoder(COMPRESSED_SIZE_MAX);
61
62        int dictSize = options.getDictSize();
63        int extraSizeBefore = getExtraSizeBefore(dictSize);
64        lzma = LZMAEncoder.getInstance(rc,
65                options.getLc(), options.getLp(), options.getPb(),
66                options.getMode(),
67                dictSize, extraSizeBefore, options.getNiceLen(),
68                options.getMatchFinder(), options.getDepthLimit());
69
70        lz = lzma.getLZEncoder();
71
72        byte[] presetDict = options.getPresetDict();
73        if (presetDict != null && presetDict.length > 0) {
74            lz.setPresetDict(dictSize, presetDict);
75            dictResetNeeded = false;
76        }
77
78        props = (options.getPb() * 5 + options.getLp()) * 9 + options.getLc();
79    }
80
81    public void write(int b) throws IOException {
82        tempBuf[0] = (byte)b;
83        write(tempBuf, 0, 1);
84    }
85
86    public void write(byte[] buf, int off, int len) throws IOException {
87        if (off < 0 || len < 0 || off + len < 0 || off + len > buf.length)
88            throw new IndexOutOfBoundsException();
89
90        if (exception != null)
91            throw exception;
92
93        if (finished)
94            throw new XZIOException("Stream finished or closed");
95
96        try {
97            while (len > 0) {
98                int used = lz.fillWindow(buf, off, len);
99                off += used;
100                len -= used;
101                pendingSize += used;
102
103                if (lzma.encodeForLZMA2())
104                    writeChunk();
105            }
106        } catch (IOException e) {
107            exception = e;
108            throw e;
109        }
110    }
111
112    private void writeChunk() throws IOException {
113        int compressedSize = rc.finish();
114        int uncompressedSize = lzma.getUncompressedSize();
115
116        assert compressedSize > 0 : compressedSize;
117        assert uncompressedSize > 0 : uncompressedSize;
118
119        // +2 because the header of a compressed chunk is 2 bytes
120        // bigger than the header of an uncompressed chunk.
121        if (compressedSize + 2 < uncompressedSize) {
122            writeLZMA(uncompressedSize, compressedSize);
123        } else {
124            lzma.reset();
125            uncompressedSize = lzma.getUncompressedSize();
126            assert uncompressedSize > 0 : uncompressedSize;
127            writeUncompressed(uncompressedSize);
128        }
129
130        pendingSize -= uncompressedSize;
131        lzma.resetUncompressedSize();
132        rc.reset();
133    }
134
135    private void writeLZMA(int uncompressedSize, int compressedSize)
136            throws IOException {
137        int control;
138
139        if (propsNeeded) {
140            if (dictResetNeeded)
141                control = 0x80 + (3 << 5);
142            else
143                control = 0x80 + (2 << 5);
144        } else {
145            if (stateResetNeeded)
146                control = 0x80 + (1 << 5);
147            else
148                control = 0x80;
149        }
150
151        control |= (uncompressedSize - 1) >>> 16;
152        outData.writeByte(control);
153
154        outData.writeShort(uncompressedSize - 1);
155        outData.writeShort(compressedSize - 1);
156
157        if (propsNeeded)
158            outData.writeByte(props);
159
160        rc.write(out);
161
162        propsNeeded = false;
163        stateResetNeeded = false;
164        dictResetNeeded = false;
165    }
166
167    private void writeUncompressed(int uncompressedSize) throws IOException {
168        while (uncompressedSize > 0) {
169            int chunkSize = Math.min(uncompressedSize, COMPRESSED_SIZE_MAX);
170            outData.writeByte(dictResetNeeded ? 0x01 : 0x02);
171            outData.writeShort(chunkSize - 1);
172            lz.copyUncompressed(out, uncompressedSize, chunkSize);
173            uncompressedSize -= chunkSize;
174            dictResetNeeded = false;
175        }
176
177        stateResetNeeded = true;
178    }
179
180    private void writeEndMarker() throws IOException {
181        assert !finished;
182
183        if (exception != null)
184            throw exception;
185
186        lz.setFinishing();
187
188        try {
189            while (pendingSize > 0) {
190                lzma.encodeForLZMA2();
191                writeChunk();
192            }
193
194            out.write(0x00);
195        } catch (IOException e) {
196            exception = e;
197            throw e;
198        }
199
200        finished = true;
201    }
202
203    public void flush() throws IOException {
204        if (exception != null)
205            throw exception;
206
207        if (finished)
208            throw new XZIOException("Stream finished or closed");
209
210        try {
211            lz.setFlushing();
212
213            while (pendingSize > 0) {
214                lzma.encodeForLZMA2();
215                writeChunk();
216            }
217
218            out.flush();
219        } catch (IOException e) {
220            exception = e;
221            throw e;
222        }
223    }
224
225    public void finish() throws IOException {
226        if (!finished) {
227            writeEndMarker();
228
229            try {
230                out.finish();
231            } catch (IOException e) {
232                exception = e;
233                throw e;
234            }
235
236            finished = true;
237        }
238    }
239
240    public void close() throws IOException {
241        if (out != null) {
242            if (!finished) {
243                try {
244                    writeEndMarker();
245                } catch (IOException e) {}
246            }
247
248            try {
249                out.close();
250            } catch (IOException e) {
251                if (exception == null)
252                    exception = e;
253            }
254
255            out = null;
256        }
257
258        if (exception != null)
259            throw exception;
260    }
261}
262