From 946b701b08876ec66a9d58c1b4ebafe327079913 Mon Sep 17 00:00:00 2001
From: Dan Applegate <dapplegate@squareup.com>
Date: Mon, 17 Oct 2022 09:15:50 -0400
Subject: [PATCH] Fix writing of last word on DWORD TPI parts (#1115)

* Fix writing of last word on DWORD TPI parts

* Add n_word_writes AVRMEM config option

* TPI word chunk mode in avr_write_mem

* Simplify addition of n_words_write mem component to grammar

Co-authored-by: Stefan Rueger <stefan.rueger@urclocks.com>
---
 src/avr.c            | 44 +++++++++++++++++++++++++++++++-------------
 src/avrdude.conf.in  |  3 +++
 src/config.c         |  3 +++
 src/developer_opts.c |  1 +
 src/doc/avrdude.texi |  1 +
 src/lexer.l          |  2 +-
 src/libavrdude.h     |  1 +
 7 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/src/avr.c b/src/avr.c
index 632a95a2..29ab2d7b 100644
--- a/src/avr.c
+++ b/src/avr.c
@@ -856,6 +856,9 @@ int avr_write_mem(const PROGRAMMER *pgm, const AVRPART *p, const AVRMEM *m, int
 
 
   if ((p->prog_modes & PM_TPI) && m->page_size > 1 && pgm->cmd_tpi) {
+    unsigned int    chunk; /* number of words for each write command */
+    unsigned int    j, writeable_chunk;
+
     if (wsize == 1) {
       /* fuse (configuration) memory: only single byte to write */
       return avr_write_byte(pgm, p, m, 0, m->buf[0]) == 0? 1: LIBAVRDUDE_GENERAL_FAILURE;
@@ -866,35 +869,50 @@ int avr_write_mem(const PROGRAMMER *pgm, const AVRPART *p, const AVRMEM *m, int
     /* setup for WORD_WRITE */
     avr_tpi_setup_rw(pgm, m, 0, TPI_NVMCMD_WORD_WRITE);
 
-    /* make sure it's aligned to a word boundary */
-    if (wsize & 0x1) {
-      wsize++;
+    /*
+     * Some TPI devices can only program 2 or 4 words (4 or 8 bytes) at a time.
+     * This is set by the n_word_writes option of the AVRMEM config section.
+     * Ensure that we align our write size to this boundary.
+     */
+    if (m->n_word_writes < 0 || m->n_word_writes > 4 || m->n_word_writes == 3) {
+      avrdude_message(MSG_INFO, "\n%s: ERROR: Unsupported n_word_writes value of %d "
+                      "configured for %s memory\n"
+                      "%sAborting write\n",
+                      progname, m->n_word_writes, m->desc, progbuf);
+      return LIBAVRDUDE_GENERAL_FAILURE;
     }
+    chunk = m->n_word_writes > 0 ? 2*m->n_word_writes : 2;
+    wsize = (wsize+chunk-1) / chunk * chunk;
 
-    /* write words, low byte first */
-    for (lastaddr = i = 0; i < wsize; i += 2) {
-      if ((m->tags[i] & TAG_ALLOCATED) != 0 ||
-          (m->tags[i + 1] & TAG_ALLOCATED) != 0) {
+    /* write words in chunks, low byte first */
+    for (lastaddr = i = 0; i < wsize; i += chunk) {
+      /* check that at least one byte in this chunk is allocated */
+      for (writeable_chunk = j = 0; !writeable_chunk && j < chunk; j++) {
+        writeable_chunk = m->tags[i+j] & TAG_ALLOCATED;
+      }
 
+      if (writeable_chunk) {
         if (lastaddr != i) {
           /* need to setup new address */
           avr_tpi_setup_rw(pgm, m, i, TPI_NVMCMD_WORD_WRITE);
           lastaddr = i;
         }
 
+        // Write each byte of the chunk. Unallocated bytes should read
+        // as 0xFF, which should no-op.
         cmd[0] = TPI_CMD_SST_PI;
-        cmd[1] = m->buf[i];
-        rc = pgm->cmd_tpi(pgm, cmd, 2, NULL, 0);
+        for (j = 0; j < chunk; j++) {
+          cmd[1] = m->buf[i+j];
+          rc = pgm->cmd_tpi(pgm, cmd, 2, NULL, 0);
+        }
 
-        cmd[1] = m->buf[i + 1];
-        rc = pgm->cmd_tpi(pgm, cmd, 2, NULL, 0);
-
-        lastaddr += 2;
+        lastaddr += chunk;
 
         while (avr_tpi_poll_nvmbsy(pgm));
       }
       report_progress(i, wsize, NULL);
     }
+
     return i;
   }
 
diff --git a/src/avrdude.conf.in b/src/avrdude.conf.in
index 195db4cc..49fee90d 100644
--- a/src/avrdude.conf.in
+++ b/src/avrdude.conf.in
@@ -148,6 +148,7 @@
 #           size            = <num> ;             # bytes
 #           page_size       = <num> ;             # bytes
 #           num_pages       = <num> ;             # numeric
+#           n_word_writes   = <num> ;             # TPI only: if set, number of words to write
 #           min_write_delay = <num> ;             # micro-seconds
 #           max_write_delay = <num> ;             # micro-seconds
 #           readback        = <num> <num> ;       # pair of byte values
@@ -13346,6 +13347,7 @@ part parent ".reduced_core_tiny"
     memory "flash"
         size               = 2048;
         page_size          = 16;
+        n_word_writes      = 2;
         offset             = 0x4000;
         blocksize          = 128;
     ;
@@ -13365,6 +13367,7 @@ part parent ".reduced_core_tiny"
     memory "flash"
         size               = 4096;
         page_size          = 64;
+        n_word_writes      = 4;
         offset             = 0x4000;
         blocksize          = 128;
     ;
diff --git a/src/config.c b/src/config.c
index d43102c9..24373cac 100644
--- a/src/config.c
+++ b/src/config.c
@@ -70,6 +70,9 @@ Component_t avr_comp[] = {
   part_comp_desc(mcuid, COMP_INT),
   part_comp_desc(n_interrupts, COMP_INT),
   part_comp_desc(n_page_erase, COMP_INT),
+
+  // AVRMEM
+  mem_comp_desc(n_word_writes, COMP_INT),
 };
 
 #define DEBUG 0
diff --git a/src/developer_opts.c b/src/developer_opts.c
index ba0bdcff..383fb12b 100644
--- a/src/developer_opts.c
+++ b/src/developer_opts.c
@@ -746,6 +746,7 @@ static void dev_part_strct(const AVRPART *p, bool tsv, const AVRPART *base, bool
     _if_memout(intcmp, m->size > 8192? "0x%x": "%d", size);
     _if_memout(intcmp, "%d", page_size);
     _if_memout(intcmp, "%d", num_pages);
+    _if_memout(intcmp, "%d", n_word_writes);
     _if_memout(intcmp, "0x%x", offset);
     _if_memout(intcmp, "%d", min_write_delay);
     _if_memout(intcmp, "%d", max_write_delay);
diff --git a/src/doc/avrdude.texi b/src/doc/avrdude.texi
index 10b44dc0..1e139234 100644
--- a/src/doc/avrdude.texi
+++ b/src/doc/avrdude.texi
@@ -1930,6 +1930,7 @@ part
         size            = <num> ;             # bytes
         page_size       = <num> ;             # bytes
         num_pages       = <num> ;             # numeric
+        n_word_writes   = <num> ;             # TPI only: if set, number of words to write
         min_write_delay = <num> ;             # micro-seconds
         max_write_delay = <num> ;             # micro-seconds
         readback        = <num> <num> ;       # pair of byte values
diff --git a/src/lexer.l b/src/lexer.l
index b7652065..b637c4b7 100644
--- a/src/lexer.l
+++ b/src/lexer.l
@@ -121,7 +121,7 @@ SIGN     [+-]
      }
 
 
-prog_modes|mcuid|n_interrupts|n_page_erase { /* Components for assignment  */
+prog_modes|mcuid|n_interrupts|n_page_erase|n_word_writes { /* Components for assignment  */
   Component_t *cp = cfg_comp_search(yytext, current_strct);
   if(!cp) {
     yyerror("Unknown component %s in %s", yytext, cfg_strct_name(current_strct));
diff --git a/src/libavrdude.h b/src/libavrdude.h
index 01f80a72..b63dcd27 100644
--- a/src/libavrdude.h
+++ b/src/libavrdude.h
@@ -309,6 +309,7 @@ typedef struct avrmem {
   int size;                   /* total memory size in bytes */
   int page_size;              /* size of memory page (if page addressed) */
   int num_pages;              /* number of pages (if page addressed) */
+  int n_word_writes;          /* TPI only: number words to write at a time */
   unsigned int offset;        /* offset in IO memory (ATxmega) */
   int min_write_delay;        /* microseconds */
   int max_write_delay;        /* microseconds */