Improve the functionality of wrapped DataWatcher and WatchableObject.

ProtocolLib/src/main/java/com/comphenix/protocol/utility/MinecraftReflection.java

@@ -78,6 +78,16 @@ public class MinecraftReflection {
 		}
 	}
 	
+	/**
+	 * Used during debugging and testing.
+	 * @param minecraftPackage - the current Minecraft package.
+	 * @param craftBukkitPackage - the current CraftBukkit package.
+	 */
+	public static void setMinecraftPackage(String minecraftPackage, String craftBukkitPackage) {
+		MINECRAFT_FULL_PACKAGE = minecraftPackage;
+		CRAFTBUKKIT_PACKAGE = craftBukkitPackage;
+	}
+	
 	/**
 	 * Retrieve the name of the root CraftBukkit package.
 	 * @return Full canonical name of the root CraftBukkit package.

ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/WrappedDataWatcher.java

@@ -6,6 +6,7 @@ import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -13,6 +14,8 @@ import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
+import javax.annotation.Nullable;
+
 import org.bukkit.entity.Entity;
 import org.bukkit.inventory.ItemStack;
 
@@ -21,14 +24,16 @@ import com.comphenix.protocol.reflect.FieldAccessException;
 import com.comphenix.protocol.reflect.FieldUtils;
 import com.comphenix.protocol.reflect.FuzzyReflection;
 import com.comphenix.protocol.utility.MinecraftReflection;
+import com.google.common.base.Function;
 import com.google.common.base.Objects;
+import com.google.common.collect.Iterators;
 
 /**
  * Wraps a DataWatcher that is used to transmit arbitrary key-value pairs with a given entity.
  * 
  * @author Kristian
  */
-public class WrappedDataWatcher {
+public class WrappedDataWatcher implements Iterable<WrappedWatchableObject> {
 
 	/**
 	 * Used to assign integer IDs to given types.
@@ -92,16 +97,31 @@ public class WrappedDataWatcher {
 	}
 	
 	/**
-	 * Create a new data watcher from a list of watchable objects.
+	 * Create a new data watcher for a list of watchable objects.
+	 * <p>
+	 * Note that the watchable objects are not cloned, and will be modified in place. Use "deepClone" if 
+	 * that is not desirable.
+	 * <p>
+	 * The {@link #removeObject(int)} method will not modify the given list, however.
+	 * 
 	 * @param watchableObjects - list of watchable objects that will be copied.
 	 * @throws FieldAccessException Unable to read watchable objects.
 	 */
 	public WrappedDataWatcher(List<WrappedWatchableObject> watchableObjects) throws FieldAccessException {
 		this();
+
+		Lock writeLock = getReadWriteLock().writeLock();
+		Map<Integer, Object> map = getWatchableObjectMap();
 		
-		// Fill the underlying map
-		for (WrappedWatchableObject watched : watchableObjects) {
-			setObject(watched.getIndex(), watched.getValue());
+		writeLock.lock();
+		
+		try {
+			// Add the watchable objects by reference
+			for (WrappedWatchableObject watched : watchableObjects) {
+				map.put(watched.getIndex(), watched.handle);
+			}			
+		} finally {
+			writeLock.unlock();
 		}
 	}
 	
@@ -234,9 +254,10 @@ public class WrappedDataWatcher {
      * @throws FieldAccessException If reflection failed.
      */
 	public List<WrappedWatchableObject> getWatchableObjects() throws FieldAccessException {
+		Lock readLock = getReadWriteLock().readLock();
+		readLock.lock();
+		
     	try {
-    		getReadWriteLock().readLock().lock();
-    		
     		List<WrappedWatchableObject> result = new ArrayList<WrappedWatchableObject>();
     		
     		// Add each watchable object to the list
@@ -250,7 +271,7 @@ public class WrappedDataWatcher {
     		return result;
     		
     	} finally {
-    		getReadWriteLock().readLock().unlock();
+    		readLock.unlock();
     	}
     }
 
@@ -270,6 +291,20 @@ public class WrappedDataWatcher {
     	}
     }
     
+    /**
+     * Clone the content of the current DataWatcher.
+     * @return A cloned data watcher.
+     */
+    public WrappedDataWatcher deepClone() {
+    	WrappedDataWatcher clone = new WrappedDataWatcher();
+    	
+    	// Make a new copy instead
+    	for (WrappedWatchableObject watchable : this) {
+    		clone.setObject(watchable.getIndex(), watchable.getValue());
+    	}
+    	return clone;
+    }
+    
     /**
      * Retrieve the number of watched objects.
      * @return Watched object count.
@@ -286,6 +321,23 @@ public class WrappedDataWatcher {
     	}
     }
     
+    /**
+     * Remove a given object from the underlying DataWatcher.
+     * @param index - index of the object to remove.
+     * @return The watchable object that was removed, or NULL If none could be found.
+     */
+    public WrappedWatchableObject removeObject(int index) {
+    	Lock writeLock = getReadWriteLock().writeLock();
+    	writeLock.lock();
+    	
+    	try {
+    		Object removed = getWatchableObjectMap().remove(index);
+    		return removed != null ? new WrappedWatchableObject(removed) : null;
+    	} finally {
+    		writeLock.unlock();
+    	}
+    }
+    
     /**
      * Set a watched byte.
      * @param index - index of the watched byte.
@@ -480,4 +532,20 @@ public class WrappedDataWatcher {
 			// Use fallback method
 		}
 	}
+
+	@Override
+	public Iterator<WrappedWatchableObject> iterator() {
+		// We'll wrap the iterator instead of creating a new list every time 
+		return Iterators.transform(getWatchableObjectMap().values().iterator(), 
+				new Function<Object, WrappedWatchableObject>() {
+			
+			@Override
+			public WrappedWatchableObject apply(@Nullable Object item) {
+				if (item != null)
+					return new WrappedWatchableObject(item);
+				else
+					return null;
+			}
+		});
+	}
 }

ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/WrappedWatchableObject.java

@@ -9,6 +9,7 @@ import com.comphenix.protocol.reflect.FieldAccessException;
 import com.comphenix.protocol.reflect.StructureModifier;
 import com.comphenix.protocol.reflect.instances.DefaultInstances;
 import com.comphenix.protocol.utility.MinecraftReflection;
+import com.google.common.base.Objects;
 
 /**
  * Represents a watchable object.
@@ -26,6 +27,9 @@ public class WrappedWatchableObject {
 	// Used to create new watchable objects
 	private static Constructor<?> watchableConstructor;
 	
+	// The watchable object class type
+	private static Class<?> watchableObjectClass;
+	
 	protected Object handle;
 	protected StructureModifier<Object> modifier;
 	
@@ -79,6 +83,12 @@ public class WrappedWatchableObject {
 		initialize();
 		this.handle = handle;
 		this.modifier = baseModifier.withTarget(handle);
+		
+		// Make sure the type is correct
+		if (!watchableObjectClass.isAssignableFrom(handle.getClass())) {
+			throw new ClassCastException("Cannot cast the class " + handle.getClass().getName() +
+										 " to " + watchableObjectClass.getName());
+		}
 	}
 	
 	/**
@@ -95,7 +105,8 @@ public class WrappedWatchableObject {
 	private static void initialize() {
 		if (!hasInitialized) {
 			hasInitialized = true;
-			baseModifier = new StructureModifier<Object>(MinecraftReflection.getWatchableObjectClass(), null, false);
+			watchableObjectClass = MinecraftReflection.getWatchableObjectClass();
+			baseModifier = new StructureModifier<Object>(watchableObjectClass, null, false);
 		}
 	}
 	
@@ -310,4 +321,34 @@ public class WrappedWatchableObject {
 			return value;
 		}
 	}
+	
+	@Override
+	public boolean equals(Object obj) {
+		// Quick checks
+		if (obj == this)
+			return true;
+		if (obj == null)
+			return false;
+		
+		if (obj instanceof WrappedWatchableObject) {
+			WrappedWatchableObject other = (WrappedWatchableObject) obj;
+			
+			return Objects.equal(getIndex(), other.getIndex()) &&
+				   Objects.equal(getTypeID(), other.getTypeID()) &&
+				   Objects.equal(getValue(), other.getValue());
+		}
+		
+		// No, this is not equivalent
+		return false;
+	}
+	
+	@Override
+	public int hashCode() {
+		return Objects.hashCode(getIndex(), getTypeID(), getValue());
+	}
+	
+	@Override
+	public String toString() {
+		return String.format("[%s: %s (%s)]", getIndex(), getValue(), getType().getSimpleName());
+	}
 }