/*
 * This file is part of Sponge, licensed under the MIT License (MIT).
 *
 * Copyright (c) SpongePowered <https://www.spongepowered.org>
 * Copyright (c) contributors
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
package org.spongepowered.common.mixin.core.world.gen;

import net.minecraft.util.LongHashMap;
import net.minecraft.world.ChunkCoordIntPair;
import net.minecraft.world.WorldServer;
import net.minecraft.world.chunk.Chunk;
import net.minecraft.world.chunk.IChunkProvider;
import net.minecraft.world.gen.ChunkProviderServer;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Redirect;
import org.spongepowered.common.interfaces.world.IMixinWorld;
import org.spongepowered.common.interfaces.world.gen.IMixinChunkProviderServer;

import javax.annotation.Nullable;

@Mixin(ChunkProviderServer.class)
public abstract class MixinChunkProviderServer implements IMixinChunkProviderServer {

    @Shadow public WorldServer worldObj;
    @Shadow private LongHashMap<Chunk> id2ChunkMap;

    @Shadow public abstract Chunk provideChunk(int x, int z);

    @Nullable
    @Override
    public Chunk getChunkIfLoaded(int x, int z) {
        return this.id2ChunkMap.getValueByKey(ChunkCoordIntPair.chunkXZ2Int(x, z));
    }

    @Redirect(method = "populate",
            at = @At(value = "INVOKE", target = "Lnet/minecraft/world/chunk/IChunkProvider;populate(Lnet/minecraft/world/chunk/IChunkProvider;II)V"))
    public void onChunkPopulate(IChunkProvider serverChunkGenerator, IChunkProvider chunkProvider, int x, int z) {
        IMixinWorld world = (IMixinWorld) this.worldObj;
        boolean capturingTerrain = world.getCauseTracker().isCapturingTerrainGen();
        boolean processingCapture = world.getCauseTracker().isProcessingCaptureCause();
        world.getCauseTracker().setProcessingCaptureCause(true);
        world.getCauseTracker().setCapturingTerrainGen(true);
        serverChunkGenerator.populate(chunkProvider, x, z);
        world.getCauseTracker().setCapturingTerrainGen(capturingTerrain);
        world.getCauseTracker().setProcessingCaptureCause(processingCapture);
    }
}

@@ -1644,63 +1644,17 @@
 xin.
-Shadow;%0Aimport org.spongepowered.asm.mixin.injection.At
+Overwrite
 ;%0Aim
@@ -1690,26 +1690,14 @@
 xin.
-injection.Redirect
+Shadow
 ;%0Aim
@@ -1879,16 +1879,24 @@
 %0A@Mixin(
+value = 
 ChunkPro
@@ -1912,16 +1912,33 @@
 er.class
+, priority = 1000
 )%0Apublic
@@ -2110,16 +2110,73 @@
 hunkMap;
+%0A    @Shadow public IChunkProvider serverChunkGenerator;;
 %0A%0A    @S
@@ -2409,36 +2409,141 @@
     
-@Redirect(method = %22
+/**%0A     * @author blood - May 9th, 2016%0A     * @reason Control terrain gen flag here to avoid leaking block captures during 
 populate
 %22,%0A 
@@ -2542,160 +2542,175 @@
 late
-%22,
+.
 %0A     
+*
      
+%0A
   
-at = @At(value = %22INVOKE%22, target = %22Lnet/minecraft/world/chunk/IC
+   * @param chunkProvider The c
 hunk
-P
+ p
 rovider
-;populate(Lnet/minecraft/world/chunk/IChunkProvider;II)V%22))
+%0A     * @param x The x coordinate of chunk%0A     * @param z The z coordinate of chunk%0A     */%0A    @Overwrite
 %0A   
@@ -2726,16 +2726,9 @@
 oid 
-onChunkP
+p
 opul
@@ -2750,76 +2750,227 @@
 der 
-serverChunkGenerator, IChunkProvider chunkProvider, int x, int z) %7B%0A
+chunkProvider, int x, int z) %7B%0A        Chunk chunk = this.provideChunk(x, z);%0A%0A        if (!chunk.isTerrainPopulated()) %7B%0A            chunk.func_150809_p();%0A%0A            if (this.serverChunkGenerator != null) %7B%0A        
     
@@ -3022,32 +3022,40 @@
 rldObj;%0A        
+   
+     
 boolean capturin
@@ -3114,32 +3114,40 @@
 nGen();%0A        
+      
+  
 boolean processi
@@ -3206,24 +3206,32 @@
 ureCause();%0A
+        
         worl
@@ -3275,32 +3275,40 @@
 ureCause(true);%0A
+        
         world.ge
@@ -3361,16 +3361,29 @@
         
+        this.
 serverCh
@@ -3426,16 +3426,66 @@
  x, z);%0A
+                chunk.setChunkModified();%0A        
         
@@ -3550,32 +3550,40 @@
 rrain);%0A        
+  
+      
 world.getCauseTr
@@ -3632,20 +3632,44 @@
 ngCapture);%0A
+            %7D%0A        %7D%0A
     %7D%0A%7D%0A

Replace MixinChunkProviderServer populate redirect with overwrite.