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.pool;
33
34import org.jf.dexlib2.Opcodes;
35import org.jf.dexlib2.ValueType;
36import org.jf.dexlib2.iface.Annotation;
37import org.jf.dexlib2.iface.AnnotationElement;
38import org.jf.dexlib2.iface.ClassDef;
39import org.jf.dexlib2.iface.Field;
40import org.jf.dexlib2.iface.reference.*;
41import org.jf.dexlib2.iface.value.*;
42import org.jf.dexlib2.writer.*;
43import org.jf.dexlib2.writer.io.DexDataStore;
44import org.jf.dexlib2.writer.io.FileDataStore;
45import org.jf.util.ExceptionWithContext;
46
47import javax.annotation.Nonnull;
48import java.io.File;
49import java.io.IOException;
50import java.util.Collection;
51import java.util.Set;
52
53public class DexPool extends DexWriter<CharSequence, StringReference, CharSequence, TypeReference,
54        MethodProtoReference, FieldReference, MethodReference, PoolClassDef,
55        Annotation, Set<? extends Annotation>,
56        TypeListPool.Key<? extends Collection<? extends CharSequence>>, Field, PoolMethod,
57        EncodedValue, AnnotationElement, StringPool, TypePool, ProtoPool, FieldPool, MethodPool, ClassPool,
58        TypeListPool, AnnotationPool, AnnotationSetPool> {
59
60    private final Markable[] sections = new Markable[] {
61            stringSection, typeSection, protoSection, fieldSection, methodSection, classSection, typeListSection,
62            annotationSection, annotationSetSection
63    };
64
65    public DexPool(Opcodes opcodes) {
66        super(opcodes);
67    }
68
69    @Nonnull @Override protected SectionProvider getSectionProvider() {
70        return new DexPoolSectionProvider();
71    }
72
73    public static void writeTo(@Nonnull DexDataStore dataStore, @Nonnull org.jf.dexlib2.iface.DexFile input)
74            throws IOException {
75        DexPool dexPool = new DexPool(input.getOpcodes());
76        for (ClassDef classDef: input.getClasses()) {
77            dexPool.internClass(classDef);
78        }
79        dexPool.writeTo(dataStore);
80    }
81
82    public static void writeTo(@Nonnull String path, @Nonnull org.jf.dexlib2.iface.DexFile input) throws IOException {
83        DexPool dexPool = new DexPool(input.getOpcodes());
84        for (ClassDef classDef: input.getClasses()) {
85            dexPool.internClass(classDef);
86        }
87        dexPool.writeTo(new FileDataStore(new File(path)));
88    }
89
90    /**
91     * Interns a class into this DexPool
92     * @param classDef The class to intern
93     */
94    public void internClass(ClassDef classDef) {
95        classSection.intern(classDef);
96    }
97
98    /**
99     * Creates a marked state that can be returned to by calling reset()
100     *
101     * This is useful to rollback the last added class if it causes a method/field/type overflow
102     */
103    public void mark() {
104        for (Markable section: sections) {
105            section.mark();
106        }
107    }
108
109    /**
110     * Resets to the last marked state
111     *
112     * This is useful to rollback the last added class if it causes a method/field/type overflow
113     */
114    public void reset() {
115        for (Markable section: sections) {
116            section.reset();
117        }
118    }
119
120    @Override protected void writeEncodedValue(@Nonnull InternalEncodedValueWriter writer,
121                                               @Nonnull EncodedValue encodedValue) throws IOException {
122        switch (encodedValue.getValueType()) {
123            case ValueType.ANNOTATION:
124                AnnotationEncodedValue annotationEncodedValue = (AnnotationEncodedValue)encodedValue;
125                writer.writeAnnotation(annotationEncodedValue.getType(), annotationEncodedValue.getElements());
126                break;
127            case ValueType.ARRAY:
128                ArrayEncodedValue arrayEncodedValue = (ArrayEncodedValue)encodedValue;
129                writer.writeArray(arrayEncodedValue.getValue());
130                break;
131            case ValueType.BOOLEAN:
132                writer.writeBoolean(((BooleanEncodedValue)encodedValue).getValue());
133                break;
134            case ValueType.BYTE:
135                writer.writeByte(((ByteEncodedValue)encodedValue).getValue());
136                break;
137            case ValueType.CHAR:
138                writer.writeChar(((CharEncodedValue)encodedValue).getValue());
139                break;
140            case ValueType.DOUBLE:
141                writer.writeDouble(((DoubleEncodedValue)encodedValue).getValue());
142                break;
143            case ValueType.ENUM:
144                writer.writeEnum(((EnumEncodedValue)encodedValue).getValue());
145                break;
146            case ValueType.FIELD:
147                writer.writeField(((FieldEncodedValue)encodedValue).getValue());
148                break;
149            case ValueType.FLOAT:
150                writer.writeFloat(((FloatEncodedValue)encodedValue).getValue());
151                break;
152            case ValueType.INT:
153                writer.writeInt(((IntEncodedValue)encodedValue).getValue());
154                break;
155            case ValueType.LONG:
156                writer.writeLong(((LongEncodedValue)encodedValue).getValue());
157                break;
158            case ValueType.METHOD:
159                writer.writeMethod(((MethodEncodedValue)encodedValue).getValue());
160                break;
161            case ValueType.NULL:
162                writer.writeNull();
163                break;
164            case ValueType.SHORT:
165                writer.writeShort(((ShortEncodedValue)encodedValue).getValue());
166                break;
167            case ValueType.STRING:
168                writer.writeString(((StringEncodedValue)encodedValue).getValue());
169                break;
170            case ValueType.TYPE:
171                writer.writeType(((TypeEncodedValue)encodedValue).getValue());
172                break;
173            default:
174                throw new ExceptionWithContext("Unrecognized value type: %d", encodedValue.getValueType());
175        }
176    }
177
178    void internEncodedValue(@Nonnull EncodedValue encodedValue) {
179        switch (encodedValue.getValueType()) {
180            case ValueType.ANNOTATION:
181                AnnotationEncodedValue annotationEncodedValue = (AnnotationEncodedValue)encodedValue;
182                typeSection.intern(annotationEncodedValue.getType());
183                for (AnnotationElement element: annotationEncodedValue.getElements()) {
184                    stringSection.intern(element.getName());
185                    internEncodedValue(element.getValue());
186                }
187                break;
188            case ValueType.ARRAY:
189                for (EncodedValue element: ((ArrayEncodedValue)encodedValue).getValue()) {
190                    internEncodedValue(element);
191                }
192                break;
193            case ValueType.STRING:
194                stringSection.intern(((StringEncodedValue)encodedValue).getValue());
195                break;
196            case ValueType.TYPE:
197                typeSection.intern(((TypeEncodedValue)encodedValue).getValue());
198                break;
199            case ValueType.ENUM:
200                fieldSection.intern(((EnumEncodedValue)encodedValue).getValue());
201                break;
202            case ValueType.FIELD:
203                fieldSection.intern(((FieldEncodedValue)encodedValue).getValue());
204                break;
205            case ValueType.METHOD:
206                methodSection.intern(((MethodEncodedValue)encodedValue).getValue());
207                break;
208        }
209    }
210
211    protected class DexPoolSectionProvider extends SectionProvider {
212        @Nonnull @Override public StringPool getStringSection() {
213            return new StringPool(DexPool.this);
214        }
215
216        @Nonnull @Override public TypePool getTypeSection() {
217            return new TypePool(DexPool.this);
218        }
219
220        @Nonnull @Override public ProtoPool getProtoSection() {
221            return new ProtoPool(DexPool.this);
222        }
223
224        @Nonnull @Override public FieldPool getFieldSection() {
225            return new FieldPool(DexPool.this);
226        }
227
228        @Nonnull @Override public MethodPool getMethodSection() {
229            return new MethodPool(DexPool.this);
230        }
231
232        @Nonnull @Override public ClassPool getClassSection() {
233            return new ClassPool(DexPool.this);
234        }
235
236        @Nonnull @Override public TypeListPool getTypeListSection() {
237            return new TypeListPool(DexPool.this);
238        }
239
240        @Nonnull @Override public AnnotationPool getAnnotationSection() {
241            return new AnnotationPool(DexPool.this);
242        }
243
244        @Nonnull @Override public AnnotationSetPool getAnnotationSetSection() {
245            return new AnnotationSetPool(DexPool.this);
246        }
247    }
248}
249