2828import com .sk89q .worldedit .math .BlockVector3 ;
2929import com .sk89q .worldedit .util .Location ;
3030import com .sk89q .worldedit .util .concurrency .LazyReference ;
31- import com .sk89q .worldedit .world .block .BlockTypes ;
3231import com .sk89q .worldedit .world .entity .EntityTypes ;
3332import net .minecraft .core .BlockPos ;
3433import net .minecraft .core .registries .Registries ;
3736import net .minecraft .server .level .ServerLevel ;
3837import net .minecraft .world .entity .Entity ;
3938import net .minecraft .world .level .WorldGenLevel ;
39+ import net .minecraft .world .level .block .Blocks ;
40+ import net .minecraft .world .level .block .EntityBlock ;
4041import net .minecraft .world .level .block .entity .BlockEntity ;
4142import net .minecraft .world .level .block .state .BlockState ;
4243import net .minecraft .world .phys .Vec3 ;
5051import java .lang .reflect .InvocationHandler ;
5152import java .lang .reflect .Method ;
5253import java .lang .reflect .Proxy ;
54+ import java .util .HashMap ;
55+ import java .util .Map ;
5356
54- public class PaperweightServerLevelDelegateProxy implements InvocationHandler {
57+ public class PaperweightServerLevelDelegateProxy implements InvocationHandler , AutoCloseable {
58+
59+ private static BlockVector3 adapt (BlockPos blockPos ) {
60+ return BlockVector3 .at (blockPos .getX (), blockPos .getY (), blockPos .getZ ());
61+ }
5562
5663 private final EditSession editSession ;
5764 private final ServerLevel serverLevel ;
5865 private final PaperweightAdapter adapter ;
66+ private final Map <BlockVector3 , BlockEntity > createdBlockEntities = new HashMap <>();
5967
6068 private PaperweightServerLevelDelegateProxy (EditSession editSession , ServerLevel serverLevel , PaperweightAdapter adapter ) {
6169 this .editSession = editSession ;
6270 this .serverLevel = serverLevel ;
6371 this .adapter = adapter ;
6472 }
6573
66- public static WorldGenLevel newInstance (EditSession editSession , ServerLevel serverLevel , PaperweightAdapter adapter ) {
67- return (WorldGenLevel ) Proxy .newProxyInstance (
68- serverLevel .getClass ().getClassLoader (),
69- serverLevel .getClass ().getInterfaces (),
70- new PaperweightServerLevelDelegateProxy (editSession , serverLevel , adapter )
74+ public record LevelAndProxy (WorldGenLevel level , PaperweightServerLevelDelegateProxy proxy ) implements AutoCloseable {
75+ @ Override
76+ public void close () throws MaxChangedBlocksException {
77+ proxy .close ();
78+ }
79+ }
80+
81+ public static LevelAndProxy newInstance (EditSession editSession , ServerLevel serverLevel , PaperweightAdapter adapter ) {
82+ PaperweightServerLevelDelegateProxy proxy = new PaperweightServerLevelDelegateProxy (editSession , serverLevel , adapter );
83+ return new LevelAndProxy (
84+ (WorldGenLevel ) Proxy .newProxyInstance (
85+ serverLevel .getClass ().getClassLoader (),
86+ serverLevel .getClass ().getInterfaces (),
87+ proxy
88+ ),
89+ proxy
7190 );
7291 }
7392
7493 @ Nullable
7594 private BlockEntity getBlockEntity (BlockPos blockPos ) {
76- BlockEntity tileEntity = this .serverLevel .getChunkAt (blockPos ).getBlockEntity (blockPos );
77- if (tileEntity == null ) {
78- return null ;
79- }
80- tileEntity .loadWithComponents (
81- (CompoundTag ) adapter .fromNative (this .editSession .getFullBlock (BlockVector3 .at (blockPos .getX (), blockPos .getY (), blockPos .getZ ())).getNbtReference ().getValue ()),
82- this .serverLevel .registryAccess ()
83- );
84-
85- return tileEntity ;
95+ // This doesn't synthesize or load from world. I think editing existing block entities without setting the block
96+ // (in the context of features) should not be supported in the first place.
97+ BlockVector3 pos = adapt (blockPos );
98+ return createdBlockEntities .get (pos );
8699 }
87100
88101 private BlockState getBlockState (BlockPos blockPos ) {
89- return adapter .adapt (this .editSession .getBlock (BlockVector3 . at (blockPos . getX (), blockPos . getY (), blockPos . getZ () )));
102+ return adapter .adapt (this .editSession .getBlock (adapt (blockPos )));
90103 }
91104
92105 private boolean setBlock (BlockPos blockPos , BlockState blockState ) {
93106 try {
94- return editSession .setBlock (BlockVector3 .at (blockPos .getX (), blockPos .getY (), blockPos .getZ ()), adapter .adapt (blockState ));
107+ handleBlockEntity (blockPos , blockState );
108+ return editSession .setBlock (adapt (blockPos ), adapter .adapt (blockState ));
95109 } catch (MaxChangedBlocksException e ) {
96110 throw new RuntimeException (e );
97111 }
98112 }
99113
100- private boolean removeBlock (BlockPos blockPos ) {
101- try {
102- return editSession .setBlock (BlockVector3 .at (blockPos .getX (), blockPos .getY (), blockPos .getZ ()), BlockTypes .AIR .getDefaultState ());
103- } catch (MaxChangedBlocksException e ) {
104- throw new RuntimeException (e );
114+ // For BlockEntity#setBlockState, not sure why it's deprecated
115+ @ SuppressWarnings ("deprecation" )
116+ private void handleBlockEntity (BlockPos blockPos , BlockState blockState ) {
117+ BlockVector3 pos = adapt (blockPos );
118+ if (blockState .hasBlockEntity ()) {
119+ if (!(blockState .getBlock () instanceof EntityBlock entityBlock )) {
120+ // This will probably never happen, as Mojang's own code assumes that
121+ // hasBlockEntity implies instanceof EntityBlock, but just to be safe...
122+ throw new AssertionError ("BlockState has block entity but block is not an EntityBlock: " + blockState );
123+ }
124+ BlockEntity newEntity = entityBlock .newBlockEntity (blockPos , blockState );
125+ if (newEntity != null ) {
126+ newEntity .setBlockState (blockState );
127+ createdBlockEntities .put (pos , newEntity );
128+ // Should we load existing NBT here? This is for feature / structure gen so it seems unnecessary.
129+ // But it would align with the behavior of the real setBlock method.
130+ return ;
131+ }
105132 }
133+ // Discard any block entity that was previously created if new block is set without block entity
134+ createdBlockEntities .remove (pos );
135+ }
136+
137+ private boolean removeBlock (BlockPos blockPos , boolean bl ) {
138+ return setBlock (blockPos , Blocks .AIR .defaultBlockState ());
106139 }
107140
108141 private boolean addEntity (Entity entity ) {
@@ -117,6 +150,20 @@ private boolean addEntity(Entity entity) {
117150 return editSession .createEntity (location , baseEntity ) != null ;
118151 }
119152
153+ @ Override
154+ public void close () throws MaxChangedBlocksException {
155+ for (Map .Entry <BlockVector3 , BlockEntity > entry : createdBlockEntities .entrySet ()) {
156+ BlockVector3 blockPos = entry .getKey ();
157+ BlockEntity blockEntity = entry .getValue ();
158+ net .minecraft .nbt .CompoundTag tag = blockEntity .saveWithId (serverLevel .registryAccess ());
159+ editSession .setBlock (
160+ blockPos ,
161+ adapter .adapt (blockEntity .getBlockState ())
162+ .toBaseBlock (LazyReference .from (() -> (LinCompoundTag ) adapter .toNative (tag )))
163+ );
164+ }
165+ }
166+
120167 private static void addMethodHandleToTable (
121168 ImmutableTable .Builder <String , MethodType , MethodHandle > table ,
122169 String methodName ,
0 commit comments