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);
}
}