Refactor MenuInfoProvider to give more concrete info

4 files changed, 75 insertions, 43 deletions

diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/client/view/ViewsImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/client/view/ViewsImpl.java
index d5e73916d..a76cce52f 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/client/view/ViewsImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/client/view/ViewsImpl.java
@@ -280,44 +280,46 @@ public class ViewsImpl implements Views {
         AbstractContainerMenu menu = Minecraft.getInstance().player.containerMenu;
         Set<EntryStack<?>> craftables = new HashSet<>();
         for (Map.Entry<CategoryIdentifier<?>, List<Display>> entry : DisplayRegistry.getInstance().getAll().entrySet()) {
-            MenuInfo<AbstractContainerMenu, Display> info = menu != null ?
-                    (MenuInfo<AbstractContainerMenu, Display>) MenuInfoRegistry.getInstance().get(entry.getKey(), menu.getClass())
-                    : null;
-            
-            class InfoContext implements MenuInfoContext<AbstractContainerMenu, LocalPlayer, Display> {
-                private Display display;
-                
-                @Override
-                public AbstractContainerMenu getMenu() {
-                    return menu;
-                }
-                
-                @Override
-                public LocalPlayer getPlayerEntity() {
-                    return Minecraft.getInstance().player;
-                }
-                
-                @Override
-                public MenuInfo<AbstractContainerMenu, Display> getContainerInfo() {
-                    return info;
-                }
+            List<Display> displays = entry.getValue();
+            for (Display display : displays) {
+                MenuInfo<AbstractContainerMenu, Display> info = menu != null ?
+                        MenuInfoRegistry.getInstance().getClient(display, menu)
+                        : null;
                 
-                @Override
-                public CategoryIdentifier<Display> getCategoryIdentifier() {
-                    return (CategoryIdentifier<Display>) entry.getKey();
+                class InfoContext implements MenuInfoContext<AbstractContainerMenu, LocalPlayer, Display> {
+                    private Display display;
+                    
+                    public InfoContext(Display display) {
+                        this.display = display;
+                    }
+                    
+                    @Override
+                    public AbstractContainerMenu getMenu() {
+                        return menu;
+                    }
+                    
+                    @Override
+                    public LocalPlayer getPlayerEntity() {
+                        return Minecraft.getInstance().player;
+                    }
+                    
+                    @Override
+                    public MenuInfo<AbstractContainerMenu, Display> getContainerInfo() {
+                        return info;
+                    }
+                    
+                    @Override
+                    public CategoryIdentifier<Display> getCategoryIdentifier() {
+                        return (CategoryIdentifier<Display>) entry.getKey();
+                    }
+                    
+                    @Override
+                    public Display getDisplay() {
+                        return display;
+                    }
                 }
                 
-                @Override
-                public Display getDisplay() {
-                    return display;
-                }
-            }
-            
-            InfoContext context = new InfoContext();
-            
-            List<Display> displays = entry.getValue();
-            for (Display display : displays) {
-                context.display = display;
+                InfoContext context = new InfoContext(display);
                 Iterable<SlotAccessor> inputSlots = info != null ? info.getInputSlots(context) : Collections.emptySet();
                 int slotsCraftable = 0;
                 List<EntryIngredient> requiredInput = display.getRequiredEntries();
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/InputSlotCrafter.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/InputSlotCrafter.java
index 8c5df5b53..aa233ca85 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/InputSlotCrafter.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/InputSlotCrafter.java
@@ -55,16 +55,24 @@ public class InputSlotCrafter<T extends AbstractContainerMenu, C extends Contain
     private Iterable<SlotAccessor> inventoryStacks;
     private ServerPlayer player;
     
-    private InputSlotCrafter(CategoryIdentifier<D> category, T container, CompoundTag display, MenuInfo<T, D> menuInfo) {
+    private InputSlotCrafter(CategoryIdentifier<D> category, T container) {
         this.category = category;
         this.container = container;
+    }
+    
+    public void setDisplay(D display) {
+        this.display = display;
+    }
+    
+    public void setMenuInfo(MenuInfo<T, D> menuInfo) {
         this.menuInfo = menuInfo;
-        this.display = menuInfo.read(this, display);
     }
     
     public static <T extends AbstractContainerMenu, C extends Container, D extends Display> InputSlotCrafter<T, C, D> start(CategoryIdentifier<D> category, T menu, ServerPlayer player, CompoundTag display, boolean hasShift) {
-        MenuInfo<T, D> menuInfo = Objects.requireNonNull(MenuInfoRegistry.getInstance().get(category, (Class<T>) menu.getClass()), "Container Info does not exist on the server!");
-        InputSlotCrafter<T, C, D> crafter = new InputSlotCrafter<>(category, menu, display, menuInfo);
+        InputSlotCrafter<T, C, D> crafter = new InputSlotCrafter<>(category, menu);
+        MenuInfo<T, D> menuInfo = Objects.requireNonNull(MenuInfoRegistry.getInstance().get(category, menu, crafter, display), "Container Info does not exist on the server!");
+        crafter.setMenuInfo(menuInfo);
+        crafter.setDisplay(menuInfo.read(crafter, display));
         crafter.fillInputSlots(player, hasShift);
         return crafter;
     }
diff --git a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
index bf84c1808..a42cf8d89 100644
--- a/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
+++ b/runtime/src/main/java/me/shedaniel/rei/impl/common/transfer/MenuInfoRegistryImpl.java
@@ -31,12 +31,17 @@ import me.shedaniel.rei.api.common.plugins.REIServerPlugin;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfo;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfoProvider;
 import me.shedaniel.rei.api.common.transfer.info.MenuInfoRegistry;
+import me.shedaniel.rei.api.common.transfer.info.MenuSerializationProviderContext;
+import net.fabricmc.api.EnvType;
+import net.fabricmc.api.Environment;
+import net.minecraft.nbt.CompoundTag;
 import net.minecraft.world.inventory.AbstractContainerMenu;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 import java.util.function.Predicate;
 
 public class MenuInfoRegistryImpl implements MenuInfoRegistry {
@@ -58,11 +63,28 @@ public class MenuInfoRegistryImpl implements MenuInfoRegistry {
     @Override
     @Nullable
     public <T extends AbstractContainerMenu, D extends Display> MenuInfo<T, D> get(CategoryIdentifier<D> category, Class<T> menuClass) {
+        return getInternal(category, menuClass, provider -> provider.provide(category, menuClass));
+    }
+    
+    @Override
+    @Nullable
+    @Environment(EnvType.CLIENT)
+    public <C extends AbstractContainerMenu, D extends Display> MenuInfo<C, D> getClient(D display, C menu) {
+        return getInternal((CategoryIdentifier<D>) display.getCategoryIdentifier(), (Class<C>) menu.getClass(), provider -> provider.provideClient(display, menu));
+    }
+    
+    @Override
+    @Nullable
+    public <C extends AbstractContainerMenu, D extends Display> MenuInfo<C, D> get(CategoryIdentifier<D> category, C menu, MenuSerializationProviderContext<C, ?, D> context, CompoundTag tag) {
+        return getInternal(category, (Class<C>) menu.getClass(), provider -> provider.provide(category, menu, context, tag));
+    }
+    
+    private <C extends AbstractContainerMenu, D extends Display> MenuInfo<C, D> getInternal(CategoryIdentifier<D> category, Class<C> menuClass, Function<MenuInfoProvider<C, D>, Optional<MenuInfo<C, D>>> function) {
         Map<Class<? extends AbstractContainerMenu>, List<MenuInfoProvider<?, ?>>> infoMap = map.get(category);
         if (infoMap != null && !infoMap.isEmpty()) {
             if (infoMap.containsKey(menuClass)) {
                 for (MenuInfoProvider<?, ?> provider : infoMap.get(menuClass)) {
-                    Optional<MenuInfo<T, D>> info = ((MenuInfoProvider<T, D>) provider).provide(category, menuClass);
+                    Optional<MenuInfo<C, D>> info = function.apply((MenuInfoProvider<C, D>) provider);
                     if (info.isPresent()) {
                         return info.get();
                     }
@@ -71,7 +93,7 @@ public class MenuInfoRegistryImpl implements MenuInfoRegistry {
             for (Map.Entry<Class<? extends AbstractContainerMenu>, List<MenuInfoProvider<?, ?>>> entry : infoMap.entrySet()) {
                 if (entry.getKey().isAssignableFrom(menuClass)) {
                     for (MenuInfoProvider<?, ?> provider : entry.getValue()) {
-                        Optional<MenuInfo<T, D>> info = ((MenuInfoProvider<T, D>) provider).provide(category, menuClass);
+                        Optional<MenuInfo<C, D>> info = function.apply((MenuInfoProvider<C, D>) provider);
                         if (info.isPresent()) {
                             return info.get();
                         }
@@ -84,7 +106,7 @@ public class MenuInfoRegistryImpl implements MenuInfoRegistry {
             if (entry.getKey().test(category) && !entry.getValue().isEmpty()) {
                 List<MenuInfoProvider<?, ?>> infoList = entry.getValue();
                 if (!infoList.isEmpty()) {
-                    Optional<MenuInfo<T, D>> info = ((MenuInfoProvider<T, D>) infoList.get(0)).provide(category, menuClass);
+                    Optional<MenuInfo<C, D>> info = function.apply((MenuInfoProvider<C, D>) infoList.get(0));
                     if (info.isPresent()) {
                         return info.get();
                     }
diff --git a/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java b/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
index 14f50a0d7..1c9703d5e 100644
--- a/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
+++ b/runtime/src/main/java/me/shedaniel/rei/plugin/autocrafting/DefaultCategoryHandler.java
@@ -61,7 +61,7 @@ public class DefaultCategoryHandler implements TransferHandler {
             return Result.createNotApplicable();
         }
         AbstractContainerMenu menu = context.getMenu();
-        MenuInfo<AbstractContainerMenu, Display> menuInfo = MenuInfoRegistry.getInstance().get((CategoryIdentifier<Display>) display.getCategoryIdentifier(), (Class<AbstractContainerMenu>) menu.getClass());
+        MenuInfo<AbstractContainerMenu, Display> menuInfo = MenuInfoRegistry.getInstance().getClient(display, menu);
         if (menuInfo == null) {
             return Result.createNotApplicable();
         }