From 76c56081815c6fb63a2d1e0af37af4aa01359b95 Mon Sep 17 00:00:00 2001
From: Harald Barth <haba@kth.se>
Date: Tue, 9 Aug 2022 13:12:04 +0200
Subject: [PATCH] Protect port registers from change during interrupt code in
 differnet way

---
 MotorDriver.cpp  | 54 +++++++++++++++++++++++-------------------------
 MotorDriver.h    | 11 ++++++++--
 TrackManager.cpp |  4 ++--
 3 files changed, 37 insertions(+), 32 deletions(-)

diff --git a/MotorDriver.cpp b/MotorDriver.cpp
index 620d996..be5e304 100644
--- a/MotorDriver.cpp
+++ b/MotorDriver.cpp
@@ -148,11 +148,17 @@ bool MotorDriver::isPWMCapable() {
 void MotorDriver::setPower(POWERMODE mode) {
   bool on=mode==POWERMODE::ON;
   if (on) {
+    noInterrupts();
     IODevice::write(powerPin,HIGH);
+    interrupts();
     if (isProgTrack)
       DCCWaveform::progTrack.clearResets();
   }
-  else IODevice::write(powerPin,LOW);
+  else {
+      noInterrupts();
+      IODevice::write(powerPin,LOW);
+      interrupts();
+  }
   powerMode=mode; 
 }
 
@@ -164,10 +170,14 @@ void MotorDriver::setPower(POWERMODE mode) {
 // (HIGH == release brake) and setBrake does
 // compensate for that.
 //
-void MotorDriver::setBrake(bool on) {
+void MotorDriver::setBrake(bool on, bool interruptContext) {
   if (brakePin == UNUSED_PIN) return;
-  if (on ^ invertBrake) setHIGH(fastBrakePin);
-  else setLOW(fastBrakePin);
+  if (!interruptContext) {noInterrupts();}
+  if (on ^ invertBrake)
+    setHIGH(fastBrakePin);
+  else
+    setLOW(fastBrakePin);
+  if (!interruptContext) {interrupts();}
 }
 
 bool MotorDriver::canMeasureCurrent() {
@@ -214,30 +224,18 @@ void MotorDriver::setDCSignal(byte speedcode) {
   if (invertBrake)
     brake=255-brake;
   analogWrite(brakePin,brake);
-  // as the port registers can be shadowed to get syncronized DCC signals
-  // we need to take care of that and we have to turn off interrupts during
-  // that time as otherwise setDCCSignal() which is called from interrupt
-  // contect can undo whatever we do here.
-  if (fastSignalPin.shadowinout != NULL) {
-    if (HAVE_PORTA(fastSignalPin.shadowinout == &PORTA)) {
-      noInterrupts();
-      HAVE_PORTA(shadowPORTA=PORTA);
-      setSignal(tDir);
-      HAVE_PORTA(PORTA=shadowPORTA);
-      interrupts();
-    } else if (HAVE_PORTB(fastSignalPin.shadowinout == &PORTB)) {
-      noInterrupts();
-      HAVE_PORTB(shadowPORTB=PORTB);
-      setSignal(tDir);
-      HAVE_PORTB(PORTB=shadowPORTB);
-      interrupts();
-    } else if (HAVE_PORTC(fastSignalPin.shadowinout == &PORTC)) {
-      noInterrupts();
-      HAVE_PORTC(shadowPORTC=PORTC);
-      setSignal(tDir);
-      HAVE_PORTC(PORTC=shadowPORTC);
-      interrupts();
-    }
+  if (HAVE_PORTA(fastSignalPin.shadowinout == &PORTA)) {
+    HAVE_PORTA(shadowPORTA=PORTA);
+    setSignal(tDir);
+    HAVE_PORTA(PORTA=shadowPORTA);
+  } else if (HAVE_PORTB(fastSignalPin.shadowinout == &PORTB)) {
+    HAVE_PORTB(shadowPORTB=PORTB);
+    setSignal(tDir);
+    HAVE_PORTB(PORTB=shadowPORTB);
+  } else if (HAVE_PORTC(fastSignalPin.shadowinout == &PORTC)) {
+    HAVE_PORTC(shadowPORTC=PORTC);
+    setSignal(tDir);
+    HAVE_PORTC(PORTC=shadowPORTC);
   } else {
     setSignal(tDir);
   }
diff --git a/MotorDriver.h b/MotorDriver.h
index 9e2ff4f..2138a40 100644
--- a/MotorDriver.h
+++ b/MotorDriver.h
@@ -105,7 +105,14 @@ class MotorDriver {
                 byte current_pin, float senseFactor, unsigned int tripMilliamps, byte faultPin);
     virtual void setPower( POWERMODE mode);
     virtual POWERMODE getPower() { return powerMode;}
-    __attribute__((always_inline)) inline void setSignal( bool high) {
+    // as the port registers can be shadowed to get syncronized DCC signals
+    // we need to take care of that and we have to turn off interrupts if
+    // we setSignal() or setBrake() or setPower() during that time as
+    // otherwise the call from interrupt context can undo whatever we do
+    // from outside interrupt
+    virtual void setBrake( bool on, bool interruptContext=false);
+  __attribute__((always_inline)) inline void setSignal( bool high, bool interruptContext=false) {
+      if (!interruptContext) {noInterrupts();}
       if (trackPWM) {
 	DCCTimer::setPWM(signalPin,high);
       }
@@ -119,6 +126,7 @@ class MotorDriver {
 	  if (dualSignal) setHIGH(fastSignalPin2);
 	}
       }
+      if (!interruptContext) {interrupts();}
     };
     inline void enableSignal(bool on) {
       if (on)
@@ -127,7 +135,6 @@ class MotorDriver {
 	pinMode(signalPin, INPUT);
     };
     inline byte getSignalPin() { return signalPin; };
-    virtual void setBrake( bool on);
     virtual void setDCSignal(byte speedByte);
     virtual int  getCurrentRaw();
     virtual int getCurrentRawInInterrupt();
diff --git a/TrackManager.cpp b/TrackManager.cpp
index ba3f8c0..6c4da7b 100644
--- a/TrackManager.cpp
+++ b/TrackManager.cpp
@@ -92,7 +92,7 @@ void TrackManager::setDCCSignal( bool on) {
   HAVE_PORTA(shadowPORTA=PORTA);
   HAVE_PORTB(shadowPORTB=PORTB);
   HAVE_PORTC(shadowPORTC=PORTC);
-  APPLY_BY_MODE(TRACK_MODE_MAIN,setSignal(on));
+  APPLY_BY_MODE(TRACK_MODE_MAIN,setSignal(on, true));
   HAVE_PORTA(PORTA=shadowPORTA);
   HAVE_PORTB(PORTB=shadowPORTB);
   HAVE_PORTC(PORTC=shadowPORTC);
@@ -110,7 +110,7 @@ void TrackManager::setPROGSignal( bool on) {
   HAVE_PORTA(shadowPORTA=PORTA);
   HAVE_PORTB(shadowPORTB=PORTB);
   HAVE_PORTC(shadowPORTC=PORTC);
-  APPLY_BY_MODE(TRACK_MODE_PROG,setSignal(on));
+  APPLY_BY_MODE(TRACK_MODE_PROG,setSignal(on, true));
   HAVE_PORTA(PORTA=shadowPORTA);
   HAVE_PORTB(PORTB=shadowPORTB);
   HAVE_PORTC(PORTC=shadowPORTC);