From ec52fa24a84d8ecb176445801b751b9583f75fbc Mon Sep 17 00:00:00 2001
From: CrazyRedMachine <crazyredmachine86@gmail.com>
Date: Mon, 20 May 2024 21:54:49 +0200
Subject: [PATCH] add BST for customs

---
 popnhax/custom_categs.cc |  7 ++++++-
 popnhax/custom_categs.h  |  2 ++
 popnhax/loader.cc        | 30 ++++++++++++++++++++++--------
 util/Module.mk           |  1 +
 util/bst.cc              | 34 ++++++++++++++++++++++++++++++++++
 util/bst.h               | 14 ++++++++++++++
 6 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 util/bst.cc
 create mode 100644 util/bst.h

diff --git a/popnhax/custom_categs.cc b/popnhax/custom_categs.cc
index 6025389..6d3eca6 100644
--- a/popnhax/custom_categs.cc
+++ b/popnhax/custom_categs.cc
@@ -19,6 +19,8 @@
 #include "minhook/hde32.h"
 #include "minhook/include/MinHook.h"
 
+#include "custom_categs.h"
+
 #define F_OK 0
 
 uint8_t g_game_version;
@@ -26,6 +28,8 @@ uint32_t g_playerdata_ptr_addr; //pointer to the playerdata memory zone (offset
 char *g_current_friendid;
 uint32_t g_current_songid;
 
+bst_t *g_customs_bst = NULL;
+
 void (*add_song_in_list)();
 //game code takes array start address from offset 0xC and the address after the list end from offset 0x10
 typedef struct songlist_s {
@@ -57,7 +61,8 @@ uint32_t favorites_struct_addr = (uint32_t)&favorites_struct;
 
 bool is_a_custom(uint32_t songid)
 {
-	return (songid >= g_min_id && (g_max_id==0 || g_max_id >= songid));
+	return (bst_search(g_customs_bst, songid) != NULL);
+//	return (songid >= g_min_id && (g_max_id==0 || g_max_id >= songid));
 }
 
 void add_song_to_favorites()
diff --git a/popnhax/custom_categs.h b/popnhax/custom_categs.h
index cd10951..40f9f25 100644
--- a/popnhax/custom_categs.h
+++ b/popnhax/custom_categs.h
@@ -3,8 +3,10 @@
 
 #include <stdint.h>
 #include "popnhax/config.h"
+#include "util/bst.h"
 
 extern uint32_t g_max_id;
+extern bst_t *g_customs_bst;
 
 bool is_excluded_folder(const char *input_filename);
 bool patch_custom_categs(const char *dllFilename, struct popnhax_config *config);
diff --git a/popnhax/loader.cc b/popnhax/loader.cc
index 0d3ba3f..ba4246d 100644
--- a/popnhax/loader.cc
+++ b/popnhax/loader.cc
@@ -4,6 +4,7 @@
 #include <windows.h>
 
 #include "imports/avs.h"
+#include "util/bst.h"
 #include "util/patch.h"
 #include "util/log.h"
 #include "xmlhelper.h"
@@ -856,16 +857,32 @@ void parse_musicdb(const char *input_filename, const char *target, struct popnha
                                 sizeof(idxStr));
             uint32_t idx = atoi(idxStr);
 
-if (idx > g_max_id)
-{
-    g_max_id = idx;
-}
             // Get an existing music entry in memory
             // If it exists, return the existing entry
             // If it doesn't exist, create a new entry in memory
             // Update the data in-place and make all parameters optional
             music_entry *m = get_music(idx);
-            bool is_fresh = m == NULL;
+            bool is_fresh = m == NULL; // ie. not part of internal songdb
+            bool is_gone = ( m != NULL && strcmp((const char*) m->title_ptr, "\x81\x5D") == 0); // removed entries all have this title (SJIS "-")
+
+            // Update customs/omni songid list
+            if ( is_fresh || is_gone )
+            {
+                //TODO: remove g_max_id entirely
+                if (idx > g_max_id)
+                {
+                    g_max_id = idx;
+                }
+
+                if ( bst_search(g_customs_bst, idx) == NULL )
+                {
+                    g_customs_bst = bst_insert(g_customs_bst, idx);
+                    //LOG("%d inserted into customs bst\n", idx);
+                    //TODO: handle custom category creation here as well while we're at it (beware: maybe we should still consider !is_gone charts as custom?)
+                } else {
+                    //LOG("%d already present in customs bst\n", idx);
+                }
+            }
 
             if (is_fresh) {
                 // Default music entry
@@ -921,9 +938,6 @@ if (idx > g_max_id)
                 }
             }
 
-            //force loading background for unilab
-            //m->mask |= 0x100;
-
             if ( config->custom_categ
               && config->custom_exclude_from_version
               && !is_excluded_folder(input_filename)
diff --git a/util/Module.mk b/util/Module.mk
index 6452371..c86a17e 100644
--- a/util/Module.mk
+++ b/util/Module.mk
@@ -2,6 +2,7 @@ libs                    += util
 
 srcpp_util              := \
 	fuzzy_search.cc \
+	bst.cc \
 	search.cc \
 	cmdline.cc \
 	patch.cc \
diff --git a/util/bst.cc b/util/bst.cc
new file mode 100644
index 0000000..2c6e22c
--- /dev/null
+++ b/util/bst.cc
@@ -0,0 +1,34 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+
+#include "bst.h"
+
+bst_t* bst_search(bst_t *root, uint32_t val)
+{
+    if( root == NULL || root->data == val )
+        return root;
+    else if( val > (root->data) )
+        return bst_search(root->right, val);
+    else
+        return bst_search(root->left,val);
+}
+
+bst_t* bst_insert(bst_t *root, uint32_t val)
+{
+    if ( root == NULL )
+    {
+        bst_t *p;
+        p = (bst_t *)malloc(sizeof(bst_t));
+        p->data = val;
+        p->left = NULL;
+        p->right = NULL;
+        return p;
+    }
+    else if ( val > root->data )
+        root->right = bst_insert(root->right, val);
+    else
+        root->left = bst_insert(root->left, val);
+
+    return root;
+}
\ No newline at end of file
diff --git a/util/bst.h b/util/bst.h
new file mode 100644
index 0000000..c8338b2
--- /dev/null
+++ b/util/bst.h
@@ -0,0 +1,14 @@
+#ifndef __BST_H__
+#define __BST_H__
+
+typedef struct bst_s
+{
+    uint32_t data;
+    struct bst_s *right;
+    struct bst_s *left;
+} bst_t;
+
+bst_t* bst_search(bst_t *root, uint32_t val);
+bst_t* bst_insert(bst_t *root, uint32_t val);
+
+#endif