From b5204037dabd8c8cf0d528fffa7e92d49cd60536 Mon Sep 17 00:00:00 2001
From: Aetias <aetias@outlook.com>
Date: Sat, 23 Sep 2023 19:05:42 +0200
Subject: [PATCH] Match compressed code

---
 Makefile              | 18 +++++++++++++++-
 tools/compress/main.c | 48 ++++++++++++++++++++++++++++++++-----------
 tools/lcf.py          |  2 +-
 3 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/Makefile b/Makefile
index ff7dbb67..da2dc611 100644
--- a/Makefile
+++ b/Makefile
@@ -9,6 +9,9 @@ CXX_FILES := $(wildcard src/*.cpp)
 ASM_OBJS = $(ASM_FILES:%.s=$(BUILD_DIR)/%.o)
 CXX_OBJS = $(CXX_FILES:%.cpp=$(BUILD_DIR)/%.o)
 
+OV_BINS := $(wildcard $(BUILD_DIR)/overlays/*.bin)
+OV_LZS = $(OV_BINS:%.bin=%.lz)
+
 MW_VER     := 2.0/sp1p5
 MW_ASM     := $(TOOLS_DIR)/mwccarm/$(MW_VER)/mwasmarm
 MW_CC      := $(TOOLS_DIR)/mwccarm/$(MW_VER)/mwccarm
@@ -20,14 +23,20 @@ CC_FLAGS  := -O1 -thumb
 LD_FLAGS  := -proc arm946e -nostdlib -nointerworking -nodead -m func_02000800 -map closure,unused -o main.bin -msgstyle gcc
 
 .PHONY: all
-all: arm9
+all: tools arm9
+
+.PHONY: tools
+tools:
+	cd $(TOOLS_DIR)/compress && $(MAKE)
 
 .PHONY: arm9
 arm9: setup $(ASM_OBJS) lcf link
+	$(MAKE) compress
 
 .PHONY: setup
 setup:
 	mkdir -p $(BUILD_DIR)
+	mkdir -p $(BUILD_DIR)/overlays
 
 .PHONY: clean
 clean:
@@ -48,3 +57,10 @@ $(CXX_OBJS): $(BUILD_DIR)/%.o: %.cpp
 .PHONY: link
 link:
 	cd $(BUILD_DIR) && LM_LICENSE_FILE=$(MW_LICENSE) $(MW_LD) $(LD_FLAGS) $(LCF_FILE) @$(OBJS_FILE)
+
+.PHONY: compress
+compress: $(OV_LZS)
+	$(TOOLS_DIR)/compress/compress -s 0x4000 -i $(BUILD_DIR)/arm9.bin -o $(BUILD_DIR)/arm9.lz
+
+$(OV_LZS): %.lz: %.bin
+	$(TOOLS_DIR)/compress/compress -p -i $< -o $@
diff --git a/tools/compress/main.c b/tools/compress/main.c
index 638b8405..bfe20d6f 100644
--- a/tools/compress/main.c
+++ b/tools/compress/main.c
@@ -30,6 +30,11 @@
 #define MAX_SUBSEQ ((MIN_SUBSEQ) + LEN_MASK)
 #define LOOKAHEAD (1 << (DIST_BITS))
 
+typedef struct {
+    uint8_t *pos;
+    uint32_t totalBytesSaved;
+} BlockInfo;
+
 bool FindSubsequence(const uint8_t *buf, const uint8_t *start, const uint8_t *end, size_t *pLen, size_t *pDist) {
     size_t bestLen = 0;
     size_t bestDist = 0;
@@ -62,17 +67,10 @@ bool Compress(const uint8_t *src, uint8_t *dst, size_t size, uint8_t **pResult,
     uint8_t *pFlags = write;
     uint8_t flags;
     size_t flagCount = 0;
-    size_t numIdentical = 0;
+    int32_t bytesSaved = 0;
+    BlockInfo *blockInfos = malloc(((size + 7) / 8) * sizeof(BlockInfo));
+    size_t numBlocks = 0;
     while (read > src) {
-        if (read - 1 == src && flagCount == 0) {
-            read -= 1;
-            *write = *read;
-            for (numIdentical = 0; write + numIdentical < writeEnd; ++numIdentical) {
-                if (read[numIdentical] != write[numIdentical]) break;
-            }
-            write -= 1;
-            break;
-        }
         flags <<= 1;
         size_t len, dist;
         if (FindSubsequence(read - 1, src, end, &len, &dist)) {
@@ -82,6 +80,7 @@ bool Compress(const uint8_t *src, uint8_t *dst, size_t size, uint8_t **pResult,
             write -= 2;
             WRITE16(write, pair);
             flags |= 1;
+            bytesSaved += len - 2;
         } else {
             // write literal
             write -= 1;
@@ -92,13 +91,38 @@ bool Compress(const uint8_t *src, uint8_t *dst, size_t size, uint8_t **pResult,
         if (flagCount == 8) {
             flagCount = 0;
             *pFlags = flags;
+            bytesSaved -= 1;
             flags = 0;
             write -= 1;
             pFlags = write;
+
+            blockInfos[numBlocks].pos = write;
+            blockInfos[numBlocks].totalBytesSaved = bytesSaved;
+            numBlocks++;
+        }
+    }
+    if (flagCount != 0) {
+        blockInfos[numBlocks].pos = write;
+        blockInfos[numBlocks].totalBytesSaved = bytesSaved;
+        numBlocks++;
+
+        *pFlags = flags << (8 - flagCount);
+    }
+    else {
+        write += 1;
+    }
+    size_t numIdentical = 0;
+    for (int32_t i = 0; i < numBlocks - 1; ++i) {
+        if (blockInfos[i].totalBytesSaved == bytesSaved) {
+            numIdentical = blockInfos[i].pos - write;
+            write += 1;
+            memcpy(write, src, numIdentical);
+            while (write[numIdentical] == read[numIdentical]) {
+                numIdentical += 1;
+            }
+            break;
         }
     }
-    if (flagCount != 0) *pFlags = flags << (8 - flagCount);
-    else write += 1;
     *pResult = write;
     *pLen = size - (write - dst);
     *pNumIdentical = numIdentical;
diff --git a/tools/lcf.py b/tools/lcf.py
index 7d3d302d..2d8b6892 100644
--- a/tools/lcf.py
+++ b/tools/lcf.py
@@ -233,7 +233,7 @@ with open(f'{BUILD}arm9_linker_script.lcf', 'w') as file:
             file.write(f'AFTER({ov.after.name})')
         elif type(ov.after) is str:
             file.write(f'AFTER({ov.after})')
-        file.write(f' > {ov.name}.bin\n')
+        file.write(f' > overlays/{ov.name}.bin\n')
     file.write('}\n')
     file.write('\n')
     file.write('SECTIONS {\n')