Fix infinite leash

src/main/scala/li/cil/oc/integration/opencomputers/DriverUpgradeLeash.scala

@@ -1,13 +1,12 @@
 package li.cil.oc.integration.opencomputers
 
-import li.cil.oc.Constants
-import li.cil.oc.api
 import li.cil.oc.api.driver.EnvironmentProvider
 import li.cil.oc.api.driver.item.HostAware
+import li.cil.oc.api.internal
 import li.cil.oc.api.network.EnvironmentHost
-import li.cil.oc.common.Slot
-import li.cil.oc.common.Tier
+import li.cil.oc.common.{Slot, Tier}
 import li.cil.oc.server.component
+import li.cil.oc.{Constants, api}
 import net.minecraft.entity.Entity
 import net.minecraft.item.ItemStack
 
@@ -18,7 +17,7 @@ object DriverUpgradeLeash extends Item with HostAware {
   override def createEnvironment(stack: ItemStack, host: EnvironmentHost) =
     if (host.world != null && host.world.isRemote) null
     else host match {
-      case entity: Entity => new component.UpgradeLeash(entity)
+      case entity: Entity with internal.Drone => new component.UpgradeLeash(entity)
       case _ => null
     }
 

src/main/scala/li/cil/oc/server/component/UpgradeLeash.scala

@@ -1,36 +1,30 @@
 package li.cil.oc.server.component
 
-import java.util
-import java.util.UUID
-
-import li.cil.oc.Constants
-import li.cil.oc.api.driver.DeviceInfo.DeviceAttribute
-import li.cil.oc.api.driver.DeviceInfo.DeviceClass
-import li.cil.oc.OpenComputers
-import li.cil.oc.api.Network
 import li.cil.oc.api.driver.DeviceInfo
-import li.cil.oc.api.machine.Arguments
-import li.cil.oc.api.machine.Callback
-import li.cil.oc.api.machine.Context
-import li.cil.oc.api.network.Node
-import li.cil.oc.api.network.Visibility
-import li.cil.oc.api.prefab
+import li.cil.oc.api.driver.DeviceInfo.{DeviceAttribute, DeviceClass}
+import li.cil.oc.api.machine.{Arguments, Callback, Context}
+import li.cil.oc.api.network.{Node, Visibility}
 import li.cil.oc.api.prefab.AbstractManagedEnvironment
+import li.cil.oc.api.{Network, internal}
 import li.cil.oc.common.EventHandler
 import li.cil.oc.util.BlockPosition
 import li.cil.oc.util.ExtendedArguments._
 import li.cil.oc.util.ExtendedNBT._
-import net.minecraft.entity.Entity
-import net.minecraft.entity.EntityLiving
-import net.minecraft.nbt.NBTTagCompound
-import net.minecraft.nbt.NBTTagString
+import li.cil.oc.{Constants, OpenComputers}
+import net.minecraft.entity.{Entity, EntityLiving}
+import net.minecraft.init.Items
+import net.minecraft.item.ItemStack
+import net.minecraft.nbt.{NBTTagCompound, NBTTagString}
 import net.minecraftforge.common.util.Constants.NBT
 
+import java.util
+import java.util.UUID
 import scala.collection.convert.WrapAsJava._
 import scala.collection.convert.WrapAsScala._
 import scala.collection.mutable
+import scala.util.control.Breaks._
 
-class UpgradeLeash(val host: Entity) extends AbstractManagedEnvironment with traits.WorldAware with DeviceInfo {
+class UpgradeLeash(val host: Entity with internal.Drone) extends AbstractManagedEnvironment with traits.WorldAware with DeviceInfo {
   override val node = Network.newNode(this, Visibility.Network).
     withComponent("leash").
     create()
@@ -49,7 +43,7 @@ class UpgradeLeash(val host: Entity) extends AbstractManagedEnvironment with tra
 
   val leashedEntities = mutable.Set.empty[UUID]
 
-  override def position = BlockPosition(host)
+  override def position = BlockPosition(host.getPosition)
 
   @Callback(doc = """function(side:number):boolean -- Tries to put an entity on the specified side of the device onto a leash.""")
   def leash(context: Context, args: Arguments): Array[AnyRef] = {
@@ -58,14 +52,22 @@ class UpgradeLeash(val host: Entity) extends AbstractManagedEnvironment with tra
     val nearBounds = position.bounds
     val farBounds = nearBounds.offset(side.getXOffset * 2.0, side.getYOffset * 2.0, side.getZOffset * 2.0)
     val bounds = nearBounds.union(farBounds)
-    entitiesInBounds[EntityLiving](classOf[EntityLiving], bounds).find(_.canBeLeashedTo(fakePlayer)) match {
-      case Some(entity) =>
-        entity.setLeashHolder(host, true)
-        leashedEntities += entity.getUniqueID
-        context.pause(0.1)
-        result(true)
-      case _ => result(Unit, "no unleashed entity")
+
+    for (index <- 0 to host.mainInventory().getSizeInventory) {
+      val stack = host.mainInventory().getStackInSlot(index)
+      if (stack.getItem == Items.LEAD) {
+        entitiesInBounds[EntityLiving](classOf[EntityLiving], bounds).find(_.canBeLeashedTo(fakePlayer)) match {
+          case Some(entity) =>
+            entity.setLeashHolder(host, true)
+            leashedEntities += entity.getUniqueID
+            stack.shrink(1)
+            context.pause(0.1)
+            result(true)
+          case _ => result(Unit, "no unleashed entity")
+        }
+      }
     }
+    result(Unit, "don't have any lead")
   }
 
   @Callback(doc = """function() -- Unleashes all currently leashed entities.""")
@@ -85,6 +87,20 @@ class UpgradeLeash(val host: Entity) extends AbstractManagedEnvironment with tra
     entitiesInBounds(classOf[EntityLiving], position.bounds.grow(5, 5, 5)).foreach(entity => {
       if (leashedEntities.contains(entity.getUniqueID) && entity.getLeashHolder == host) {
         entity.clearLeashed(true, false)
+        breakable {
+          for (index <- 0 to host.mainInventory().getSizeInventory) {
+            val itemStack = host.mainInventory().getStackInSlot(index)
+            if (itemStack == ItemStack.EMPTY) {
+              host.mainInventory().setInventorySlotContents(index, new ItemStack(Items.LEAD))
+              break()
+            }
+
+            if (itemStack.getItem == Items.LEAD) {
+              itemStack.grow(1)
+              break()
+            }
+          }
+        }
       }
     })
     leashedEntities.clear()