Block State Item Models

MixinItemModelShaper.java

@Mixin(ItemModelShaper.class)
public class MixinItemModelShaper {

    @Shadow @Final private ModelManager modelManager;
    @Unique
    private static final Int2ObjectMap<ModelResourceLocation> blockstatemodels$cache = new Int2ObjectOpenHashMap<>();

    @Inject(method = "getItemModel(Lnet/minecraft/world/item/ItemStack;)Lnet/minecraft/client/resources/model/BakedModel;", at = @At("HEAD"), cancellable = true)
    public void onGetModel(ItemStack stack, CallbackInfoReturnable<BakedModel> cir)  {
        final Item item = stack.getItem();
        if (item instanceof BlockItem blockItem && blockstatemodels$canShowBlockStateModel(item)) {
            if (stack.hasTag() && stack.getTag().contains("BlockStateTag", Tag.TAG_COMPOUND)) {
                final CompoundTag tag = stack.getTag().getCompound("BlockStateTag");
                final int hash = tag.hashCode();
                if (!blockstatemodels$cache.containsKey(hash)) {
                    final Block block = blockItem.getBlock();
                    final Map<String, String> states = new HashMap<>();
                    final BlockState state = block.defaultBlockState();
                    //noinspection rawtypes
                    for (Property property : state.getProperties()) states.put(property.getName(), blockstatemodels$getPropertyValue(property, state));
                    for (String key : tag.getAllKeys()) states.put(key, tag.getString(key));

                    final String[] stringStates = states.entrySet().stream().map(entry -> entry.getKey() + "=" + entry.getValue()).toArray(String[]::new);
                    Arrays.sort(stringStates);
                    final ResourceLocation id = blockstatemodels$hasValue(block);
                    if (id != null) {
                        blockstatemodels$cache.put(hash, new ModelResourceLocation(id, String.join(",", stringStates)));
                    }
                }
                if (blockstatemodels$cache.containsKey(hash)) {
                    final ModelResourceLocation location = blockstatemodels$cache.get(hash);
                    if (location != null) {
                        final BakedModel model = this.modelManager.getModel(location);
                        if (model != this.modelManager.getMissingModel()) {
                            cir.setReturnValue(model);
                        } else {
                            blockstatemodels$cache.put(hash, null);
                        }
                    }
                }
            }
        }
    }

    @Unique
    @SuppressWarnings({"rawtypes", "unchecked"})
    private static String blockstatemodels$getPropertyValue(Property property, BlockState state) {
        return property.getName(state.getValue(property));
    }

    @Unique
    private static boolean blockstatemodels$canShowBlockStateModel(Item item) {
        if (item instanceof BlockItem blockItem) {
            if (item instanceof ModelItemExtension extension && extension.canShowBlockStateModel()) {
                return true;
            }
            return blockItem.getBlock() instanceof ModelItemExtension extension && extension.canShowBlockStateModel();
        }
        return false;
    }

    @Unique
    private static ResourceLocation blockstatemodels$hasValue(Block block) {
        ResourceLocation key = Registry.BLOCK.getKey(block);
        if (key == Registry.BLOCK.getDefaultKey() && Registry.BLOCK.get(key) != block) return null;
        return key;
    }
}