From e56871f5efa28420904bd8ec013a4b424186bb67 Mon Sep 17 00:00:00 2001
From: "NGnius (Graham)" <ngniusness@gmail.com>
Date: Mon, 3 Aug 2020 13:24:35 -0400
Subject: [PATCH] Fix wire entity queries when its not submitted yet

---
 GamecraftModdingAPI/Blocks/SignalEngine.cs | 29 ++++++++++++++++++++--
 GamecraftModdingAPI/Blocks/Wire.cs         | 24 +++++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/GamecraftModdingAPI/Blocks/SignalEngine.cs b/GamecraftModdingAPI/Blocks/SignalEngine.cs
index e710e24..e961423 100644
--- a/GamecraftModdingAPI/Blocks/SignalEngine.cs
+++ b/GamecraftModdingAPI/Blocks/SignalEngine.cs
@@ -39,7 +39,7 @@ namespace GamecraftModdingAPI.Blocks
 
         // implementations for block wiring
 
-        public EGID CreateNewWire(EGID startBlock, byte startPort, EGID endBlock, byte endPort)
+        public WireEntityStruct CreateNewWire(EGID startBlock, byte startPort, EGID endBlock, byte endPort)
         {
 	        EGID wireEGID = new EGID(WiresExclusiveGroups.NewWireEntityId, NamedExclusiveGroup<WiresGroup>.Group);
 	        EntityComponentInitializer wireInitializer = Factory.BuildEntity<WireEntityDescriptor>(wireEGID);
@@ -49,8 +49,9 @@ namespace GamecraftModdingAPI.Blocks
 				sourcePortUsage = startPort,
 				destinationBlockEGID = endBlock,
 				destinationPortUsage = endPort,
+				ID = wireEGID
 			});
-	        return wireEGID;
+	        return wireInitializer.Get<WireEntityStruct>();
         }
 
         public ref WireEntityStruct GetWire(EGID wire)
@@ -209,11 +210,35 @@ namespace GamecraftModdingAPI.Blocks
 			return new EGID(ports.firstInputID + portUsage, NamedExclusiveGroup<InputPortsGroup>.Group);
 		}
 		
+		public EGID MatchBlockInputToPort(EGID block, byte portUsage, out bool exists)
+		{
+			if (!entitiesDB.Exists<BlockPortsStruct>(block))
+			{
+				exists = false;
+				return default;
+			}
+			exists = true;
+			BlockPortsStruct ports = entitiesDB.QueryEntity<BlockPortsStruct>(block);
+			return new EGID(ports.firstInputID + portUsage, NamedExclusiveGroup<InputPortsGroup>.Group);
+		}
+		
 		public EGID MatchBlockOutputToPort(Block block, byte portUsage, out bool exists)
 		{
 			BlockPortsStruct ports = GetFromDbOrInitData<BlockPortsStruct>(block, block.Id, out exists);
 			return new EGID(ports.firstOutputID + portUsage, NamedExclusiveGroup<OutputPortsGroup>.Group);
 		}
+		
+		public EGID MatchBlockOutputToPort(EGID block, byte portUsage, out bool exists)
+		{
+			if (!entitiesDB.Exists<BlockPortsStruct>(block))
+			{
+				exists = false;
+				return default;
+			}
+			exists = true;
+			BlockPortsStruct ports = entitiesDB.QueryEntity<BlockPortsStruct>(block);
+			return new EGID(ports.firstOutputID + portUsage, NamedExclusiveGroup<OutputPortsGroup>.Group);
+		}
 
 		public ref WireEntityStruct MatchPortToWire(EGID portID, EGID blockID, out bool exists)
 		{
diff --git a/GamecraftModdingAPI/Blocks/Wire.cs b/GamecraftModdingAPI/Blocks/Wire.cs
index 9a73135..6c427f6 100644
--- a/GamecraftModdingAPI/Blocks/Wire.cs
+++ b/GamecraftModdingAPI/Blocks/Wire.cs
@@ -26,8 +26,8 @@ namespace GamecraftModdingAPI.Blocks
         
         public static Wire Connect(SignalingBlock start, byte startPort, SignalingBlock end, byte endPort)
         {
-            EGID wireEgid = signalEngine.CreateNewWire(start.Id, startPort, end.Id, endPort);
-            return new Wire(start, end, startPort, endPort, wireEgid, false);
+            WireEntityStruct wire = signalEngine.CreateNewWire(start.Id, startPort, end.Id, endPort);
+            return new Wire(wire);
         }
 
         /// <summary>
@@ -155,6 +155,22 @@ namespace GamecraftModdingAPI.Blocks
             this.startBlockEGID = wire.sourceBlockEGID;
             this.endBlockEGID = wire.destinationBlockEGID;
             this.inputToOutput = false;
+            endPortEGID = signalEngine.MatchBlockInputToPort(wire.destinationBlockEGID, wire.destinationPortUsage, out bool exists);
+            if (!exists) throw new WireInvalidException("Wire end port not found");
+            startPortEGID = signalEngine.MatchBlockOutputToPort(wire.sourceBlockEGID, wire.sourcePortUsage, out exists);
+            if (!exists) throw new WireInvalidException("Wire start port not found");
+        }
+
+        internal Wire(WireEntityStruct wire)
+        {
+            this.wireEGID = wire.ID;
+            this.startBlockEGID = wire.sourceBlockEGID;
+            this.endBlockEGID = wire.destinationBlockEGID;
+            inputToOutput = false;
+            endPortEGID = signalEngine.MatchBlockInputToPort(wire.destinationBlockEGID, wire.destinationPortUsage, out bool exists);
+            if (!exists) throw new WireInvalidException("Wire end port not found");
+            startPortEGID = signalEngine.MatchBlockOutputToPort(wire.sourceBlockEGID, wire.sourcePortUsage, out exists);
+            if (!exists) throw new WireInvalidException("Wire start port not found");
         }
 
         /// <summary>
@@ -328,9 +344,9 @@ namespace GamecraftModdingAPI.Blocks
         {
             if (signalEngine.Exists<WireEntityStruct>(wireEGID))
             {
-                return $"{nameof(Id)}: {Id}, Start{nameof(Start.Id)}: {Start.Id}, End{nameof(Id)}: {End.Id}, ({Start.Type}::{StartPort} aka {Start.PortName(StartPort, inputToOutput)}) -> ({End.Type}::{EndPort} aka {End.PortName(EndPort, !inputToOutput)})";
+                return $"{nameof(Id)}: {Id}, Start{nameof(Start.Id)}: {Start.Id}, End{nameof(End.Id)}: {End.Id}, ({Start.Type}::{StartPort} aka {Start.PortName(StartPort, inputToOutput)}) -> ({End.Type}::{EndPort} aka {End.PortName(EndPort, !inputToOutput)})";
             }
-            return $"{nameof(Id)}: {Id}, Start{nameof(Start.Id)}: {Start.Id}, End{nameof(Id)}: {End.Id}, ({Start.Type}::{StartPort}) -> ({End.Type}::{EndPort})";
+            return $"{nameof(Id)}: {Id}, Start{nameof(Start.Id)}: {Start.Id}, End{nameof(End.Id)}: {End.Id}, ({Start.Type} -> {End.Type})";
         }
 
         internal static void Init() { }