From c8bef0bf5e02a0b1fc9e07a81c0b047b28f8408a Mon Sep 17 00:00:00 2001
From: MCUdude <hansibull@gmail.com>
Date: Thu, 19 Jan 2023 23:08:00 +0100
Subject: [PATCH] Improvements to HardwareSerial.cpp

---
 HardwareSerial.cpp | 34 ++++++++++++++++++++++------------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/HardwareSerial.cpp b/HardwareSerial.cpp
index 048ea69..46d3602 100755
--- a/HardwareSerial.cpp
+++ b/HardwareSerial.cpp
@@ -26,6 +26,7 @@
 #include <stdio.h>
 #include <string.h>
 #include <inttypes.h>
+#include <util/atomic.h>
 #include "Arduino.h"
 
 #include "HardwareSerial.h"
@@ -76,6 +77,13 @@ void serialEventRun(void)
 #endif
 }
 
+// macro to guard critical sections when needed for large TX buffer sizes
+ #if (SERIAL_TX_BUFFER_SIZE>256)
+ #define TX_BUFFER_ATOMIC ATOMIC_BLOCK(ATOMIC_RESTORESTATE)
+ #else
+ #define TX_BUFFER_ATOMIC
+ #endif
+
 // Actual interrupt handlers //////////////////////////////////////////////////////////////
 
 void HardwareSerial::_tx_udr_empty_irq(void)
@@ -174,15 +182,13 @@ int HardwareSerial::read(void)
 
 int HardwareSerial::availableForWrite(void)
 {
-#if (SERIAL_TX_BUFFER_SIZE>256)
-  uint8_t oldSREG = SREG;
-  cli();
-#endif
-  tx_buffer_index_t head = _tx_buffer_head;
-  tx_buffer_index_t tail = _tx_buffer_tail;
-#if (SERIAL_TX_BUFFER_SIZE>256)
-  SREG = oldSREG;
-#endif
+  tx_buffer_index_t head;
+  tx_buffer_index_t tail;
+
+  TX_BUFFER_ATOMIC {
+    head = _tx_buffer_head;
+    tail = _tx_buffer_tail;
+  }
   if (head >= tail) return SERIAL_TX_BUFFER_SIZE - 1 - head + tail;
   return tail - head - 1;
 }
@@ -237,9 +243,13 @@ size_t HardwareSerial::write(uint8_t c)
   }
 
   _tx_buffer[_tx_buffer_head] = c;
-  _tx_buffer_head = i;
-
-  *_ucsrb |= _BV(UDRIE0);
+  // make atomic to prevent execution of ISR between setting the
+  // head pointer and setting the interrupt flag resulting in buffer
+  // retransmission
+  ATOMIC_BLOCK(ATOMIC_RESTORESTATE) {
+    _tx_buffer_head = i;
+    *_ucsrb |= _BV(UDRIE0);
+  }
   
   return 1;
 }