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 "StringPool.h"
18#include "util/BigBuffer.h"
19#include "util/StringPiece.h"
20#include "util/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            indexIter = 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    // Reassign the indices.
246    const size_t len = mStrings.size();
247    for (size_t index = 0; index < len; index++) {
248        mStrings[index]->index = index;
249    }
250}
251
252void StringPool::sort(const std::function<bool(const Entry&, const Entry&)>& cmp) {
253    std::sort(std::begin(mStrings), std::end(mStrings),
254            [&cmp](const std::unique_ptr<Entry>& a, const std::unique_ptr<Entry>& b) -> bool {
255                return cmp(*a, *b);
256            }
257    );
258
259    // Assign the indices.
260    const size_t len = mStrings.size();
261    for (size_t index = 0; index < len; index++) {
262        mStrings[index]->index = index;
263    }
264
265    // Reorder the styles.
266    std::sort(std::begin(mStyles), std::end(mStyles),
267            [](const std::unique_ptr<StyleEntry>& lhs,
268               const std::unique_ptr<StyleEntry>& rhs) -> bool {
269                return lhs->str.getIndex() < rhs->str.getIndex();
270            }
271    );
272}
273
274template <typename T>
275static T* encodeLength(T* data, size_t length) {
276    static_assert(std::is_integral<T>::value, "wat.");
277
278    constexpr size_t kMask = 1 << ((sizeof(T) * 8) - 1);
279    constexpr size_t kMaxSize = kMask - 1;
280    if (length > kMaxSize) {
281        *data++ = kMask | (kMaxSize & (length >> (sizeof(T) * 8)));
282    }
283    *data++ = length;
284    return data;
285}
286
287template <typename T>
288static size_t encodedLengthUnits(size_t length) {
289    static_assert(std::is_integral<T>::value, "wat.");
290
291    constexpr size_t kMask = 1 << ((sizeof(T) * 8) - 1);
292    constexpr size_t kMaxSize = kMask - 1;
293    return length > kMaxSize ? 2 : 1;
294}
295
296
297bool StringPool::flatten(BigBuffer* out, const StringPool& pool, bool utf8) {
298    const size_t startIndex = out->size();
299    android::ResStringPool_header* header = out->nextBlock<android::ResStringPool_header>();
300    header->header.type = android::RES_STRING_POOL_TYPE;
301    header->header.headerSize = sizeof(*header);
302    header->stringCount = pool.size();
303    if (utf8) {
304        header->flags |= android::ResStringPool_header::UTF8_FLAG;
305    }
306
307    uint32_t* indices = pool.size() != 0 ? out->nextBlock<uint32_t>(pool.size()) : nullptr;
308
309    uint32_t* styleIndices = nullptr;
310    if (!pool.mStyles.empty()) {
311        header->styleCount = pool.mStyles.back()->str.getIndex() + 1;
312        styleIndices = out->nextBlock<uint32_t>(header->styleCount);
313    }
314
315    const size_t beforeStringsIndex = out->size();
316    header->stringsStart = beforeStringsIndex - startIndex;
317
318    for (const auto& entry : pool) {
319        *indices = out->size() - beforeStringsIndex;
320        indices++;
321
322        if (utf8) {
323            std::string encoded = util::utf16ToUtf8(entry->value);
324
325            const size_t totalSize = encodedLengthUnits<char>(entry->value.size())
326                    + encodedLengthUnits<char>(encoded.length())
327                    + encoded.size() + 1;
328
329            char* data = out->nextBlock<char>(totalSize);
330
331            // First encode the actual UTF16 string length.
332            data = encodeLength(data, entry->value.size());
333
334            // Now encode the size of the converted UTF8 string.
335            data = encodeLength(data, encoded.length());
336            strncpy(data, encoded.data(), encoded.size());
337        } else {
338            const size_t totalSize = encodedLengthUnits<char16_t>(entry->value.size())
339                    + entry->value.size() + 1;
340
341            char16_t* data = out->nextBlock<char16_t>(totalSize);
342
343            // Encode the actual UTF16 string length.
344            data = encodeLength(data, entry->value.size());
345            const size_t byteLength = entry->value.size() * sizeof(char16_t);
346
347            // NOTE: For some reason, strncpy16(data, entry->value.data(), entry->value.size())
348            // truncates the string.
349            memcpy(data, entry->value.data(), byteLength);
350
351            // The null-terminating character is already here due to the block of data being set
352            // to 0s on allocation.
353        }
354    }
355
356    out->align4();
357
358    if (!pool.mStyles.empty()) {
359        const size_t beforeStylesIndex = out->size();
360        header->stylesStart = beforeStylesIndex - startIndex;
361
362        size_t currentIndex = 0;
363        for (const auto& entry : pool.mStyles) {
364            while (entry->str.getIndex() > currentIndex) {
365                styleIndices[currentIndex++] = out->size() - beforeStylesIndex;
366
367                uint32_t* spanOffset = out->nextBlock<uint32_t>();
368                *spanOffset = android::ResStringPool_span::END;
369            }
370            styleIndices[currentIndex++] = out->size() - beforeStylesIndex;
371
372            android::ResStringPool_span* span =
373                    out->nextBlock<android::ResStringPool_span>(entry->spans.size());
374            for (const auto& s : entry->spans) {
375                span->name.index = s.name.getIndex();
376                span->firstChar = s.firstChar;
377                span->lastChar = s.lastChar;
378                span++;
379            }
380
381            uint32_t* spanEnd = out->nextBlock<uint32_t>();
382            *spanEnd = android::ResStringPool_span::END;
383        }
384
385        // The error checking code in the platform looks for an entire
386        // ResStringPool_span structure worth of 0xFFFFFFFF at the end
387        // of the style block, so fill in the remaining 2 32bit words
388        // with 0xFFFFFFFF.
389        const size_t paddingLength = sizeof(android::ResStringPool_span)
390                - sizeof(android::ResStringPool_span::name);
391        uint8_t* padding = out->nextBlock<uint8_t>(paddingLength);
392        memset(padding, 0xff, paddingLength);
393        out->align4();
394    }
395    header->header.size = out->size() - startIndex;
396    return true;
397}
398
399bool StringPool::flattenUtf8(BigBuffer* out, const StringPool& pool) {
400    return flatten(out, pool, true);
401}
402
403bool StringPool::flattenUtf16(BigBuffer* out, const StringPool& pool) {
404    return flatten(out, pool, false);
405}
406
407} // namespace aapt
408