Add a -m flag to the //distr command to filter included blocks

worldedit-core/src/main/java/com/sk89q/worldedit/EditSession.java

@@ -2281,7 +2281,18 @@ public int makeForest(Region region, double density, TreeGenerator.TreeType tree
      * @return the results
      */
     public List<Countable<BlockState>> getBlockDistribution(Region region, boolean separateStates) {
-        BlockDistributionCounter count = new BlockDistributionCounter(this, separateStates);
+        return getBlockDistribution(region, Masks.alwaysTrue(), separateStates);
+    }
+
+    /**
+     * Get the block distribution inside a region.
+     *
+     * @param region a region
+     * @param mask a mask to filter the blocks to count
+     * @return the results
+     */
+    public List<Countable<BlockState>> getBlockDistribution(Region region, @Nullable Mask mask, boolean separateStates) {
+        BlockDistributionCounter count = new BlockDistributionCounter(this, mask, separateStates);
         RegionVisitor visitor = new RegionVisitor(region, count);
         Operations.completeBlindly(visitor);
         return count.getDistribution();

worldedit-core/src/main/java/com/sk89q/worldedit/command/SelectionCommands.java

@@ -701,19 +701,21 @@ public void distr(Actor actor, World world, LocalSession session,
                       @Switch(name = 'd', desc = "Separate blocks by state")
                           boolean separateStates,
                       @ArgFlag(name = 'p', desc = "Gets page from a previous distribution.")
-                          Integer page) throws WorldEditException {
+                          Integer page,
+                      @ArgFlag(name = 'm', desc = "Only include blocks matching the given mask")
+                          Mask sourceMask) throws WorldEditException {
         List<Countable<BlockState>> distribution;
 
         if (page == null) {
             if (clipboardDistr) {
                 Clipboard clipboard = session.getClipboard().getClipboard(); // throws if missing
-                BlockDistributionCounter count = new BlockDistributionCounter(clipboard, separateStates);
+                BlockDistributionCounter count = new BlockDistributionCounter(clipboard, sourceMask, separateStates);
                 RegionVisitor visitor = new RegionVisitor(clipboard.getRegion(), count);
                 Operations.completeBlindly(visitor);
                 distribution = count.getDistribution();
             } else {
                 try (EditSession editSession = session.createEditSession(actor)) {
-                    distribution = editSession.getBlockDistribution(session.getSelection(world), separateStates);
+                    distribution = editSession.getBlockDistribution(session.getSelection(world), sourceMask, separateStates);
                 }
             }
             session.setLastDistribution(distribution);

worldedit-core/src/main/java/com/sk89q/worldedit/function/block/BlockDistributionCounter.java

@@ -22,9 +22,12 @@
 import com.sk89q.worldedit.WorldEditException;
 import com.sk89q.worldedit.extent.Extent;
 import com.sk89q.worldedit.function.RegionFunction;
+import com.sk89q.worldedit.function.mask.Mask;
+import com.sk89q.worldedit.function.mask.Masks;
 import com.sk89q.worldedit.math.BlockVector3;
 import com.sk89q.worldedit.util.Countable;
 import com.sk89q.worldedit.world.block.BlockState;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -35,18 +38,28 @@
 public class BlockDistributionCounter implements RegionFunction {
 
     private final Extent extent;
+    private final Mask mask;
     private final boolean separateStates;
 
     private final List<Countable<BlockState>> distribution = new ArrayList<>();
     private final Map<BlockState, Countable<BlockState>> map = new HashMap<>();
 
     public BlockDistributionCounter(Extent extent, boolean separateStates) {
+        this(extent, null, separateStates);
+    }
+
+    public BlockDistributionCounter(Extent extent, @Nullable Mask mask, boolean separateStates) {
         this.extent = extent;
+        this.mask = mask == null ? Masks.alwaysTrue() : mask;
         this.separateStates = separateStates;
     }
 
     @Override
     public boolean apply(BlockVector3 position) throws WorldEditException {
+        if (!mask.test(position)) {
+            return true;
+        }
+
         BlockState blk = extent.getBlock(position);
         if (!separateStates) {
             blk = blk.getBlockType().getDefaultState();