Refactor node persistence

src/main/scala/ocelot/desktop/node/Node.scala

@@ -1,5 +1,6 @@
 package ocelot.desktop.node
 
+import ocelot.desktop.OcelotDesktop
 import ocelot.desktop.color.{Color, RGBAColor}
 import ocelot.desktop.geometry.{Rect2D, Size2D, Vector2D}
 import ocelot.desktop.graphics.Graphics
@@ -12,13 +13,17 @@ import ocelot.desktop.ui.widget.{InputDialog, Widget, WorkspaceView}
 import ocelot.desktop.util.DrawUtils
 import ocelot.desktop.util.animation.ColorAnimation
 import org.lwjgl.input.Keyboard
-import totoro.ocelot.brain.entity.traits.Environment
+import totoro.ocelot.brain.entity.traits.{Entity, Environment, SidedEnvironment}
 import totoro.ocelot.brain.nbt.NBTTagCompound
 import totoro.ocelot.brain.network
+import totoro.ocelot.brain.util.Direction
 
 import scala.collection.mutable.ArrayBuffer
 
-trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
+abstract class Node(val entity: Entity with Environment) extends Widget with DragHandler with ClickHandler with HoverHandler {
+  if (!OcelotDesktop.workspace.getEntitiesIter.contains(entity))
+    OcelotDesktop.workspace.add(entity)
+
   var workspaceView: WorkspaceView = _
 
   protected val MovingHighlight: RGBAColor = RGBAColor(240, 250, 240)
@@ -47,6 +52,15 @@ trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
   }
 
   def save(nbt: NBTTagCompound): Unit = {
+    val address = entity match {
+      case env: SidedEnvironment =>
+        env.sidedNode(Direction.South).address
+      case _ => if (entity.node == null) null else entity.node.address
+    }
+
+    if (address != null && address.nonEmpty)
+      nbt.setString("address", address)
+
     val posTag = new NBTTagCompound
     position.save(posTag)
     nbt.setTag("pos", posTag)
@@ -106,9 +120,9 @@ trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
       ))
     }
 
-    if (exposeAddress && environment.node != null && environment.node.address != null) {
+    if (exposeAddress && entity.node != null && entity.node.address != null) {
       menu.addEntry(new ContextMenuEntry("Copy address", () => {
-        UiHandler.clipboard = environment.node.address
+        UiHandler.clipboard = entity.node.address
       }))
     }
 
@@ -133,15 +147,13 @@ trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
     }
   }
 
-  def environment: Environment
-
   def icon: String = "icons/NA"
 
   def iconColor: Color = RGBAColor(255, 255, 255)
 
   def ports: Array[NodePort] = Array(NodePort())
 
-  def getNodeByPort(port: NodePort): network.Node = environment.node
+  def getNodeByPort(port: NodePort): network.Node = entity.node
 
   def connections: Iterator[(NodePort, Node, NodePort)] = _connections.iterator
 
@@ -259,7 +271,7 @@ trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
   }
 
   def labelOrAddress: String = {
-    label.orElse(Option(environment.node.address)).getOrElse("unknown")
+    label.orElse(Option(entity.node.address)).getOrElse("unknown")
   }
 
   override def draw(g: Graphics): Unit = {
@@ -291,5 +303,6 @@ trait Node extends Widget with DragHandler with ClickHandler with HoverHandler {
   def dispose(): Unit = {
     disconnectFromAll()
     window.foreach(_.hide())
+    OcelotDesktop.workspace.remove(entity)
   }
 }

src/main/scala/ocelot/desktop/node/NodeRegistry.scala

@@ -14,7 +14,7 @@ object NodeRegistry {
 
   for (tier <- 0 to 2) {
     register(NodeType("Screen" + tier, "nodes/Screen", tier, () => {
-      new ScreenNode(new Screen(tier))
+      new ScreenNode(new Screen(tier)).setup()
     }))
   }
 
@@ -24,7 +24,7 @@ object NodeRegistry {
 
   for (tier <- 0 to 3) {
     register(NodeType("Computer" + tier, "nodes/Computer", tier, () => {
-      new ComputerNode(new Case(tier))
+      new ComputerNode(new Case(tier)).setup()
     }))
   }
 

src/main/scala/ocelot/desktop/node/nodes/CableNode.scala

@@ -1,34 +1,10 @@
 package ocelot.desktop.node.nodes
 
-import ocelot.desktop.OcelotDesktop
 import ocelot.desktop.geometry.Size2D
 import ocelot.desktop.node.Node
 import totoro.ocelot.brain.entity.Cable
-import totoro.ocelot.brain.entity.traits.Environment
-import totoro.ocelot.brain.nbt.NBTTagCompound
-
-class CableNode(cable: Cable) extends Node {
-  OcelotDesktop.workspace.add(cable)
-
-  def this(nbt: NBTTagCompound) {
-    this({
-      val address = nbt.getString("address")
-      if (address == "")
-        new Cable
-      else
-        OcelotDesktop.workspace.entityByAddress(address).get.asInstanceOf[Cable]
-    })
-    super.load(nbt)
-  }
-
-  override def save(nbt: NBTTagCompound): Unit = {
-    super.save(nbt)
-    if (cable.node.address != null)
-      nbt.setString("address", cable.node.address)
-  }
-
-  override def environment: Environment = cable
 
+class CableNode(val cable: Cable) extends Node(cable) {
   override def icon: String = "nodes/Cable"
 
   override def minimumSize: Size2D = Size2D(36, 36)

src/main/scala/ocelot/desktop/node/nodes/ComputerNode.scala

@@ -1,6 +1,5 @@
 package ocelot.desktop.node.nodes
 
-import ocelot.desktop.OcelotDesktop
 import ocelot.desktop.audio.{Audio, SoundCategory, SoundSource}
 import ocelot.desktop.color.Color
 import ocelot.desktop.graphics.Graphics
@@ -11,15 +10,14 @@ import ocelot.desktop.ui.widget.contextmenu.{ContextMenu, ContextMenuEntry, Cont
 import ocelot.desktop.ui.widget.slot._
 import ocelot.desktop.util.{ResourceManager, TierColor}
 import org.lwjgl.input.Keyboard
-import totoro.ocelot.brain.entity.traits.{Computer, Entity, Floppy, GenericCPU, Inventory}
+import totoro.ocelot.brain.entity.traits.{Entity, Floppy, GenericCPU, Inventory}
 import totoro.ocelot.brain.entity.{CPU, Case, EEPROM, GraphicsCard, HDDManaged, HDDUnmanaged, Memory}
 import totoro.ocelot.brain.loot.Loot
-import totoro.ocelot.brain.nbt.NBTTagCompound
 import totoro.ocelot.brain.util.Tier
 
 import scala.reflect.ClassTag
 
-class ComputerNode(val computer: Case, setup: Boolean = true) extends Node {
+class ComputerNode(val computer: Case) extends Node(computer) {
   var eepromSlot: EEPROMSlot = _
   var cpuSlot: CPUSlot = _
   var memorySlots: Array[MemorySlot] = _
@@ -30,34 +28,19 @@ class ComputerNode(val computer: Case, setup: Boolean = true) extends Node {
   private val soundComputerRunning = new SoundSource(Audio.MachineComputerRunning, SoundCategory.Environment, true)
 
   setupSlots()
+  refitSlots()
 
-  if (setup) {
+  def setup(): ComputerNode = {
     cpuSlot.owner.put(new CPU(computer.tier.min(2)))
     memorySlots(0).owner.put(new Memory(computer.tier.min(2) * 2 + 1))
     memorySlots(1).owner.put(new Memory(computer.tier.min(2) * 2 + 1))
     cardSlots(0).owner.put(new GraphicsCard(computer.tier.min(1)))
     floppySlot.map(_.owner).foreach(_.put(Loot.OpenOsFloppy.create()))
     eepromSlot.owner.put(Loot.LuaBiosEEPROM.create())
-    OcelotDesktop.workspace.add(computer)
+    refitSlots()
+    this
   }
 
-  refitSlots()
-
-  def this(nbt: NBTTagCompound) {
-    this({
-      val address = nbt.getString("address")
-      OcelotDesktop.workspace.entityByAddress(address).get.asInstanceOf[Case]
-    }, setup = false)
-    super.load(nbt)
-  }
-
-  override def save(nbt: NBTTagCompound): Unit = {
-    super.save(nbt)
-    nbt.setString("address", computer.node.address)
-  }
-
-  override def environment: Computer = computer
-
   override val icon: String = "nodes/Computer"
   override def iconColor: Color = TierColor.get(computer.tier)
 

src/main/scala/ocelot/desktop/node/nodes/DiskDriveNode.scala

@@ -1,37 +1,18 @@
 package ocelot.desktop.node.nodes
 
-import ocelot.desktop.OcelotDesktop
 import ocelot.desktop.color.IntColor
 import ocelot.desktop.graphics.Graphics
 import ocelot.desktop.node.Node
 import ocelot.desktop.ui.widget.slot.FloppySlot
 import totoro.ocelot.brain.entity.FloppyDiskDrive
-import totoro.ocelot.brain.entity.traits.{Environment, Floppy}
+import totoro.ocelot.brain.entity.traits.Floppy
 import totoro.ocelot.brain.loot.Loot
-import totoro.ocelot.brain.nbt.NBTTagCompound
 import totoro.ocelot.brain.util.DyeColor
 
-class DiskDriveNode(val diskDrive: FloppyDiskDrive) extends Node {
-  OcelotDesktop.workspace.add(diskDrive)
-
+class DiskDriveNode(val diskDrive: FloppyDiskDrive) extends Node(diskDrive) {
   val slot: FloppySlot = new FloppySlot(diskDrive.inventory(0))
   slot.item = floppy.getOrElse(Loot.OpenOsFloppy.create())
 
-  def this(nbt: NBTTagCompound) {
-    this({
-      val address = nbt.getString("address")
-      OcelotDesktop.workspace.entityByAddress(address).get.asInstanceOf[FloppyDiskDrive]
-    })
-    super.load(nbt)
-  }
-
-  override def save(nbt: NBTTagCompound): Unit = {
-    super.save(nbt)
-    nbt.setString("address", diskDrive.node.address)
-  }
-
-  override def environment: Environment = diskDrive
-
   override def icon: String = "nodes/DiskDrive"
 
   override protected val canOpen = true

src/main/scala/ocelot/desktop/node/nodes/RelayNode.scala

@@ -1,30 +1,11 @@
 package ocelot.desktop.node.nodes
 
-import ocelot.desktop.OcelotDesktop
 import ocelot.desktop.node.{Node, NodePort}
 import totoro.ocelot.brain.entity.Relay
-import totoro.ocelot.brain.nbt.NBTTagCompound
 import totoro.ocelot.brain.network
 import totoro.ocelot.brain.util.Direction
 
-class RelayNode(relay: Relay) extends Node {
-  OcelotDesktop.workspace.add(relay)
-
-  def this(nbt: NBTTagCompound) {
-    this({
-      val address = nbt.getString("address")
-      OcelotDesktop.workspace.entityByAddress(address).get.asInstanceOf[Relay]
-    })
-    super.load(nbt)
-  }
-
-  override def save(nbt: NBTTagCompound): Unit = {
-    super.save(nbt)
-    nbt.setString("address", relay.sidedNode(Direction.South).address)
-  }
-
-  override def environment: Relay = relay
-
+class RelayNode(val relay: Relay) extends Node(relay) {
   override val icon: String = "nodes/Relay"
 
   override def ports: Array[NodePort] = Array(

src/main/scala/ocelot/desktop/node/nodes/ScreenNode.scala

@@ -6,32 +6,40 @@ import ocelot.desktop.graphics.Graphics
 import ocelot.desktop.node.Node
 import ocelot.desktop.ui.widget.contextmenu.{ContextMenu, ContextMenuEntry}
 import ocelot.desktop.util.TierColor
-import totoro.ocelot.brain.entity.traits.Environment
 import totoro.ocelot.brain.entity.{Keyboard, Screen}
 import totoro.ocelot.brain.nbt.NBTTagCompound
 
-class ScreenNode(var screen: Screen, setup: Boolean = true) extends Node {
-  if (setup) {
-    val keyboard = new Keyboard
-    OcelotDesktop.workspace.add(screen)
-    OcelotDesktop.workspace.add(keyboard)
-    screen.connect(keyboard)
-  }
+class ScreenNode(val screen: Screen) extends Node(screen) {
+  private var keyboard: Option[Keyboard] = None
 
-  def this(nbt: NBTTagCompound) {
-    this({
-      val address = nbt.getString("address")
-      OcelotDesktop.workspace.entityByAddress(address).get.asInstanceOf[Screen]
-    }, setup = false)
+  override def load(nbt: NBTTagCompound): Unit = {
     super.load(nbt)
+    if (nbt.hasKey("string")) {
+      keyboard = OcelotDesktop.workspace.entityByAddress(nbt.getString("string")).map(_.asInstanceOf[Keyboard])
+    }
   }
 
   override def save(nbt: NBTTagCompound): Unit = {
     super.save(nbt)
-    nbt.setString("address", screen.node.address)
+    if (keyboard.isDefined) {
+      nbt.setString("keyboard", keyboard.get.node.address)
+    }
   }
 
-  override def environment: Environment = screen
+  override def dispose(): Unit = {
+    super.dispose()
+    if (keyboard.isDefined) {
+      OcelotDesktop.workspace.remove(keyboard.get)
+    }
+  }
+
+  def setup(): ScreenNode = {
+    val kbd = new Keyboard
+    OcelotDesktop.workspace.add(kbd)
+    screen.connect(kbd)
+    keyboard = Some(kbd)
+    this
+  }
 
   override def icon: String = "nodes/Screen"
 

src/main/scala/ocelot/desktop/ui/widget/WorkspaceView.scala

@@ -11,9 +11,9 @@ import ocelot.desktop.ui.layout.{CopyLayout, Layout}
 import ocelot.desktop.ui.widget.window.{NodeSelector, ProfilerWindow, WindowPool}
 import ocelot.desktop.util.DrawUtils
 import ocelot.desktop.util.animation.ValueAnimation
-import ocelot.desktop.{ColorScheme, Settings}
+import ocelot.desktop.{ColorScheme, OcelotDesktop, Settings}
 import org.lwjgl.input.Keyboard
-import totoro.ocelot.brain.entity.traits.{Environment, SidedEnvironment}
+import totoro.ocelot.brain.entity.traits.{Entity, Environment, SidedEnvironment}
 import totoro.ocelot.brain.entity.{Case, Screen}
 import totoro.ocelot.brain.nbt.ExtendedNBT._
 import totoro.ocelot.brain.nbt.{NBT, NBTBase, NBTTagCompound}
@@ -58,9 +58,18 @@ class WorkspaceView extends Widget with DragHandler with ClickHandler with Hover
     nbt.getTagList("nodes", NBT.TAG_COMPOUND).foreach((nbt: NBTTagCompound) => {
       val clazzName = nbt.getString("class")
       val clazz = Class.forName(clazzName)
-      val node = clazz.getConstructor(classOf[NBTTagCompound]).newInstance(nbt).asInstanceOf[Node]
+
+      val address = nbt.getString("address")
+      var entity = OcelotDesktop.workspace.entityByAddress(address).orNull
+      if (entity == null) {
+        val entityClass = Class.forName(nbt.getString("entityClass"))
+        entity = entityClass.getConstructor().newInstance().asInstanceOf[Entity]
+      }
+
+      val node = clazz.getConstructor(entity.getClass).newInstance(entity).asInstanceOf[Node]
       node.workspaceView = this
       node.root = _root
+      node.load(nbt)
       nodes += node
     })
 
@@ -78,7 +87,7 @@ class WorkspaceView extends Widget with DragHandler with ClickHandler with Hover
   }
 
   private def findNodeByAddress(nodes: Iterable[Node], address: String): Option[Node] =
-    nodes.find(_.environment match {
+    nodes.find(_.entity match {
       case sided: SidedEnvironment => Direction.values.exists(dir => sided.canConnect(dir) && sided.sidedNode(dir).address == address)
       case env => env.node.address == address
     })
@@ -91,6 +100,7 @@ class WorkspaceView extends Widget with DragHandler with ClickHandler with Hover
     nbt.setTagList("nodes", nodes.map(node => {
       val nbt = new NBTTagCompound
       nbt.setString("class", node.getClass.getName)
+      nbt.setString("entityClass", node.entity.getClass.getName)
       node.save(nbt)
       nbt: NBTBase
     }).asJava)
@@ -100,9 +110,9 @@ class WorkspaceView extends Widget with DragHandler with ClickHandler with Hover
     for (nodeA <- nodes) {
       for ((portA, nodeB, portB) <- nodeA.connections) {
         val byteA = portA.toByte
-        val addrA = getNode(nodeA.environment, portA.direction).address
+        val addrA = getNode(nodeA.entity, portA.direction).address
         val byteB = portB.toByte
-        val addrB = getNode(nodeB.environment, portB.direction).address
+        val addrB = getNode(nodeB.entity, portB.direction).address
         val pair = (byteA, addrA, byteB, addrB)
         val pairRev = (byteB, addrB, byteA, addrA)
         if (!pairs.contains(pair) && !pairs.contains(pairRev)) pairs += pair
@@ -153,8 +163,8 @@ class WorkspaceView extends Widget with DragHandler with ClickHandler with Hover
   }
 
   def createDefaultWorkspace(): Unit = {
-    addNode(new ComputerNode(new Case(Tier.Three)), Vector2D(68, 68))
-    addNode(new ScreenNode(new Screen(Tier.Two)), Vector2D(204, 136))
+    addNode(new ComputerNode(new Case(Tier.Three)).setup(), Vector2D(68, 68))
+    addNode(new ScreenNode(new Screen(Tier.Two)).setup(), Vector2D(204, 136))
     nodes(0).connect(NodePort(), nodes(1), NodePort())
   }