From a8fd8c517211d4554ffd82a371f749a9db6d4e64 Mon Sep 17 00:00:00 2001
From: Fazana <52551480+FazanaJ@users.noreply.github.com>
Date: Thu, 30 Sep 2021 15:24:44 +0100
Subject: [PATCH] fix memory.c and add ucode_small flag

---
 include/object_constants.h   |  1 +
 src/boot/memory.c            |  9 ++++-----
 src/engine/behavior_script.c | 10 ++++++++++
 src/engine/graph_node.h      |  1 +
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/include/object_constants.h b/include/object_constants.h
index 7b9a3526..8bfe26dc 100644
--- a/include/object_constants.h
+++ b/include/object_constants.h
@@ -49,6 +49,7 @@
 #define OBJ_FLAG_SILHOUETTE                         (1 << 19) // 0x00080000
 #define OBJ_FLAG_OCCLUDE_SILHOUETTE                 (1 << 20) // 0x00100000
 #define OBJ_FLAG_OPACITY_FROM_CAMERA_DIST           (1 << 21) // 0x00200000
+#define OBJ_FLAG_UCODE_SMALL                        (1 << 22) // 0x00400000
 #define OBJ_FLAG_HITBOX_WAS_SET                     (1 << 30) // 0x40000000
 
 /* oHeldState */
diff --git a/src/boot/memory.c b/src/boot/memory.c
index e0914775..7d44ac60 100644
--- a/src/boot/memory.c
+++ b/src/boot/memory.c
@@ -68,7 +68,6 @@ struct MemoryPool *gEffectsMemoryPool;
 
 
 uintptr_t sSegmentTable[32];
-uintptr_t sSegmentROMTable[32];
 u32 sPoolFreeSpace;
 u8 *sPoolStart;
 u8 *sPoolEnd;
@@ -339,13 +338,13 @@ void *load_segment(s32 segment, u8 *srcStart, u8 *srcEnd, u32 side, u8 *bssStart
         addr = dynamic_dma_read(srcStart, srcEnd, side, TLB_PAGE_SIZE, (uintptr_t)bssEnd - (uintptr_t)bssStart);
         if (addr != NULL) {
             u8 *realAddr = (u8 *)ALIGN((uintptr_t)addr, TLB_PAGE_SIZE);
-            set_segment_base_addr(segment, realAddr); sSegmentROMTable[segment] = (uintptr_t) srcStart;
+            set_segment_base_addr(segment, realAddr);
             mapTLBPages(segment << 24, VIRTUAL_TO_PHYSICAL(realAddr), (srcEnd - srcStart) + ((uintptr_t)bssEnd - (uintptr_t)bssStart), segment);
         }
     } else {
         addr = dynamic_dma_read(srcStart, srcEnd, side, 0, 0);
         if (addr != NULL) {
-            set_segment_base_addr(segment, addr); sSegmentROMTable[segment] = (uintptr_t) srcStart;
+            set_segment_base_addr(segment, addr);
         }
     }
     #if PUPPYPRINT_DEBUG
@@ -420,7 +419,7 @@ void *load_segment_decompress(s32 segment, u8 *srcStart, u8 *srcEnd) {
             decompress(compressed, dest);
 #endif
             osSyncPrintf("end decompress\n");
-            set_segment_base_addr(segment, dest); sSegmentROMTable[segment] = (uintptr_t) srcStart;
+            set_segment_base_addr(segment, dest);
             main_pool_free(compressed);
         }
     }
@@ -459,7 +458,7 @@ void *load_segment_decompress_heap(u32 segment, u8 *srcStart, u8 *srcEnd) {
 #elif MIO0
         decompress(compressed, gDecompressionHeap);
 #endif
-        set_segment_base_addr(segment, gDecompressionHeap); sSegmentROMTable[segment] = (uintptr_t) srcStart;
+        set_segment_base_addr(segment, gDecompressionHeap);
         main_pool_free(compressed);
     }
     return gDecompressionHeap;
diff --git a/src/engine/behavior_script.c b/src/engine/behavior_script.c
index 44eb3838..2ea40dad 100644
--- a/src/engine/behavior_script.c
+++ b/src/engine/behavior_script.c
@@ -920,6 +920,7 @@ void cur_obj_update(void) {
     f32 distanceFromMario;
     BhvCommandProc bhvCmdProc;
     s32 bhvProcResult;
+    s32 objListIndex;
 
     // Calculate the distance from the object to Mario.
     if (objFlags & OBJ_FLAG_COMPUTE_DIST_TO_MARIO) {
@@ -995,6 +996,15 @@ void cur_obj_update(void) {
     COND_BIT((!(objFlags & OBJ_FLAG_UCODE_LARGE       )), gCurrentObject->header.gfx.node.flags, GRAPH_RENDER_UCODE_REJ         );
     COND_BIT((  objFlags & OBJ_FLAG_SILHOUETTE         ), gCurrentObject->header.gfx.node.flags, GRAPH_RENDER_SILHOUETTE        );
     COND_BIT((  objFlags & OBJ_FLAG_OCCLUDE_SILHOUETTE ), gCurrentObject->header.gfx.node.flags, GRAPH_RENDER_OCCLUDE_SILHOUETTE);
+    BehaviorScript *bhvScript = segmented_to_virtual(gCurrentObject->behavior);
+    if ((bhvScript[0] >> 24) == 0) {
+        objListIndex = (bhvScript[0] >> 16) & 0xFFFF;
+    }
+    if (objListIndex == OBJ_LIST_SURFACE && !(objFlags & OBJ_FLAG_UCODE_SMALL))
+    {
+        gCurrentObject->header.gfx.node.flags &= ~GRAPH_RENDER_UCODE_REJ;
+        gCurrentObject->header.gfx.node.flags |= GRAPH_RENDER_UCODE_ZEX;
+    }
 
 #ifdef OBJ_OPACITY_BY_CAM_DIST
     if (objFlags & OBJ_FLAG_OPACITY_FROM_CAMERA_DIST) {
diff --git a/src/engine/graph_node.h b/src/engine/graph_node.h
index cc033f5c..f112dc9f 100644
--- a/src/engine/graph_node.h
+++ b/src/engine/graph_node.h
@@ -18,6 +18,7 @@
 #define GRAPH_RENDER_SILHOUETTE         (1 << 6) // 0x0040
 #define GRAPH_RENDER_OCCLUDE_SILHOUETTE (1 << 7) // 0x0080
 #define GRAPH_RENDER_UCODE_REJ          (1 << 8) // 0x0100
+#define GRAPH_RENDER_UCODE_ZEX          (1 << 9) // 0x0200
 
 // The amount of bits to use for the above flags out of a s16 variable.
 // The remaining bits to the left are used for the render layers.