Merge bundle tables without copying the base table

src/main/java/com/reandroid/apk/ApkBundle.java

@@ -51,7 +51,7 @@ public ApkModule mergeModules(boolean force) throws IOException {
         if(base == null){
             base = getLargestTableModule();
         }
-        result.merge(base, force);
+        result.merge(base, force, false);
         AndroidManifestBlockMerger manifestMerger = getManifestMerger();
         if (manifestMerger != null) {
             manifestMerger.reset();
@@ -69,7 +69,7 @@ public ApkModule mergeModules(boolean force) throws IOException {
             if(signatureBlock == null){
                 signatureBlock = asb;
             }
-            result.merge(module, force);
+            result.merge(module, force, false);
             if (manifestMerger != null) {
                 manifestMerger.merge(module.getAndroidManifest());
             }
@@ -107,17 +107,19 @@ private String generateMergedModuleName(){
         }
         return name;
     }
-    private ApkModule getLargestTableModule(){
+    private ApkModule getLargestTableModule() throws IOException {
+        // Compared by entry size, so the tables are not all parsed just to pick one
         ApkModule apkModule=null;
-        int chunkSize=0;
+        long largest=0;
         for(ApkModule module:getApkModuleList()){
             if(!module.hasTableBlock()){
                 continue;
             }
-            TableBlock tableBlock=module.getTableBlock();
-            int size=tableBlock.getHeaderBlock().getChunkSize();
-            if(apkModule==null || size>chunkSize){
-                chunkSize=size;
+            TableBlock loaded=module.getLoadedTableBlock();
+            long size=loaded != null ? loaded.getHeaderBlock().getChunkSize()
+                    : module.getInputSource(TableBlock.FILE_NAME).getLength();
+            if(apkModule==null || size>largest){
+                largest=size;
                 apkModule=module;
             }
         }

src/main/java/com/reandroid/apk/ApkModule.java

@@ -1154,9 +1154,7 @@ TableBlock loadTableBlock() throws IOException {
             tableBlock = (TableBlock) ((BlockInputSource<?>) inputSource).getBlock();
         } else {
             setTableOriginalSource(inputSource);
-            InputStream inputStream = inputSource.openStream();
-            tableBlock = TableBlock.load(inputStream);
-            inputStream.close();
+            tableBlock = parseTableBlock(inputSource);
         }
         BlockInputSource<TableBlock> blockInputSource = new BlockInputSource<>(
                 inputSource.getName(), tableBlock);
@@ -1220,13 +1218,20 @@ public void merge(ApkModule module) throws IOException {
         merge(module, false);
     }
     public void merge(ApkModule module, boolean force) throws IOException {
+        merge(module, force, true);
+    }
+    /**
+     * @param refreshTable whether to refresh the merged table afterwards. {@link ApkBundle}
+     *                     merges every module first and refreshes the table once at the end.
+     */
+    void merge(ApkModule module, boolean force, boolean refreshTable) throws IOException {
         if (module == null || module == this) {
             return;
         }
         logMessage("Merging: " + module.getModuleName());
         validateMerge(module, force);
         mergeDexFiles(module);
-        mergeTable(module);
+        mergeTable(module, refreshTable);
         mergeFiles(module);
         getUncompressedFiles().merge(module.getUncompressedFiles());
         mergeFusedModules(module);
@@ -1291,20 +1296,50 @@ private void validateMerge(ApkModule apkModule, boolean force) throws IOExceptio
         }
         logMessage(msg);
     }
-    private void mergeTable(ApkModule module) {
+    private void mergeTable(ApkModule module, boolean refresh) throws IOException {
         if (!module.hasTableBlock()) {
             return;
         }
+        TableBlock coming = module.getLoadedTableBlock();
+        if (coming == null) {
+            // Parsed for this merge only, so nothing else holds it: it can become the merged
+            // table as is, or be dropped once merged
+            coming = module.loadDetachedTableBlock();
+            if (coming != null && !hasTableBlock()) {
+                addInputSource(new BlockInputSource<>(TableBlock.FILE_NAME, coming));
+                return;
+            }
+            if (coming == null) {
+                coming = module.getTableBlock();
+            }
+        }
         TableBlock exist;
         if (!hasTableBlock()) {
-            exist=new TableBlock();
-            BlockInputSource<TableBlock> inputSource=new BlockInputSource<>(TableBlock.FILE_NAME, exist);
-            addInputSource(inputSource);
+            exist = new TableBlock();
+            addInputSource(new BlockInputSource<>(TableBlock.FILE_NAME, exist));
         } else {
-            exist=getTableBlock();
+            exist = getTableBlock();
+        }
+        exist.merge(coming, refresh);
+    }
+    /**
+     * Parses resources.arsc without caching it on this module, or returns null when the
+     * table source is already a block rather than an archive entry.
+     */
+    private TableBlock loadDetachedTableBlock() throws IOException {
+        InputSource inputSource = getInputSource(TableBlock.FILE_NAME);
+        if (inputSource == null || inputSource instanceof BlockInputSource) {
+            return null;
+        }
+        return parseTableBlock(inputSource);
+    }
+    private static TableBlock parseTableBlock(InputSource inputSource) throws IOException {
+        InputStream inputStream = inputSource.openStream();
+        try {
+            return TableBlock.load(inputStream);
+        } finally {
+            inputStream.close();
         }
-        TableBlock coming=module.getTableBlock();
-        exist.merge(coming);
     }
     private void mergeFiles(ApkModule module) {
         ZipEntryMap entryMapExist = getZipEntryMap();
@@ -1328,7 +1363,8 @@ private void mergeFiles(ApkModule module) {
         }
     }
     private void mergeDexFiles(ApkModule module) {
-        UncompressedFiles uncompressedFiles = module.getUncompressedFiles();
+        UncompressedFiles comingUncompressed = module.getUncompressedFiles();
+        UncompressedFiles uncompressed = getUncompressedFiles();
         List<DexFileInputSource> existList = listDexFiles();
         List<DexFileInputSource> comingList = module.listDexFiles();
         ZipEntryMap zipEntryMap = getZipEntryMap();
@@ -1342,8 +1378,13 @@ private void mergeDexFiles(ApkModule module) {
             }
         }
         for (DexFileInputSource source : comingList) {
-            uncompressedFiles.removePath(source.getAlias());
             String name = DexFileInputSource.getDexName(index);
+            if (comingUncompressed.isUncompressed(source.getAlias())) {
+                // Keep a stored dex stored under its new name, so it is copied as is
+                // instead of being deflated while writing
+                uncompressed.addPath(name);
+            }
+            comingUncompressed.removePath(source.getAlias());
             DexFileInputSource add = new DexFileInputSource(name, source.getInputSource());
             zipEntryMap.add(add);
             logMessage("Added [" + module.getModuleName() +"] "

src/main/java/com/reandroid/arsc/chunk/TableBlock.java

@@ -795,12 +795,21 @@ public void fromJson(JSONObject json) {
         refresh();
     }
     public void merge(TableBlock tableBlock){
+        merge(tableBlock, true);
+    }
+    /**
+     * @param refresh whether to refresh this table afterwards. A caller merging several tables
+     *                can skip it for all but the last, as refreshing a large table is costly.
+     */
+    public void merge(TableBlock tableBlock, boolean refresh){
         if(tableBlock == null || tableBlock == this){
             return;
         }
         getStringPool().merge(tableBlock.getStringPool());
         getPackageArray().merge(tableBlock.getPackageArray());
-        refresh();
+        if(refresh){
+            refresh();
+        }
     }
     @Override
     public byte[] getBytes(){

src/test/java/com/reandroid/apk/ApkBundleMergeTest.java

@@ -0,0 +1,94 @@
+package com.reandroid.apk;
+
+import com.reandroid.TestUtils;
+import com.reandroid.archive.ByteInputSource;
+import com.reandroid.archive.ZipEntryMap;
+import com.reandroid.arsc.chunk.PackageBlock;
+import com.reandroid.arsc.chunk.TableBlock;
+import com.reandroid.arsc.chunk.xml.AndroidManifestBlock;
+import com.reandroid.arsc.value.Entry;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+
+public class ApkBundleMergeTest {
+
+    private static final String PACKAGE = "com.reandroid.merge";
+
+    @Test
+    public void testMergeKeepsBaseResourcesAndAddsSplitConfigs() throws IOException {
+        ApkModule merged = merge();
+
+        PackageBlock packageBlock = merged.getTableBlock().pickOne();
+        Assert.assertEquals("Hello", value(packageBlock, "", "hello"));
+        Assert.assertEquals("Hallo", value(packageBlock, "-de", "hello"));
+        Assert.assertEquals("Only in base", value(packageBlock, "", "base_only"));
+    }
+
+    @Test
+    public void testMergeKeepsStoredDexStored() throws IOException {
+        ApkModule merged = merge();
+
+        Assert.assertTrue(merged.getUncompressedFiles().isUncompressed("classes.dex"));
+        Assert.assertTrue(merged.getUncompressedFiles().isUncompressed("classes2.dex"));
+    }
+
+    /** Merges, writes and reloads the result, as callers use it */
+    private static ApkModule merge() throws IOException {
+        ApkBundle bundle = new ApkBundle();
+        bundle.addModule(createBase());
+        bundle.addModule(createSplit());
+        ApkModule merged = bundle.mergeModules();
+        File file = new File(TestUtils.getTempDir(), "bundle_merged.apk");
+        merged.writeApk(file);
+        merged.close();
+        bundle.close();
+        return ApkModule.loadApkFile(file);
+    }
+
+    private static ApkModule createBase() throws IOException {
+        TableBlock tableBlock = new TableBlock();
+        PackageBlock packageBlock = tableBlock.newPackage(0x7f, PACKAGE);
+        packageBlock.getOrCreate("", "string", "hello").setValueAsString("Hello");
+        packageBlock.getOrCreate("", "string", "base_only").setValueAsString("Only in base");
+        tableBlock.refreshFull();
+        return createModule("base", tableBlock, null);
+    }
+
+    private static ApkModule createSplit() throws IOException {
+        TableBlock tableBlock = new TableBlock();
+        PackageBlock packageBlock = tableBlock.newPackage(0x7f, PACKAGE);
+        // Created first, so it gets the id "hello" has in the base
+        packageBlock.getOrCreate("-de", "string", "hello").setValueAsString("Hallo");
+        tableBlock.refreshFull();
+        return createModule("split_config.de", tableBlock, "config.de");
+    }
+
+    private static ApkModule createModule(String name, TableBlock tableBlock, String split) {
+        ApkModule module = new ApkModule(name, new ZipEntryMap());
+        module.setLoadDefaultFramework(false);
+
+        AndroidManifestBlock manifest = new AndroidManifestBlock();
+        manifest.setPackageName(PACKAGE);
+        if (split != null) {
+            manifest.setSplit(split, true);
+        }
+        manifest.refresh();
+        module.add(new ByteInputSource(manifest.getBytes(), AndroidManifestBlock.FILE_NAME));
+
+        // A table the module has not parsed yet, as when loaded from an archive
+        module.add(new ByteInputSource(tableBlock.getBytes(), TableBlock.FILE_NAME));
+
+        module.add(new ByteInputSource(new byte[]{'d', 'e', 'x', '\n'}, "classes.dex"));
+        module.getUncompressedFiles().addPath("classes.dex");
+        return module;
+    }
+
+    private static String value(PackageBlock packageBlock, String qualifiers, String name) {
+        Entry entry = packageBlock.getEntry(qualifiers, "string", name);
+        Assert.assertNotNull(qualifiers + "/" + name, entry);
+        return entry.getResValue().getValueAsString();
+    }
+}