1/*
2 * Copyright 2013, Google Inc.
3 * All rights reserved.
4 *
5 * Redistribution and use in source and binary forms, with or without
6 * modification, are permitted provided that the following conditions are
7 * met:
8 *
9 *     * Redistributions of source code must retain the above copyright
10 * notice, this list of conditions and the following disclaimer.
11 *     * Redistributions in binary form must reproduce the above
12 * copyright notice, this list of conditions and the following disclaimer
13 * in the documentation and/or other materials provided with the
14 * distribution.
15 *     * Neither the name of Google Inc. nor the names of its
16 * contributors may be used to endorse or promote products derived from
17 * this software without specific prior written permission.
18 *
19 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30 */
31
32package org.jf.dexlib2.writer;
33
34import com.google.common.collect.Ordering;
35import com.google.common.primitives.Ints;
36import org.jf.dexlib2.Opcode;
37import org.jf.dexlib2.Opcodes;
38import org.jf.dexlib2.ReferenceType;
39import org.jf.dexlib2.iface.instruction.DualReferenceInstruction;
40import org.jf.dexlib2.iface.instruction.ReferenceInstruction;
41import org.jf.dexlib2.iface.instruction.SwitchElement;
42import org.jf.dexlib2.iface.instruction.formats.*;
43import org.jf.dexlib2.iface.reference.FieldReference;
44import org.jf.dexlib2.iface.reference.MethodProtoReference;
45import org.jf.dexlib2.iface.reference.MethodReference;
46import org.jf.dexlib2.iface.reference.Reference;
47import org.jf.dexlib2.iface.reference.StringReference;
48import org.jf.dexlib2.iface.reference.TypeReference;
49import org.jf.util.ExceptionWithContext;
50
51import javax.annotation.Nonnull;
52import java.io.IOException;
53import java.util.Comparator;
54import java.util.List;
55
56public class InstructionWriter<StringRef extends StringReference, TypeRef extends TypeReference,
57        FieldRefKey extends FieldReference, MethodRefKey extends MethodReference,
58        ProtoRefKey extends MethodProtoReference> {
59    @Nonnull private final Opcodes opcodes;
60    @Nonnull private final DexDataWriter writer;
61    @Nonnull private final StringSection<?, StringRef> stringSection;
62    @Nonnull private final TypeSection<?, ?, TypeRef> typeSection;
63    @Nonnull private final FieldSection<?, ?, FieldRefKey, ?> fieldSection;
64    @Nonnull private final MethodSection<?, ?, ?, MethodRefKey, ?> methodSection;
65    @Nonnull private final ProtoSection<?, ?, ProtoRefKey, ?> protoSection;
66
67    @Nonnull static <StringRef extends StringReference, TypeRef extends TypeReference, FieldRefKey extends FieldReference,
68            MethodRefKey extends MethodReference, ProtoRefKey extends MethodProtoReference>
69            InstructionWriter<StringRef, TypeRef, FieldRefKey, MethodRefKey, ProtoRefKey>
70            makeInstructionWriter(
71                @Nonnull Opcodes opcodes,
72                @Nonnull DexDataWriter writer,
73                @Nonnull StringSection<?, StringRef> stringSection,
74                @Nonnull TypeSection<?, ?, TypeRef> typeSection,
75                @Nonnull FieldSection<?, ?, FieldRefKey, ?> fieldSection,
76                @Nonnull MethodSection<?, ?, ?, MethodRefKey, ?> methodSection,
77                @Nonnull ProtoSection<?, ?, ProtoRefKey, ?> protoSection) {
78        return new InstructionWriter<StringRef, TypeRef, FieldRefKey, MethodRefKey, ProtoRefKey>(
79                opcodes, writer, stringSection, typeSection, fieldSection, methodSection, protoSection);
80    }
81
82    InstructionWriter(@Nonnull Opcodes opcodes,
83                      @Nonnull DexDataWriter writer,
84                      @Nonnull StringSection<?, StringRef> stringSection,
85                      @Nonnull TypeSection<?, ?, TypeRef> typeSection,
86                      @Nonnull FieldSection<?, ?, FieldRefKey, ?> fieldSection,
87                      @Nonnull MethodSection<?, ?, ?, MethodRefKey, ?> methodSection,
88                      @Nonnull ProtoSection<?, ?, ProtoRefKey, ?> protoSection) {
89        this.opcodes = opcodes;
90        this.writer = writer;
91        this.stringSection = stringSection;
92        this.typeSection = typeSection;
93        this.fieldSection = fieldSection;
94        this.methodSection = methodSection;
95        this.protoSection = protoSection;
96    }
97
98    private short getOpcodeValue(Opcode opcode) {
99        Short value = opcodes.getOpcodeValue(opcode);
100        if (value == null) {
101            throw new ExceptionWithContext("Instruction %s is invalid for api %d", opcode.name, opcodes.api);
102        }
103        return value;
104    }
105
106    public void write(@Nonnull Instruction10t instruction) {
107        try {
108            writer.write(getOpcodeValue(instruction.getOpcode()));
109            writer.write(instruction.getCodeOffset());
110        } catch (IOException ex) {
111            throw new RuntimeException(ex);
112        }
113    }
114
115    public void write(@Nonnull Instruction10x instruction) {
116        try {
117            writer.write(getOpcodeValue(instruction.getOpcode()));
118            writer.write(0);
119        } catch (IOException ex) {
120            throw new RuntimeException(ex);
121        }
122    }
123
124    public void write(@Nonnull Instruction11n instruction) {
125        try {
126            writer.write(getOpcodeValue(instruction.getOpcode()));
127            writer.write(packNibbles(instruction.getRegisterA(), instruction.getNarrowLiteral()));
128        } catch (IOException ex) {
129            throw new RuntimeException(ex);
130        }
131    }
132
133    public void write(@Nonnull Instruction11x instruction) {
134        try {
135            writer.write(getOpcodeValue(instruction.getOpcode()));
136            writer.write(instruction.getRegisterA());
137        } catch (IOException ex) {
138            throw new RuntimeException(ex);
139        }
140    }
141
142    public void write(@Nonnull Instruction12x instruction) {
143        try {
144            writer.write(getOpcodeValue(instruction.getOpcode()));
145            writer.write(packNibbles(instruction.getRegisterA(), instruction.getRegisterB()));
146        } catch (IOException ex) {
147            throw new RuntimeException(ex);
148        }
149    }
150
151    public void write(@Nonnull Instruction20bc instruction) {
152        try {
153            writer.write(getOpcodeValue(instruction.getOpcode()));
154            writer.write(instruction.getVerificationError());
155            writer.writeUshort(getReferenceIndex(instruction));
156        } catch (IOException ex) {
157            throw new RuntimeException(ex);
158        }
159    }
160
161    public void write(@Nonnull Instruction20t instruction) {
162        try {
163            writer.write(getOpcodeValue(instruction.getOpcode()));
164            writer.write(0);
165            writer.writeShort(instruction.getCodeOffset());
166        } catch (IOException ex) {
167            throw new RuntimeException(ex);
168        }
169    }
170
171    public void write(@Nonnull Instruction21c instruction) {
172        try {
173            writer.write(getOpcodeValue(instruction.getOpcode()));
174            writer.write(instruction.getRegisterA());
175            writer.writeUshort(getReferenceIndex(instruction));
176        } catch (IOException ex) {
177            throw new RuntimeException(ex);
178        }
179    }
180
181    public void write(@Nonnull Instruction21ih instruction) {
182        try {
183            writer.write(getOpcodeValue(instruction.getOpcode()));
184            writer.write(instruction.getRegisterA());
185            writer.writeShort(instruction.getHatLiteral());
186        } catch (IOException ex) {
187            throw new RuntimeException(ex);
188        }
189    }
190
191    public void write(@Nonnull Instruction21lh instruction) {
192        try {
193            writer.write(getOpcodeValue(instruction.getOpcode()));
194            writer.write(instruction.getRegisterA());
195            writer.writeShort(instruction.getHatLiteral());
196        } catch (IOException ex) {
197            throw new RuntimeException(ex);
198        }
199    }
200
201    public void write(@Nonnull Instruction21s instruction) {
202        try {
203            writer.write(getOpcodeValue(instruction.getOpcode()));
204            writer.write(instruction.getRegisterA());
205            writer.writeShort(instruction.getNarrowLiteral());
206        } catch (IOException ex) {
207            throw new RuntimeException(ex);
208        }
209    }
210
211    public void write(@Nonnull Instruction21t instruction) {
212        try {
213            writer.write(getOpcodeValue(instruction.getOpcode()));
214            writer.write(instruction.getRegisterA());
215            writer.writeShort(instruction.getCodeOffset());
216        } catch (IOException ex) {
217            throw new RuntimeException(ex);
218        }
219    }
220
221    public void write(@Nonnull Instruction22b instruction) {
222        try {
223            writer.write(getOpcodeValue(instruction.getOpcode()));
224            writer.write(instruction.getRegisterA());
225            writer.write(instruction.getRegisterB());
226            writer.write(instruction.getNarrowLiteral());
227        } catch (IOException ex) {
228            throw new RuntimeException(ex);
229        }
230    }
231
232    public void write(@Nonnull Instruction22c instruction) {
233        try {
234            writer.write(getOpcodeValue(instruction.getOpcode()));
235            writer.write(packNibbles(instruction.getRegisterA(), instruction.getRegisterB()));
236            writer.writeUshort(getReferenceIndex(instruction));
237        } catch (IOException ex) {
238            throw new RuntimeException(ex);
239        }
240    }
241
242    public void write(@Nonnull Instruction22s instruction) {
243        try {
244            writer.write(getOpcodeValue(instruction.getOpcode()));
245            writer.write(packNibbles(instruction.getRegisterA(), instruction.getRegisterB()));
246            writer.writeShort(instruction.getNarrowLiteral());
247        } catch (IOException ex) {
248            throw new RuntimeException(ex);
249        }
250    }
251
252    public void write(@Nonnull Instruction22t instruction) {
253        try {
254            writer.write(getOpcodeValue(instruction.getOpcode()));
255            writer.write(packNibbles(instruction.getRegisterA(), instruction.getRegisterB()));
256            writer.writeShort(instruction.getCodeOffset());
257        } catch (IOException ex) {
258            throw new RuntimeException(ex);
259        }
260    }
261
262    public void write(@Nonnull Instruction22x instruction) {
263        try {
264            writer.write(getOpcodeValue(instruction.getOpcode()));
265            writer.write(instruction.getRegisterA());
266            writer.writeUshort(instruction.getRegisterB());
267        } catch (IOException ex) {
268            throw new RuntimeException(ex);
269        }
270    }
271
272    public void write(@Nonnull Instruction23x instruction) {
273        try {
274            writer.write(getOpcodeValue(instruction.getOpcode()));
275            writer.write(instruction.getRegisterA());
276            writer.write(instruction.getRegisterB());
277            writer.write(instruction.getRegisterC());
278        } catch (IOException ex) {
279            throw new RuntimeException(ex);
280        }
281    }
282
283    public void write(@Nonnull Instruction30t instruction) {
284        try {
285            writer.write(getOpcodeValue(instruction.getOpcode()));
286            writer.write(0);
287            writer.writeInt(instruction.getCodeOffset());
288        } catch (IOException ex) {
289            throw new RuntimeException(ex);
290        }
291    }
292
293    public void write(@Nonnull Instruction31c instruction) {
294        try {
295            writer.write(getOpcodeValue(instruction.getOpcode()));
296            writer.write(instruction.getRegisterA());
297            writer.writeInt(getReferenceIndex(instruction));
298        } catch (IOException ex) {
299            throw new RuntimeException(ex);
300        }
301    }
302
303    public void write(@Nonnull Instruction31i instruction) {
304        try {
305            writer.write(getOpcodeValue(instruction.getOpcode()));
306            writer.write(instruction.getRegisterA());
307            writer.writeInt(instruction.getNarrowLiteral());
308        } catch (IOException ex) {
309            throw new RuntimeException(ex);
310        }
311    }
312
313    public void write(@Nonnull Instruction31t instruction) {
314        try {
315            writer.write(getOpcodeValue(instruction.getOpcode()));
316            writer.write(instruction.getRegisterA());
317            writer.writeInt(instruction.getCodeOffset());
318        } catch (IOException ex) {
319            throw new RuntimeException(ex);
320        }
321    }
322
323    public void write(@Nonnull Instruction32x instruction) {
324        try {
325            writer.write(getOpcodeValue(instruction.getOpcode()));
326            writer.write(0);
327            writer.writeUshort(instruction.getRegisterA());
328            writer.writeUshort(instruction.getRegisterB());
329        } catch (IOException ex) {
330            throw new RuntimeException(ex);
331        }
332    }
333
334    public void write(@Nonnull Instruction35c instruction) {
335        try {
336            writer.write(getOpcodeValue(instruction.getOpcode()));
337            writer.write(packNibbles(instruction.getRegisterG(), instruction.getRegisterCount()));
338            writer.writeUshort(getReferenceIndex(instruction));
339            writer.write(packNibbles(instruction.getRegisterC(), instruction.getRegisterD()));
340            writer.write(packNibbles(instruction.getRegisterE(), instruction.getRegisterF()));
341        } catch (IOException ex) {
342            throw new RuntimeException(ex);
343        }
344    }
345
346    public void write(@Nonnull Instruction3rc instruction) {
347        try {
348            writer.write(getOpcodeValue(instruction.getOpcode()));
349            writer.write(instruction.getRegisterCount());
350            writer.writeUshort(getReferenceIndex(instruction));
351            writer.writeUshort(instruction.getStartRegister());
352        } catch (IOException ex) {
353            throw new RuntimeException(ex);
354        }
355    }
356
357    public void write(@Nonnull Instruction45cc instruction) {
358        try {
359            writer.write(getOpcodeValue(instruction.getOpcode()));
360            writer.write(packNibbles(instruction.getRegisterG(), instruction.getRegisterCount()));
361            writer.writeUshort(getReferenceIndex(instruction));
362            writer.write(packNibbles(instruction.getRegisterC(), instruction.getRegisterD()));
363            writer.write(packNibbles(instruction.getRegisterE(), instruction.getRegisterF()));
364            writer.writeUshort(getReference2Index(instruction));
365        } catch (IOException ex) {
366            throw new RuntimeException(ex);
367        }
368    }
369
370    public void write(@Nonnull Instruction4rcc instruction) {
371        try {
372            writer.write(getOpcodeValue(instruction.getOpcode()));
373            writer.write(instruction.getRegisterCount());
374            writer.writeUshort(getReferenceIndex(instruction));
375            writer.writeUshort(instruction.getStartRegister());
376            writer.writeUshort(getReference2Index(instruction));
377        } catch (IOException ex) {
378            throw new RuntimeException(ex);
379        }
380    }
381
382    public void write(@Nonnull Instruction51l instruction) {
383        try {
384            writer.write(getOpcodeValue(instruction.getOpcode()));
385            writer.write(instruction.getRegisterA());
386            writer.writeLong(instruction.getWideLiteral());
387        } catch (IOException ex) {
388            throw new RuntimeException(ex);
389        }
390    }
391
392    public void write(@Nonnull ArrayPayload instruction) {
393        try {
394            writer.writeUshort(getOpcodeValue(instruction.getOpcode()));
395            writer.writeUshort(instruction.getElementWidth());
396            List<Number> elements = instruction.getArrayElements();
397            writer.writeInt(elements.size());
398            switch (instruction.getElementWidth()) {
399                case 1:
400                    for (Number element: elements) {
401                        writer.write(element.byteValue());
402                    }
403                    break;
404                case 2:
405                    for (Number element: elements) {
406                        writer.writeShort(element.shortValue());
407                    }
408                    break;
409                case 4:
410                    for (Number element: elements) {
411                        writer.writeInt(element.intValue());
412                    }
413                    break;
414                case 8:
415                    for (Number element: elements) {
416                        writer.writeLong(element.longValue());
417                    }
418                    break;
419            }
420            if ((writer.getPosition() & 1) != 0) {
421                writer.write(0);
422            }
423        } catch (IOException ex) {
424            throw new RuntimeException(ex);
425        }
426    }
427
428    public void write(@Nonnull SparseSwitchPayload instruction) {
429        try {
430            writer.writeUbyte(0);
431            writer.writeUbyte(getOpcodeValue(instruction.getOpcode()) >> 8);
432            List<? extends SwitchElement> elements = Ordering.from(switchElementComparator).immutableSortedCopy(
433                    instruction.getSwitchElements());
434            writer.writeUshort(elements.size());
435            for (SwitchElement element: elements) {
436                writer.writeInt(element.getKey());
437            }
438            for (SwitchElement element: elements) {
439                writer.writeInt(element.getOffset());
440            }
441        } catch (IOException ex) {
442            throw new RuntimeException(ex);
443        }
444    }
445
446    private final Comparator<SwitchElement> switchElementComparator = new Comparator<SwitchElement>() {
447        @Override public int compare(SwitchElement element1, SwitchElement element2) {
448            return Ints.compare(element1.getKey(), element2.getKey());
449        }
450    };
451
452    public void write(@Nonnull PackedSwitchPayload instruction) {
453        try {
454            writer.writeUbyte(0);
455            writer.writeUbyte(getOpcodeValue(instruction.getOpcode()) >> 8);
456            List<? extends SwitchElement> elements = instruction.getSwitchElements();
457            writer.writeUshort(elements.size());
458            if (elements.size() == 0) {
459                writer.writeInt(0);
460            } else {
461                writer.writeInt(elements.get(0).getKey());
462                for (SwitchElement element: elements) {
463                    writer.writeInt(element.getOffset());
464                }
465            }
466        } catch (IOException ex) {
467            throw new RuntimeException(ex);
468        }
469    }
470
471    private static int packNibbles(int a, int b) {
472        return (b << 4) | a;
473    }
474
475    private int getReferenceIndex(ReferenceInstruction referenceInstruction) {
476        return getReferenceIndex(referenceInstruction.getReferenceType(),
477                referenceInstruction.getReference());
478    }
479
480    private int getReference2Index(DualReferenceInstruction referenceInstruction) {
481        return getReferenceIndex(referenceInstruction.getReferenceType2(),
482                referenceInstruction.getReference2());
483    }
484
485    private int getReferenceIndex(int referenceType, Reference reference) {
486        switch (referenceType) {
487            case ReferenceType.FIELD:
488                return fieldSection.getItemIndex((FieldRefKey) reference);
489            case ReferenceType.METHOD:
490                return methodSection.getItemIndex((MethodRefKey) reference);
491            case ReferenceType.STRING:
492                return stringSection.getItemIndex((StringRef) reference);
493            case ReferenceType.TYPE:
494                return typeSection.getItemIndex((TypeRef) reference);
495            case ReferenceType.METHOD_PROTO:
496                return protoSection.getItemIndex((ProtoRefKey) reference);
497            default:
498                throw new ExceptionWithContext("Unknown reference type: %d",  referenceType);
499        }
500    }
501}
502