1/*
2 * Copyright (C) 2015 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#include "BigBuffer.h"
18#include "StringPiece.h"
19#include "StringPool.h"
20#include "Util.h"
21
22#include <algorithm>
23#include <androidfw/ResourceTypes.h>
24#include <memory>
25#include <string>
26
27namespace aapt {
28
29StringPool::Ref::Ref() : mEntry(nullptr) {
30}
31
32StringPool::Ref::Ref(const StringPool::Ref& rhs) : mEntry(rhs.mEntry) {
33    if (mEntry != nullptr) {
34        mEntry->ref++;
35    }
36}
37
38StringPool::Ref::Ref(StringPool::Entry* entry) : mEntry(entry) {
39    if (mEntry != nullptr) {
40        mEntry->ref++;
41    }
42}
43
44StringPool::Ref::~Ref() {
45    if (mEntry != nullptr) {
46        mEntry->ref--;
47    }
48}
49
50StringPool::Ref& StringPool::Ref::operator=(const StringPool::Ref& rhs) {
51    if (rhs.mEntry != nullptr) {
52        rhs.mEntry->ref++;
53    }
54
55    if (mEntry != nullptr) {
56        mEntry->ref--;
57    }
58    mEntry = rhs.mEntry;
59    return *this;
60}
61
62const std::u16string* StringPool::Ref::operator->() const {
63    return &mEntry->value;
64}
65
66const std::u16string& StringPool::Ref::operator*() const {
67    return mEntry->value;
68}
69
70size_t StringPool::Ref::getIndex() const {
71    return mEntry->index;
72}
73
74const StringPool::Context& StringPool::Ref::getContext() const {
75    return mEntry->context;
76}
77
78StringPool::StyleRef::StyleRef() : mEntry(nullptr) {
79}
80
81StringPool::StyleRef::StyleRef(const StringPool::StyleRef& rhs) : mEntry(rhs.mEntry) {
82    if (mEntry != nullptr) {
83        mEntry->ref++;
84    }
85}
86
87StringPool::StyleRef::StyleRef(StringPool::StyleEntry* entry) : mEntry(entry) {
88    if (mEntry != nullptr) {
89        mEntry->ref++;
90    }
91}
92
93StringPool::StyleRef::~StyleRef() {
94    if (mEntry != nullptr) {
95        mEntry->ref--;
96    }
97}
98
99StringPool::StyleRef& StringPool::StyleRef::operator=(const StringPool::StyleRef& rhs) {
100    if (rhs.mEntry != nullptr) {
101        rhs.mEntry->ref++;
102    }
103
104    if (mEntry != nullptr) {
105        mEntry->ref--;
106    }
107    mEntry = rhs.mEntry;
108    return *this;
109}
110
111const StringPool::StyleEntry* StringPool::StyleRef::operator->() const {
112    return mEntry;
113}
114
115const StringPool::StyleEntry& StringPool::StyleRef::operator*() const {
116    return *mEntry;
117}
118
119size_t StringPool::StyleRef::getIndex() const {
120    return mEntry->str.getIndex();
121}
122
123const StringPool::Context& StringPool::StyleRef::getContext() const {
124    return mEntry->str.getContext();
125}
126
127StringPool::Ref StringPool::makeRef(const StringPiece16& str) {
128    return makeRefImpl(str, Context{}, true);
129}
130
131StringPool::Ref StringPool::makeRef(const StringPiece16& str, const Context& context) {
132    return makeRefImpl(str, context, true);
133}
134
135StringPool::Ref StringPool::makeRefImpl(const StringPiece16& str, const Context& context,
136        bool unique) {
137    if (unique) {
138        auto iter = mIndexedStrings.find(str);
139        if (iter != std::end(mIndexedStrings)) {
140            return Ref(iter->second);
141        }
142    }
143
144    Entry* entry = new Entry();
145    entry->value = str.toString();
146    entry->context = context;
147    entry->index = mStrings.size();
148    entry->ref = 0;
149    mStrings.emplace_back(entry);
150    mIndexedStrings.insert(std::make_pair(StringPiece16(entry->value), entry));
151    return Ref(entry);
152}
153
154StringPool::StyleRef StringPool::makeRef(const StyleString& str) {
155    return makeRef(str, Context{});
156}
157
158StringPool::StyleRef StringPool::makeRef(const StyleString& str, const Context& context) {
159    Entry* entry = new Entry();
160    entry->value = str.str;
161    entry->context = context;
162    entry->index = mStrings.size();
163    entry->ref = 0;
164    mStrings.emplace_back(entry);
165    mIndexedStrings.insert(std::make_pair(StringPiece16(entry->value), entry));
166
167    StyleEntry* styleEntry = new StyleEntry();
168    styleEntry->str = Ref(entry);
169    for (const aapt::Span& span : str.spans) {
170        styleEntry->spans.emplace_back(Span{makeRef(span.name),
171                span.firstChar, span.lastChar});
172    }
173    styleEntry->ref = 0;
174    mStyles.emplace_back(styleEntry);
175    return StyleRef(styleEntry);
176}
177
178StringPool::StyleRef StringPool::makeRef(const StyleRef& ref) {
179    Entry* entry = new Entry();
180    entry->value = *ref.mEntry->str;
181    entry->context = ref.mEntry->str.mEntry->context;
182    entry->index = mStrings.size();
183    entry->ref = 0;
184    mStrings.emplace_back(entry);
185    mIndexedStrings.insert(std::make_pair(StringPiece16(entry->value), entry));
186
187    StyleEntry* styleEntry = new StyleEntry();
188    styleEntry->str = Ref(entry);
189    for (const Span& span : ref.mEntry->spans) {
190        styleEntry->spans.emplace_back(Span{ makeRef(*span.name), span.firstChar, span.lastChar });
191    }
192    styleEntry->ref = 0;
193    mStyles.emplace_back(styleEntry);
194    return StyleRef(styleEntry);
195}
196
197void StringPool::merge(StringPool&& pool) {
198    mIndexedStrings.insert(pool.mIndexedStrings.begin(), pool.mIndexedStrings.end());
199    pool.mIndexedStrings.clear();
200    std::move(pool.mStrings.begin(), pool.mStrings.end(), std::back_inserter(mStrings));
201    pool.mStrings.clear();
202    std::move(pool.mStyles.begin(), pool.mStyles.end(), std::back_inserter(mStyles));
203    pool.mStyles.clear();
204
205    // Assign the indices.
206    const size_t len = mStrings.size();
207    for (size_t index = 0; index < len; index++) {
208        mStrings[index]->index = index;
209    }
210}
211
212void StringPool::hintWillAdd(size_t stringCount, size_t styleCount) {
213    mStrings.reserve(mStrings.size() + stringCount);
214    mStyles.reserve(mStyles.size() + styleCount);
215}
216
217void StringPool::prune() {
218    const auto iterEnd = std::end(mIndexedStrings);
219    auto indexIter = std::begin(mIndexedStrings);
220    while (indexIter != iterEnd) {
221        if (indexIter->second->ref <= 0) {
222            mIndexedStrings.erase(indexIter++);
223        } else {
224            ++indexIter;
225        }
226    }
227
228    auto endIter2 = std::remove_if(std::begin(mStrings), std::end(mStrings),
229            [](const std::unique_ptr<Entry>& entry) -> bool {
230                return entry->ref <= 0;
231            }
232    );
233
234    auto endIter3 = std::remove_if(std::begin(mStyles), std::end(mStyles),
235            [](const std::unique_ptr<StyleEntry>& entry) -> bool {
236                return entry->ref <= 0;
237            }
238    );
239
240    // Remove the entries at the end or else we'll be accessing
241    // a deleted string from the StyleEntry.
242    mStrings.erase(endIter2, std::end(mStrings));
243    mStyles.erase(endIter3, std::end(mStyles));
244}
245
246void StringPool::sort(const std::function<bool(const Entry&, const Entry&)>& cmp) {
247    std::sort(std::begin(mStrings), std::end(mStrings),
248            [&cmp](const std::unique_ptr<Entry>& a, const std::unique_ptr<Entry>& b) -> bool {
249                return cmp(*a, *b);
250            }
251    );
252
253    // Assign the indices.
254    const size_t len = mStrings.size();
255    for (size_t index = 0; index < len; index++) {
256        mStrings[index]->index = index;
257    }
258
259    // Reorder the styles.
260    std::sort(std::begin(mStyles), std::end(mStyles),
261            [](const std::unique_ptr<StyleEntry>& lhs,
262               const std::unique_ptr<StyleEntry>& rhs) -> bool {
263                return lhs->str.getIndex() < rhs->str.getIndex();
264            }
265    );
266}
267
268template <typename T>
269static T* encodeLength(T* data, size_t length) {
270    static_assert(std::is_integral<T>::value, "wat.");
271
272    constexpr size_t kMask = 1 << ((sizeof(T) * 8) - 1);
273    constexpr size_t kMaxSize = kMask - 1;
274    if (length > kMaxSize) {
275        *data++ = kMask | (kMaxSize & (length >> (sizeof(T) * 8)));
276    }
277    *data++ = length;
278    return data;
279}
280
281template <typename T>
282static size_t encodedLengthUnits(size_t length) {
283    static_assert(std::is_integral<T>::value, "wat.");
284
285    constexpr size_t kMask = 1 << ((sizeof(T) * 8) - 1);
286    constexpr size_t kMaxSize = kMask - 1;
287    return length > kMaxSize ? 2 : 1;
288}
289
290
291bool StringPool::flatten(BigBuffer* out, const StringPool& pool, bool utf8) {
292    const size_t startIndex = out->size();
293    android::ResStringPool_header* header = out->nextBlock<android::ResStringPool_header>();
294    header->header.type = android::RES_STRING_POOL_TYPE;
295    header->header.headerSize = sizeof(*header);
296    header->stringCount = pool.size();
297    if (utf8) {
298        header->flags |= android::ResStringPool_header::UTF8_FLAG;
299    }
300
301    uint32_t* indices = pool.size() != 0 ? out->nextBlock<uint32_t>(pool.size()) : nullptr;
302
303    uint32_t* styleIndices = nullptr;
304    if (!pool.mStyles.empty()) {
305        header->styleCount = pool.mStyles.back()->str.getIndex() + 1;
306        styleIndices = out->nextBlock<uint32_t>(header->styleCount);
307    }
308
309    const size_t beforeStringsIndex = out->size();
310    header->stringsStart = beforeStringsIndex - startIndex;
311
312    for (const auto& entry : pool) {
313        *indices = out->size() - beforeStringsIndex;
314        indices++;
315
316        if (utf8) {
317            std::string encoded = util::utf16ToUtf8(entry->value);
318
319            const size_t totalSize = encodedLengthUnits<char>(entry->value.size())
320                    + encodedLengthUnits<char>(encoded.length())
321                    + encoded.size() + 1;
322
323            char* data = out->nextBlock<char>(totalSize);
324
325            // First encode the actual UTF16 string length.
326            data = encodeLength(data, entry->value.size());
327
328            // Now encode the size of the converted UTF8 string.
329            data = encodeLength(data, encoded.length());
330            strncpy(data, encoded.data(), encoded.size());
331        } else {
332            const size_t totalSize = encodedLengthUnits<char16_t>(entry->value.size())
333                    + entry->value.size() + 1;
334
335            char16_t* data = out->nextBlock<char16_t>(totalSize);
336
337            // Encode the actual UTF16 string length.
338            data = encodeLength(data, entry->value.size());
339            strncpy16(data, entry->value.data(), entry->value.size());
340        }
341    }
342
343    out->align4();
344
345    if (!pool.mStyles.empty()) {
346        const size_t beforeStylesIndex = out->size();
347        header->stylesStart = beforeStylesIndex - startIndex;
348
349        size_t currentIndex = 0;
350        for (const auto& entry : pool.mStyles) {
351            while (entry->str.getIndex() > currentIndex) {
352                styleIndices[currentIndex++] = out->size() - beforeStylesIndex;
353
354                uint32_t* spanOffset = out->nextBlock<uint32_t>();
355                *spanOffset = android::ResStringPool_span::END;
356            }
357            styleIndices[currentIndex++] = out->size() - beforeStylesIndex;
358
359            android::ResStringPool_span* span =
360                    out->nextBlock<android::ResStringPool_span>(entry->spans.size());
361            for (const auto& s : entry->spans) {
362                span->name.index = s.name.getIndex();
363                span->firstChar = s.firstChar;
364                span->lastChar = s.lastChar;
365                span++;
366            }
367
368            uint32_t* spanEnd = out->nextBlock<uint32_t>();
369            *spanEnd = android::ResStringPool_span::END;
370        }
371
372        // The error checking code in the platform looks for an entire
373        // ResStringPool_span structure worth of 0xFFFFFFFF at the end
374        // of the style block, so fill in the remaining 2 32bit words
375        // with 0xFFFFFFFF.
376        const size_t paddingLength = sizeof(android::ResStringPool_span)
377                - sizeof(android::ResStringPool_span::name);
378        uint8_t* padding = out->nextBlock<uint8_t>(paddingLength);
379        memset(padding, 0xff, paddingLength);
380        out->align4();
381    }
382    header->header.size = out->size() - startIndex;
383    return true;
384}
385
386bool StringPool::flattenUtf8(BigBuffer* out, const StringPool& pool) {
387    return flatten(out, pool, true);
388}
389
390bool StringPool::flattenUtf16(BigBuffer* out, const StringPool& pool) {
391    return flatten(out, pool, false);
392}
393
394} // namespace aapt
395