diff --git a/src/main/java/carpet/script/annotation/Locator.java b/src/main/java/carpet/script/annotation/Locator.java index f5fe1840..bbaafa94 100644 --- a/src/main/java/carpet/script/annotation/Locator.java +++ b/src/main/java/carpet/script/annotation/Locator.java @@ -18,6 +18,8 @@ import carpet.script.bundled.Module; import carpet.script.value.BlockValue; import carpet.script.value.FunctionValue; import carpet.script.value.Value; +import net.minecraft.block.BlockState; +import net.minecraft.util.math.BlockPos; import static java.lang.annotation.ElementType.PARAMETER; import static java.lang.annotation.ElementType.TYPE_USE; @@ -33,7 +35,7 @@ public interface Locator /** *

Represents that the annotated argument must be gotten by passing the arguments in there into a {@link BlockArgument} locator.

* - *

Must be used in either {@link BlockValue} or {@link BlockArgument} parameters

+ *

Must be used in either {@link BlockArgument}, {@link BlockValue} or {@link BlockPos} parameters

*/ @Documented @Retention(RUNTIME) @@ -139,7 +141,7 @@ public interface Locator private static class BlockLocator extends AbstractLocator { - private final boolean returnBlockValue; + private final java.util.function.Function returnFunction; private final boolean acceptString; private final boolean anyString; private final boolean optional; @@ -149,11 +151,24 @@ public interface Locator this.acceptString = annotation.acceptString(); this.anyString = annotation.anyString(); this.optional = annotation.optional(); - this.returnBlockValue = type == BlockValue.class; - if (returnBlockValue && (anyString || optional)) + if (type != BlockArgument.class && (anyString || optional)) throw new IllegalArgumentException("Can only use anyString or optional parameters of Locator.Block if targeting a BlockArgument"); - if (!returnBlockValue && type != BlockArgument.class) - throw new IllegalArgumentException("Locator.Block can only be used against BlockArgument or BlockValue types!"); + this.returnFunction = getReturnFunction(type); + if (returnFunction == null) + throw new IllegalArgumentException("Locator.Block can only be used against BlockArgument, BlockValue, BlockPos or BlockState types!"); + } + + @SuppressWarnings("unchecked") + private static java.util.function.Function getReturnFunction(Class type) { + if (type == BlockArgument.class) + return r -> (R) r; + if (type == BlockValue.class) + return r -> (R) r.block; + if (type == BlockPos.class) + return r -> (R) r.block.getPos(); + if (type == BlockState.class) + return r -> (R) r.block.getBlockState(); + return null; } @Override @@ -166,7 +181,7 @@ public interface Locator public R checkAndConvert(Iterator valueIterator, Context context, Context.Type theLazyT) { BlockArgument locator = BlockArgument.findIn((CarpetContext) context, valueIterator, 0, acceptString, optional, anyString); - return (R) (returnBlockValue ? locator.block : locator); + return returnFunction.apply(locator); } }