From c51b445e410b6ca2942cb074d26ccb5786e1cb31 Mon Sep 17 00:00:00 2001
From: Asbelos <asbelos@btinternet.com>
Date: Thu, 11 Mar 2021 13:35:47 +0000
Subject: [PATCH] unjoin automatically

---
 DCC.cpp        | 17 ++++++-----------
 WiThrottle.cpp | 49 ++++++++++++++++++++++++++-----------------------
 2 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/DCC.cpp b/DCC.cpp
index 6fe07c1..854f9b2 100644
--- a/DCC.cpp
+++ b/DCC.cpp
@@ -659,6 +659,7 @@ void DCC::ackManagerLoop() {
     // (typically waiting for a reset counter or ACK waiting, or when all finished.)
     switch (opcode) {
       case BASELINE:
+      setProgTrackSyncMain(false);
 	  if (DCCWaveform::progTrack.getPowerMode() == POWERMODE::OFF) {
         if (Diag::ACK) DIAG(F("\nAuto Prog power on"));
         DCCWaveform::progTrack.setPowerMode(POWERMODE::ON);
@@ -725,23 +726,20 @@ void DCC::ackManagerLoop() {
      case ITC0:
      case ITC1:   // If True Callback(0 or 1)  (if prevous WACK got an ACK)
         if (ackReceived) {
-            ackManagerProg = NULL; // all done now
-	          callback(opcode==ITC0?0:1);
+            callback(opcode==ITC0?0:1);
             return;
           }
         break;
         
       case ITCB:   // If True callback(byte)
           if (ackReceived) {
-            ackManagerProg = NULL; // all done now
             callback(ackManagerByte);
             return;
           }
         break;
 
-      case ITCB7:   // If True callback(byte & 0xF)
+      case ITCB7:   // If True callback(byte & 0x7F)
           if (ackReceived) {
-            ackManagerProg = NULL; // all done now
             callback(ackManagerByte & 0x7F);
             return;
           }
@@ -749,15 +747,13 @@ void DCC::ackManagerLoop() {
         
       case NAKFAIL:   // If nack callback(-1)
           if (!ackReceived) {
-            ackManagerProg = NULL; // all done now
-	          callback(-1);
+            callback(-1);
             return;
           }
         break;
         
       case FAIL:  // callback(-1)
-           ackManagerProg = NULL;
-	         callback(-1);
+           callback(-1);
            return;
            
       case STARTMERGE:
@@ -801,7 +797,6 @@ void DCC::ackManagerLoop() {
           
      case COMBINELOCOID: 
           // ackManagerStash is  cv17, ackManagerByte is CV 18
-          ackManagerProg=NULL;
           callback( ackManagerByte + ((ackManagerStash - 192) << 8));
           return;            
 
@@ -817,7 +812,6 @@ void DCC::ackManagerLoop() {
           break;     
      default: 
           DIAG(F("\n!! ackOp %d FAULT!!"),opcode);
-          ackManagerProg=NULL;
           callback( -1);
           return;        
     
@@ -826,6 +820,7 @@ void DCC::ackManagerLoop() {
   }
 }
 void DCC::callback(int value) {
+    ackManagerProg=NULL;  // no more steps to execute
     if (DCCWaveform::progTrack.autoPowerOff) {
       if (Diag::ACK) DIAG(F("\nAuto Prog power off"));
       DCCWaveform::progTrack.doAutoPowerOff();
diff --git a/WiThrottle.cpp b/WiThrottle.cpp
index c08b2b7..72814f8 100644
--- a/WiThrottle.cpp
+++ b/WiThrottle.cpp
@@ -208,24 +208,6 @@ int WiThrottle::getLocoId(byte * cmd) {
     return getInt(cmd+1); 
 }
 
-RingStream * WiThrottle::stashStream;
-WiThrottle * WiThrottle::stashInstance;
-byte         WiThrottle::stashClient;
-char         WiThrottle::stashThrottleChar;
-
-void WiThrottle::getLocoCallback(int locoid) {
-  stashStream->mark(stashClient);
-  if (locoid<0) StringFormatter::send(stashStream,F("HMNo loco found on prog track\n"));
-  else {
-    char addcmd[20]={'M',stashThrottleChar,'+',(locoid>127)?'L':'S'};
-    itoa(locoid,addcmd+4,10);
-    stashInstance->multithrottle(stashStream, (byte *)addcmd);
-    DCCWaveform::progTrack.setPowerMode(POWERMODE::ON);
-    DCC::setProgTrackSyncMain(true);  // <1 JOIN> so we can drive loco away
-  }
-  stashStream->commit();
-}
-
 void WiThrottle::multithrottle(RingStream * stream, byte * cmd){ 
           char throttleChar=cmd[1];
           int locoid=getLocoId(cmd+3); // -1 for *
@@ -237,14 +219,15 @@ void WiThrottle::multithrottle(RingStream * stream, byte * cmd){
        switch(cmd[2]) {
           case '+':  // add loco request
                 if (cmd[3]=='*') { 
-                  // M+* means get loco from prog track
+                  // M+* means get loco from prog track, then join tracks ready to drive away
+                  // Stash the things the callback will need later
                   stashStream= stream;
                   stashClient=stream->peekTargetMark();
                   stashThrottleChar=throttleChar;
                   stashInstance=this;
-                  DCC::setProgTrackSyncMain(false);  // remove JOIN so we can read prog
-                  DCC::getLocoId(getLocoCallback);                    
-                  return; // return nothing in stream as response comes later 
+                  // ask DCC to call us back when the loco id has been read
+                  DCC::getLocoId(getLocoCallback); // will remove any previous join                    
+                  return; // return nothing in stream as response is sent later in the callback 
                 }
                 //return error if address zero requested
                 if (locoid==0) { 
@@ -396,4 +379,24 @@ void WiThrottle::checkHeartbeat() {
 
 char WiThrottle::LorS(int cab) {
     return (cab<127)?'S':'L';
-} 
+}
+
+// Drive Away feature. Callback handling
+ 
+RingStream * WiThrottle::stashStream;
+WiThrottle * WiThrottle::stashInstance;
+byte         WiThrottle::stashClient;
+char         WiThrottle::stashThrottleChar;
+
+void WiThrottle::getLocoCallback(int locoid) {
+  stashStream->mark(stashClient);
+  if (locoid<0) StringFormatter::send(stashStream,F("HMNo loco found on prog track\n"));
+  else {
+    char addcmd[20]={'M',stashThrottleChar,'+',LorS(locoid) };
+    itoa(locoid,addcmd+4,10);
+    stashInstance->multithrottle(stashStream, (byte *)addcmd);
+    DCCWaveform::progTrack.setPowerMode(POWERMODE::ON);
+    DCC::setProgTrackSyncMain(true);  // <1 JOIN> so we can drive loco away
+  }
+  stashStream->commit();
+}