Make pinned displays inherit existing displays' origins

4 files changed, 120 insertions, 1 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/favorites/history/DisplayHistoryManager.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/favorites/history/DisplayHistoryManager.java
index 859646bb1..c6b13f78a 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/favorites/history/DisplayHistoryManager.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/gui/widget/favorites/history/DisplayHistoryManager.java
@@ -24,6 +24,7 @@
 package me.shedaniel.rei.impl.client.gui.widget.favorites.history;
 
 import com.google.common.collect.Iterables;
+import it.unimi.dsi.fastutil.objects.Reference2ObjectLinkedOpenHashMap;
 import me.shedaniel.math.Rectangle;
 import me.shedaniel.rei.api.client.registry.category.CategoryRegistry;
 import me.shedaniel.rei.api.common.category.CategoryIdentifier;
@@ -31,9 +32,12 @@ import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.display.DisplaySerializerRegistry;
 import me.shedaniel.rei.api.common.plugins.PluginManager;
 import me.shedaniel.rei.impl.client.config.ConfigManagerImpl;
+import me.shedaniel.rei.impl.client.registry.display.DisplayKey;
+import me.shedaniel.rei.impl.client.registry.display.DisplaysHolder;
 import me.shedaniel.rei.impl.common.InternalLogger;
 import net.minecraft.Util;
 import net.minecraft.nbt.CompoundTag;
+import net.minecraft.resources.ResourceLocation;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
@@ -41,8 +45,33 @@ import java.util.*;
 public class DisplayHistoryManager {
     public static final DisplayHistoryManager INSTANCE = new DisplayHistoryManager();
     private Map<String, DisplayEntry> entries = new LinkedHashMap<>();
+    private final Map<Display, DisplayEntry> displayToEntries = new Reference2ObjectLinkedOpenHashMap<>();
     private long lastCheckTime = -1;
     
+    @Nullable
+    public Object getPossibleOrigin(DisplaysHolder holder, Display display) {
+        DisplayEntry entry = displayToEntries.get(display);
+        if (entry == null) return null;
+        Optional<ResourceLocation> location = display.getDisplayLocation();
+        if (location.isEmpty()) return null;
+        Set<Display> displays = holder.getDisplaysByKey(DisplayKey.create(display.getCategoryIdentifier(), location.get()));
+        for (Display d : displays) {
+            if (!displayToEntries.containsKey(d)) {
+                Object origin = holder.getDisplayOrigin(d);
+                
+                if (origin != null) {
+                    return origin;
+                }
+            }
+        }
+        
+        return null;
+    }
+    
+    public Map<Display, DisplayEntry> getDisplayToEntries() {
+        return displayToEntries;
+    }
+    
     public Collection<DisplayEntry> getEntries(DisplayHistoryWidget parent) {
         if ((lastCheckTime == -1 || Util.getMillis() - lastCheckTime > 4000) && !PluginManager.areAnyReloading()) {
             updateEntries(parent);
@@ -56,12 +85,14 @@ public class DisplayHistoryManager {
         List<CompoundTag> displayHistory = ConfigManagerImpl.getInstance().getConfig().getDisplayHistory();
         Map<String, DisplayEntry> copy = new LinkedHashMap<>(entries);
         entries.clear();
+        displayToEntries.clear();
         for (CompoundTag tag : displayHistory) {
             String uuid = tag.getString("DisplayHistoryUUID");
             
             DisplayEntry entry = copy.get(uuid);
             if (entry != null) {
                 entries.put(entry.getUuid().toString(), entry);
+                displayToEntries.put(entry.getDisplay(), entry);
             } else if (tag.getBoolean("DisplayHistoryContains")) {
                 try {
                     CategoryIdentifier<?> categoryIdentifier = CategoryIdentifier.of(tag.getString("DisplayHistoryCategory"));
@@ -70,6 +101,7 @@ public class DisplayHistoryManager {
                         DisplayEntry newEntry = new DisplayEntry(parent, display, null);
                         newEntry.setUuid(UUID.fromString(uuid));
                         entries.put(newEntry.getUuid().toString(), newEntry);
+                        displayToEntries.put(newEntry.getDisplay(), newEntry);
                     }
                 } catch (Exception e) {
                     InternalLogger.getInstance().warn("Failed to read display history entry", e);
@@ -80,6 +112,7 @@ public class DisplayHistoryManager {
     
     public void removeEntry(DisplayEntry entry) {
         this.entries.remove(entry.getUuid().toString());
+        this.displayToEntries.remove(entry.getDisplay());
         List<CompoundTag> displayHistory = ConfigManagerImpl.getInstance().getConfig().getDisplayHistory();
         displayHistory.removeIf(tag -> tag.getString("DisplayHistoryUUID").equals(entry.getUuid().toString()));
         save();
@@ -92,6 +125,7 @@ public class DisplayHistoryManager {
             DisplayEntry entry = iterator.next();
             if (entry.getDisplay() == display) {
                 displayHistory.removeIf(tag -> tag.getString("DisplayHistoryUUID").equals(entry.getUuid().toString()));
+                this.displayToEntries.remove(entry.getDisplay());
                 iterator.remove();
             }
         }
@@ -100,10 +134,15 @@ public class DisplayHistoryManager {
         copy.put(newEntry.getUuid().toString(), newEntry);
         copy.putAll(this.entries);
         this.entries = copy;
+        this.displayToEntries.clear();
+        for (DisplayEntry entry : this.entries.values()) {
+            this.displayToEntries.put(entry.getDisplay(), entry);
+        }
         while (entries.size() >= 10) {
             DisplayEntry entry = Iterables.get(entries.values(), entries.size() - 1);
             displayHistory.removeIf(tag -> tag.getString("DisplayHistoryUUID").equals(entry.getUuid().toString()));
             this.entries.remove(entry.getUuid().toString());
+            this.displayToEntries.remove(entry.getDisplay());
         }
         
         CompoundTag compoundTag = new CompoundTag();
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayKey.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayKey.java
new file mode 100644
index 000000000..e56b70a3e
--- /dev/null
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplayKey.java
@@ -0,0 +1,60 @@
+/*
+ * This file is licensed under the MIT License, part of Roughly Enough Items.
+ * Copyright (c) 2018, 2019, 2020, 2021, 2022, 2023 shedaniel
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package me.shedaniel.rei.impl.client.registry.display;
+
+import com.google.common.collect.Maps;
+import me.shedaniel.rei.api.common.category.CategoryIdentifier;
+import net.minecraft.resources.ResourceLocation;
+
+import java.util.Collections;
+import java.util.Map;
+
+public class DisplayKey {
+    private static final Map<String, DisplayKey> VALUES = Collections.synchronizedMap(Maps.newIdentityHashMap());
+    private final CategoryIdentifier<?> categoryIdentifier;
+    private final ResourceLocation location;
+    
+    public static DisplayKey create(CategoryIdentifier<?> categoryIdentifier, ResourceLocation location) {
+        String string = (categoryIdentifier + ":" + location).intern();
+        return VALUES.computeIfAbsent(string, $ -> new DisplayKey(categoryIdentifier, location));
+    }
+    
+    private DisplayKey(CategoryIdentifier<?> categoryIdentifier, ResourceLocation location) {
+        this.categoryIdentifier = categoryIdentifier;
+        this.location = location;
+    }
+    
+    @Override
+    public String toString() {
+        return "DisplayKey[" + this.categoryIdentifier + " / " + this.location + "]";
+    }
+    
+    public CategoryIdentifier<?> categoryIdentifier() {
+        return this.categoryIdentifier;
+    }
+    
+    public ResourceLocation location() {
+        return this.location;
+    }
+}
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolder.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolder.java
index 07b5879cf..2480053d1 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolder.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolder.java
@@ -50,6 +50,8 @@ public interface DisplaysHolder {
     
     void endReload();
     
+    Set<Display> getDisplaysByKey(DisplayKey key);
+    
     boolean isCached(Display display);
     
     Set<Display> getDisplaysNotCached();
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
index f1f930639..b2bdce14d 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/registry/display/DisplaysHolderImpl.java
@@ -36,7 +36,9 @@ import me.shedaniel.rei.api.common.display.Display;
 import me.shedaniel.rei.api.common.entry.EntryIngredient;
 import me.shedaniel.rei.api.common.entry.EntryStack;
 import me.shedaniel.rei.api.common.util.EntryStacks;
+import me.shedaniel.rei.impl.client.gui.widget.favorites.history.DisplayHistoryManager;
 import me.shedaniel.rei.impl.common.InternalLogger;
+import net.minecraft.resources.ResourceLocation;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.jetbrains.annotations.Nullable;
 
@@ -45,6 +47,7 @@ import java.util.concurrent.ConcurrentHashMap;
 
 public class DisplaysHolderImpl implements DisplaysHolder {
     private final boolean cache;
+    private final SetMultimap<DisplayKey, Display> displaysByKey = Multimaps.newSetMultimap(new IdentityHashMap<>(), ReferenceOpenHashSet::new);
     private final Map<CategoryIdentifier<?>, DisplaysList> displays = new ConcurrentHashMap<>();
     private final Map<CategoryIdentifier<?>, List<Display>> unmodifiableDisplays;
     private final WeakHashMap<Display, Object> displaysBase = new WeakHashMap<>();
@@ -77,6 +80,10 @@ public class DisplaysHolderImpl implements DisplaysHolder {
     public void add(Display display, @Nullable Object origin) {
         this.displays.computeIfAbsent(display.getCategoryIdentifier(), location -> new DisplaysList())
                 .add(display);
+        Optional<ResourceLocation> location = display.getDisplayLocation();
+        if (location.isPresent()) {
+            this.displaysByKey.put(DisplayKey.create(display.getCategoryIdentifier(), location.get()), display);
+        }
         this.displayCount.increment();
         if (origin != null) {
             synchronized (this.displaysBase) {
@@ -137,6 +144,11 @@ public class DisplaysHolderImpl implements DisplaysHolder {
     }
     
     @Override
+    public Set<Display> getDisplaysByKey(DisplayKey key) {
+        return this.displaysByKey.get(key);
+    }
+    
+    @Override
     public boolean isCached(Display display) {
         return this.cache && this.displaysCached.contains(display);
     }
@@ -160,8 +172,14 @@ public class DisplaysHolderImpl implements DisplaysHolder {
     @Nullable
     public Object getDisplayOrigin(Display display) {
         synchronized (this.displaysBase) {
-            return this.displaysBase.get(display);
+            Object origin = this.displaysBase.get(display);
+            
+            if (origin != null) {
+                return origin;
+            }
         }
+        
+        return DisplayHistoryManager.INSTANCE.getPossibleOrigin(this, display);
     }
     
     private static class DisplaysList extends ArrayList<Display> {