From f4dab5adac1465eb2b608a3a98b8bb06c460f5fb Mon Sep 17 00:00:00 2001
From: Franck Jullien <franck.jullien@gmail.com>
Date: Fri, 17 Dec 2021 20:45:41 +0100
Subject: [PATCH] bios: add write size option to cmd mem_write

---
 litex/soc/software/bios/cmds/cmd_mem.c | 31 +++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/litex/soc/software/bios/cmds/cmd_mem.c b/litex/soc/software/bios/cmds/cmd_mem.c
index d2bb04fff..2a5b33546 100644
--- a/litex/soc/software/bios/cmds/cmd_mem.c
+++ b/litex/soc/software/bios/cmds/cmd_mem.c
@@ -2,6 +2,7 @@
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <stdint.h>
 #include <libbase/memtest.h>
 
 #include <generated/csr.h>
@@ -69,17 +70,20 @@ define_command(mem_read, mem_read_handler, "Read address space", MEM_CMDS);
 static void mem_write_handler(int nb_params, char **params)
 {
 	char *c;
-	unsigned int *addr;
+	void *addr;
 	unsigned int value;
 	unsigned int count;
+	unsigned int size;
 	unsigned int i;
 
 	if (nb_params < 2) {
-		printf("mem_write <address> <value> [count]");
+		printf("mem_write <address> <value> [count] [size]");
 		return;
 	}
 
-	addr = (unsigned int *)strtoul(params[0], &c, 0);
+	size = 4;
+	addr = (void *)strtoul(params[0], &c, 0);
+
 	if (*c != 0) {
 		printf("Incorrect address");
 		return;
@@ -101,8 +105,25 @@ static void mem_write_handler(int nb_params, char **params)
 		}
 	}
 
-	for (i = 0; i < count; i++)
-		*addr++ = value;
+	if (nb_params == 4)
+		size = strtoul(params[3], &c, 0);
+
+	for (i = 0; i < count; i++) {
+		switch (size) {
+		case 1:
+			*(uint8_t *)addr++ = value;
+			break;
+		case 2:
+			*(uint16_t *)addr++ = value;
+			break;
+		case 4:
+			*(uint32_t *)addr++ = value;
+			break;
+		default:
+			printf("Incorrect size");
+			return;
+		}
+	}
 }
 
 define_command(mem_write, mem_write_handler, "Write address space", MEM_CMDS);