nix-scripts/artiq-fast/pkgs/openocd-jtagspi.diff

diff --git a/doc/openocd.texi b/doc/openocd.texi
index 138922d08..ad9f10d2e 100644
--- a/doc/openocd.texi
+++ b/doc/openocd.texi
@@ -5565,6 +5565,10 @@ will not work. These include all @command{*_image} and
 functionality is available through the @command{flash write_bank},
 @command{flash read_bank}, and @command{flash verify_bank} commands.
 
+According to device size, 1- to 4-byte addresses are sent. However, some
+flash chips additionally have to be switched to 4-byte addresses by an extra
+command, see below.
+
 @itemize
 @item @var{ir} ... is loaded into the JTAG IR to map the flash as the JTAG DR.
 For the bitstreams generated from @file{xilinx_bscan_spi.py} this is the
@@ -5577,6 +5581,29 @@ set _XILINX_USER1 0x02
 flash bank $_FLASHNAME spi 0x0 0 0 0 \
            $_TARGETNAME $_XILINX_USER1
 @end example
+
+@deffn Command {jtagspi set} bank_id name total_size page_size read_cmd unused pprg_cmd mass_erase_cmd sector_size sector_erase_cmd
+Sets flash parameters: @var{name} human readable string, @var{total_size}
+size in bytes, @var{page_size} is write page size. @var{read_cmd} and @var{pprg_cmd}
+are commands for read and page program, respectively. @var{mass_erase_cmd},
+@var{sector_size} and @var{sector_erase_cmd} are optional.
+@example
+jtagspi set 0 w25q128 0x1000000 0x100 0x03 0 0x02 0xC7 0x10000 0xD8
+@end example
+@end deffn
+
+@deffn Command {jtagspi cmd} bank_id resp_num cmd_byte ...
+Sends command @var{cmd_byte} and at most 20 following bytes and reads
+@var{resp_num} bytes afterwards. E.g. for 'Enter 4-byte address mode'
+@example
+jtagspi cmd 0 0 0xB7
+@end example
+@end deffn
+
+@deffn Command {jtagspi always_4byte} bank_id [ on | off ]
+Some devices use 4-byte addresses for all commands except the legacy 0x03 read
+regardless of device size. This command controls the corresponding hack.
+@end deffn
 @end deffn
 
 @deffn {Flash Driver} {xcf}
diff --git a/src/flash/nor/jtagspi.c b/src/flash/nor/jtagspi.c
index dc49fda61..e9a643d12 100644
--- a/src/flash/nor/jtagspi.c
+++ b/src/flash/nor/jtagspi.c
@@ -29,9 +29,12 @@
 
 struct jtagspi_flash_bank {
 	struct jtag_tap *tap;
-	const struct flash_device *dev;
+	struct flash_device dev;
+	char devname[32];
 	bool probed;
+	bool always_4byte;			/* use always 4-byte address except for basic read 0x03 */
 	uint32_t ir;
+	unsigned int addr_len;		/* address length in bytes */
 };
 
 FLASH_BANK_COMMAND_HANDLER(jtagspi_flash_bank_command)
@@ -42,10 +45,11 @@ FLASH_BANK_COMMAND_HANDLER(jtagspi_flash_bank_command)
 		return ERROR_COMMAND_SYNTAX_ERROR;
 
 	info = malloc(sizeof(struct jtagspi_flash_bank));
-	if (!info) {
+	if (info == NULL) {
 		LOG_ERROR("no memory for flash bank info");
 		return ERROR_FAIL;
 	}
+	bank->sectors = NULL;
 	bank->driver_priv = info;
 
 	info->tap = NULL;
@@ -69,70 +73,59 @@ static void jtagspi_set_ir(struct flash_bank *bank)
 	jtag_add_ir_scan(info->tap, &field, TAP_IDLE);
 }
 
-static void flip_u8(uint8_t *in, uint8_t *out, int len)
+static void flip_u8(const uint8_t *in, uint8_t *out, unsigned int len)
 {
-	for (int i = 0; i < len; i++)
+	for (unsigned int i = 0; i < len; i++)
 		out[i] = flip_u32(in[i], 8);
 }
 
 static int jtagspi_cmd(struct flash_bank *bank, uint8_t cmd,
-		uint32_t *addr, uint8_t *data, int len)
+		uint8_t *write_buffer, unsigned int write_len, uint8_t *data_buffer, int data_len)
 {
-	struct jtagspi_flash_bank *info = bank->driver_priv;
+	assert(write_buffer || write_len == 0);
+	assert(data_buffer || data_len == 0);
+
 	struct scan_field fields[6];
-	uint8_t marker = 1;
-	uint8_t xfer_bits_buf[4];
-	uint8_t addr_buf[3];
-	uint8_t *data_buf;
-	uint32_t xfer_bits;
-	int is_read, lenb, n;
 
-	/* LOG_DEBUG("cmd=0x%02x len=%i", cmd, len); */
+	LOG_DEBUG("cmd=0x%02x write_len=%d data_len=%d", cmd, write_len, data_len);
 
-	is_read = (len < 0);
+	/* negative data_len == read operation */
+	const bool is_read = (data_len < 0);
 	if (is_read)
-		len = -len;
-
-	n = 0;
+		data_len = -data_len;
 
+	int n = 0;
+	const uint8_t marker = 1;
 	fields[n].num_bits = 1;
 	fields[n].out_value = &marker;
 	fields[n].in_value = NULL;
 	n++;
 
-	xfer_bits = 8 + len - 1;
-	/* cmd + read/write - 1 due to the counter implementation */
-	if (addr)
-		xfer_bits += 24;
-	h_u32_to_be(xfer_bits_buf, xfer_bits);
-	flip_u8(xfer_bits_buf, xfer_bits_buf, 4);
-	fields[n].num_bits = 32;
-	fields[n].out_value = xfer_bits_buf;
+	/* transfer length = cmd + address + read/write,
+	 * -1 due to the counter implementation */
+	uint8_t xfer_bits[4];
+	h_u32_to_be(xfer_bits, ((sizeof(cmd) + write_len + data_len) * CHAR_BIT) - 1);
+	flip_u8(xfer_bits, xfer_bits, sizeof(xfer_bits));
+	fields[n].num_bits = sizeof(xfer_bits) * CHAR_BIT;
+	fields[n].out_value = xfer_bits;
 	fields[n].in_value = NULL;
 	n++;
 
-	cmd = flip_u32(cmd, 8);
-	fields[n].num_bits = 8;
+	flip_u8(&cmd, &cmd, sizeof(cmd));
+	fields[n].num_bits = sizeof(cmd) * CHAR_BIT;
 	fields[n].out_value = &cmd;
 	fields[n].in_value = NULL;
 	n++;
 
-	if (addr) {
-		h_u24_to_be(addr_buf, *addr);
-		flip_u8(addr_buf, addr_buf, 3);
-		fields[n].num_bits = 24;
-		fields[n].out_value = addr_buf;
+	if (write_len) {
+		flip_u8(write_buffer, write_buffer, write_len);
+		fields[n].num_bits = write_len * CHAR_BIT;
+		fields[n].out_value = write_buffer;
 		fields[n].in_value = NULL;
 		n++;
 	}
 
-	lenb = DIV_ROUND_UP(len, 8);
-	data_buf = malloc(lenb);
-	if (lenb > 0) {
-		if (!data_buf) {
-			LOG_ERROR("no memory for spi buffer");
-			return ERROR_FAIL;
-		}
+	if (data_len > 0) {
 		if (is_read) {
 			fields[n].num_bits = jtag_tap_count_enabled();
 			fields[n].out_value = NULL;
@@ -140,78 +133,313 @@ static int jtagspi_cmd(struct flash_bank *bank, uint8_t cmd,
 			n++;
 
 			fields[n].out_value = NULL;
-			fields[n].in_value = data_buf;
+			fields[n].in_value = data_buffer;
 		} else {
-			flip_u8(data, data_buf, lenb);
-			fields[n].out_value = data_buf;
+			flip_u8(data_buffer, data_buffer, data_len);
+			fields[n].out_value = data_buffer;
 			fields[n].in_value = NULL;
 		}
-		fields[n].num_bits = len;
+		fields[n].num_bits = data_len * CHAR_BIT;
 		n++;
 	}
 
 	jtagspi_set_ir(bank);
 	/* passing from an IR scan to SHIFT-DR clears BYPASS registers */
+	struct jtagspi_flash_bank *info = bank->driver_priv;
 	jtag_add_dr_scan(info->tap, n, fields, TAP_IDLE);
 	int retval = jtag_execute_queue();
 
 	if (is_read)
-		flip_u8(data_buf, data, lenb);
-	free(data_buf);
+		flip_u8(data_buffer, data_buffer, data_len);
 	return retval;
 }
 
+COMMAND_HANDLER(jtagspi_handle_set)
+{
+	struct flash_bank *bank = NULL;
+	struct jtagspi_flash_bank *info = NULL;
+	struct flash_sector *sectors = NULL;
+	uint32_t temp;
+	unsigned int index = 1;
+	int retval;
+
+	LOG_DEBUG("%s", __func__);
+
+	/* there are 6 mandatory arguments:
+	 * devname, size_in_bytes, pagesize, read_cmd, unused, pprog_cmd */
+	if (index + 6 > CMD_ARGC) {
+		command_print(CMD, "jtagspi: not enough arguments");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	retval = CALL_COMMAND_HANDLER(flash_command_get_bank, 0, &bank);
+	if (ERROR_OK != retval)
+		return retval;
+	info = bank->driver_priv;
+
+	/* invalidate all old info */
+	if (info->probed) {
+		bank->size = 0;
+		bank->num_sectors = 0;
+		if (bank->sectors)
+			free(bank->sectors);
+		bank->sectors = NULL;
+		info->always_4byte = false;
+		info->probed = false;
+	}
+	memset(&info->dev, 0, sizeof(info->dev));
+
+	strncpy(info->devname, CMD_ARGV[index++], sizeof(info->devname) - 1);
+	info->devname[sizeof(info->devname) - 1] = '\0';
+
+	COMMAND_PARSE_NUMBER(u32, CMD_ARGV[index++], temp);
+	info->dev.size_in_bytes = temp;
+	if ((temp & (temp - 1)) || (temp < (1UL << 8))) {
+		command_print(CMD, "jtagspi: device size must be 2^n with n >= 8");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	COMMAND_PARSE_NUMBER(u32, CMD_ARGV[index++], temp);
+	info->dev.pagesize = temp;
+	if (info->dev.pagesize == 0)
+		info->dev.pagesize = SPIFLASH_DEF_PAGESIZE;
+	if ((temp & (temp - 1)) || (temp > info->dev.size_in_bytes)) {
+		command_print(CMD, "jtagspi: page size must be 2^n and <= device size");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], info->dev.read_cmd);
+	if ((info->dev.read_cmd != 0x03) &&
+		(info->dev.read_cmd != 0x13)) {
+		command_print(CMD, "jtagspi: only 0x03/0x13 READ allowed");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], info->dev.qread_cmd);
+
+	COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], info->dev.pprog_cmd);
+	if ((info->dev.pprog_cmd != 0x02) &&
+		(info->dev.pprog_cmd != 0x12)) {
+		command_print(CMD, "jtagspi: only 0x02/0x12 PPRG allowed");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	/* remaining params are optional */
+	if (index < CMD_ARGC)
+		COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], info->dev.chip_erase_cmd);
+	else
+		info->dev.chip_erase_cmd = 0x00;
+
+	if (index < CMD_ARGC) {
+		COMMAND_PARSE_NUMBER(u32, CMD_ARGV[index++], temp);
+		info->dev.sectorsize = temp;
+		if ((info->dev.sectorsize > info->dev.size_in_bytes) ||
+			(info->dev.sectorsize < info->dev.pagesize) || (temp & (temp - 1))) {
+			command_print(CMD, "jtagspi: sector size must be 2^n and <= device size");
+			return ERROR_COMMAND_SYNTAX_ERROR;
+		}
+
+		if (index < CMD_ARGC)
+			COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], info->dev.erase_cmd);
+		else {
+			command_print(CMD, "jtagspi: erase command missing");
+			return ERROR_COMMAND_SYNTAX_ERROR;
+		}
+	} else {
+		/* no sector size / sector erase cmd given, treat whole bank as a single sector */
+		info->dev.erase_cmd = 0x00;
+		info->dev.sectorsize = info->dev.size_in_bytes;
+	}
+
+	if (index < CMD_ARGC) {
+		command_print(CMD, "jtagspi: extra arguments");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	/* set correct size value */
+	bank->size = info->dev.size_in_bytes;
+
+	/* calculate address length in bytes */
+	if (bank->size <= (1UL << 8))
+		info->addr_len = 1;
+	else if (bank->size <= (1UL << 16))
+		info->addr_len = 2;
+	else if (bank->size <= (1UL << 24))
+		info->addr_len = 3;
+	else {
+		info->addr_len = 4;
+		LOG_WARNING("4-byte addresses needed, might need extra command to enable");
+	}
+
+	/* create and fill sectors array */
+	bank->num_sectors =
+		info->dev.size_in_bytes / info->dev.sectorsize;
+	sectors = malloc(sizeof(struct flash_sector) * bank->num_sectors);
+	if (sectors == NULL) {
+		LOG_ERROR("Not enough memory");
+		return ERROR_FAIL;
+	}
+
+	for (unsigned int sector = 0; sector < bank->num_sectors; sector++) {
+		sectors[sector].offset = sector * (info->dev.sectorsize);
+		sectors[sector].size = info->dev.sectorsize;
+		sectors[sector].is_erased = -1;
+		sectors[sector].is_protected = 0;
+	}
+
+	bank->sectors = sectors;
+	info->dev.name = info->devname;
+	if (info->dev.size_in_bytes / 4096)
+		LOG_INFO("flash \'%s\' id = unknown\nflash size = %" PRIu32 " kbytes",
+			info->dev.name, info->dev.size_in_bytes / 1024);
+	else
+		LOG_INFO("flash \'%s\' id = unknown\nflash size = %" PRIu32 " bytes",
+			info->dev.name, info->dev.size_in_bytes);
+	info->probed = true;
+
+	return ERROR_OK;
+}
+
+COMMAND_HANDLER(jtagspi_handle_cmd)
+{
+	struct flash_bank *bank;
+	unsigned int index = 1;
+	const int max = 21;
+	uint8_t num_write, num_read, write_buffer[max], read_buffer[1 << CHAR_BIT];
+	uint8_t data, *ptr;
+	char temp[4], output[(2 + max + (1 << CHAR_BIT)) * 3 + 8];
+	int retval;
+
+	LOG_DEBUG("%s", __func__);
+
+	if (CMD_ARGC < 3) {
+		command_print(CMD, "jtagspi: not enough arguments");
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	num_write = CMD_ARGC - 2;
+	if (num_write > max) {
+		LOG_ERROR("at most %d bytes may be send", max);
+		return ERROR_COMMAND_SYNTAX_ERROR;
+	}
+
+	retval = CALL_COMMAND_HANDLER(flash_command_get_bank, 0, &bank);
+	if (ERROR_OK != retval)
+		return retval;
+
+	COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index++], num_read);
+
+	snprintf(output, sizeof(output), "spi: ");
+	for (ptr = &write_buffer[0] ; index < CMD_ARGC; index++) {
+		COMMAND_PARSE_NUMBER(u8, CMD_ARGV[index], data);
+		*ptr++ = data;
+		snprintf(temp, sizeof(temp), "%02" PRIx8 " ", data);
+		strncat(output, temp, sizeof(output) - strlen(output) - 1);
+	}
+	strncat(output, "-> ", sizeof(output) - strlen(output) - 1);
+
+	/* process command */
+	ptr = &read_buffer[0];
+	jtagspi_cmd(bank, write_buffer[0], &write_buffer[1], num_write - 1, ptr, -num_read);
+	if (retval != ERROR_OK)
+		return retval;
+
+	for ( ; num_read > 0; num_read--) {
+		snprintf(temp, sizeof(temp), "%02" PRIx8 " ", *ptr++);
+		strncat(output, temp, sizeof(output) - strlen(output) - 1);
+	}
+	command_print(CMD, "%s", output);
+
+	return ERROR_OK;
+}
+
+COMMAND_HANDLER(jtagspi_handle_always_4byte)
+{
+	struct flash_bank *bank;
+	struct jtagspi_flash_bank *jtagspi_info;
+	int retval;
+
+	LOG_DEBUG("%s", __func__);
+
+	if ((CMD_ARGC != 1) && (CMD_ARGC != 2))
+		return ERROR_COMMAND_SYNTAX_ERROR;
+
+	retval = CALL_COMMAND_HANDLER(flash_command_get_bank, 0, &bank);
+	if (ERROR_OK != retval)
+		return retval;
+
+	jtagspi_info = bank->driver_priv;
+
+	if (CMD_ARGC == 1)
+		command_print(CMD, jtagspi_info->always_4byte ? "on" : "off");
+	else
+		COMMAND_PARSE_BOOL(CMD_ARGV[1], jtagspi_info->always_4byte, "on", "off");
+
+	return ERROR_OK;
+}
+
 static int jtagspi_probe(struct flash_bank *bank)
 {
 	struct jtagspi_flash_bank *info = bank->driver_priv;
 	struct flash_sector *sectors;
+	const struct flash_device *p;
 	uint8_t in_buf[3];
 	uint32_t id, sectorsize;
 
-	if (info->probed)
+	if (bank->sectors) {
 		free(bank->sectors);
+		bank->sectors = NULL;
+	}
 	info->probed = false;
 
-	if (!bank->target->tap) {
+	if (bank->target->tap == NULL) {
 		LOG_ERROR("Target has no JTAG tap");
 		return ERROR_FAIL;
 	}
 	info->tap = bank->target->tap;
 
-	jtagspi_cmd(bank, SPIFLASH_READ_ID, NULL, in_buf, -24);
+	jtagspi_cmd(bank, SPIFLASH_READ_ID, NULL, 0, in_buf, -3);
 	/* the table in spi.c has the manufacturer byte (first) as the lsb */
 	id = le_to_h_u24(in_buf);
 
-	info->dev = NULL;
-	for (const struct flash_device *p = flash_devices; p->name ; p++)
+	memset(&info->dev, 0, sizeof(info->dev));
+	for (p = flash_devices; p->name ; p++)
 		if (p->device_id == id) {
-			info->dev = p;
+			memcpy(&info->dev, p, sizeof(info->dev));
 			break;
 		}
 
-	if (!(info->dev)) {
-		LOG_ERROR("Unknown flash device (ID 0x%08" PRIx32 ")", id);
+	if (!(p->name)) {
+		LOG_ERROR("Unknown flash device (ID 0x%06" PRIx32 ")", id & 0xFFFFFF);
 		return ERROR_FAIL;
 	}
 
-	LOG_INFO("Found flash device \'%s\' (ID 0x%08" PRIx32 ")",
-		info->dev->name, info->dev->device_id);
+	LOG_INFO("Found flash device \'%s\' (ID 0x%06" PRIx32 ")",
+		info->dev.name, info->dev.device_id & 0xFFFFFF);
 
 	/* Set correct size value */
-	bank->size = info->dev->size_in_bytes;
-	if (bank->size <= (1UL << 16))
-		LOG_WARNING("device needs 2-byte addresses - not implemented");
-	if (bank->size > (1UL << 24))
-		LOG_WARNING("device needs paging or 4-byte addresses - not implemented");
+	bank->size = info->dev.size_in_bytes;
+
+	/* calculate address length in bytes */
+	if (bank->size <= (1UL << 8))
+		info->addr_len = 1;
+	else if (bank->size <= (1UL << 16))
+		info->addr_len = 2;
+	else if (bank->size <= (1UL << 24))
+		info->addr_len = 3;
+	else {
+		info->addr_len = 4;
+		LOG_WARNING("4-byte addresses needed, might need extra command to enable");
+	}
 
 	/* if no sectors, treat whole bank as single sector */
-	sectorsize = info->dev->sectorsize ?
-		info->dev->sectorsize : info->dev->size_in_bytes;
+	sectorsize = info->dev.sectorsize ?
+		info->dev.sectorsize : info->dev.size_in_bytes;
 
 	/* create and fill sectors array */
-	bank->num_sectors = info->dev->size_in_bytes / sectorsize;
+	bank->num_sectors = info->dev.size_in_bytes / sectorsize;
 	sectors = malloc(sizeof(struct flash_sector) * bank->num_sectors);
-	if (!sectors) {
+	if (sectors == NULL) {
 		LOG_ERROR("not enough memory");
 		return ERROR_FAIL;
 	}
@@ -228,27 +456,35 @@ static int jtagspi_probe(struct flash_bank *bank)
 	return ERROR_OK;
 }
 
+static int jtagspi_auto_probe(struct flash_bank *bank)
+{
+	struct jtagspi_flash_bank *info = bank->driver_priv;
+
+	if (info->probed)
+		return ERROR_OK;
+	return jtagspi_probe(bank);
+}
+
 static int jtagspi_read_status(struct flash_bank *bank, uint32_t *status)
 {
 	uint8_t buf;
-	int err = jtagspi_cmd(bank, SPIFLASH_READ_STATUS, NULL, &buf, -8);
+	int err = jtagspi_cmd(bank, SPIFLASH_READ_STATUS, NULL, 0, &buf, -1);
 	if (err == ERROR_OK) {
 		*status = buf;
-		/* LOG_DEBUG("status=0x%08" PRIx32, *status); */
+		LOG_DEBUG("status=0x%02" PRIx8, *status);
 	}
-
 	return err;
 }
 
 static int jtagspi_wait(struct flash_bank *bank, int timeout_ms)
 {
-	uint32_t status;
 	int64_t t0 = timeval_ms();
 	int64_t dt;
 
 	do {
 		dt = timeval_ms() - t0;
 
+		uint32_t status = (uint32_t)-1;
 		int retval = jtagspi_read_status(bank, &status);
 		if (retval != ERROR_OK)
 			return retval;
@@ -266,16 +502,15 @@ static int jtagspi_wait(struct flash_bank *bank, int timeout_ms)
 
 static int jtagspi_write_enable(struct flash_bank *bank)
 {
-	uint32_t status;
-
-	jtagspi_cmd(bank, SPIFLASH_WRITE_ENABLE, NULL, NULL, 0);
+	jtagspi_cmd(bank, SPIFLASH_WRITE_ENABLE, NULL, 0, NULL, 0);
 
+	uint32_t status = (uint32_t)-1;
 	int retval = jtagspi_read_status(bank, &status);
 	if (retval != ERROR_OK)
 		return retval;
 
 	if ((status & SPIFLASH_WE_BIT) == 0) {
-		LOG_ERROR("Cannot enable write to flash. Status=0x%08" PRIx32, status);
+		LOG_ERROR("Cannot enable write to flash. Status=0x%02" PRIx8, status);
 		return ERROR_FAIL;
 	}
 	return ERROR_OK;
@@ -287,28 +522,51 @@ static int jtagspi_bulk_erase(struct flash_bank *bank)
 	int retval;
 	int64_t t0 = timeval_ms();
 
-	if (info->dev->chip_erase_cmd == 0x00)
+	if (info->dev.chip_erase_cmd == 0x00)
 		return ERROR_FLASH_OPER_UNSUPPORTED;
 
 	retval = jtagspi_write_enable(bank);
 	if (retval != ERROR_OK)
 		return retval;
-	jtagspi_cmd(bank, info->dev->chip_erase_cmd, NULL, NULL, 0);
-	retval = jtagspi_wait(bank, bank->num_sectors*JTAGSPI_MAX_TIMEOUT);
+
+	jtagspi_cmd(bank, info->dev.chip_erase_cmd, NULL, 0, NULL, 0);
+	if (retval != ERROR_OK)
+		return retval;
+
+	retval = jtagspi_wait(bank, bank->num_sectors * JTAGSPI_MAX_TIMEOUT);
 	LOG_INFO("took %" PRId64 " ms", timeval_ms() - t0);
 	return retval;
 }
 
+static uint8_t *fill_addr(uint32_t addr, unsigned int addr_len, uint8_t *buffer)
+{
+	for (buffer += addr_len; addr_len > 0; --addr_len) {
+		*--buffer = addr;
+		addr >>= 8;
+	}
+
+	return buffer;
+}
+
 static int jtagspi_sector_erase(struct flash_bank *bank, unsigned int sector)
 {
 	struct jtagspi_flash_bank *info = bank->driver_priv;
 	int retval;
+	uint8_t addr[sizeof(uint32_t)];
 	int64_t t0 = timeval_ms();
 
 	retval = jtagspi_write_enable(bank);
 	if (retval != ERROR_OK)
 		return retval;
-	jtagspi_cmd(bank, info->dev->erase_cmd, &bank->sectors[sector].offset, NULL, 0);
+
+	/* ATXP032/064/128 use always 4-byte addresses except for 0x03 read */
+	unsigned int addr_len = info->always_4byte ? 4 : info->addr_len;
+
+	retval = jtagspi_cmd(bank, info->dev.erase_cmd, fill_addr(bank->sectors[sector].offset, addr_len, addr),
+			addr_len, NULL, 0);
+	if (retval != ERROR_OK)
+		return retval;
+
 	retval = jtagspi_wait(bank, JTAGSPI_MAX_TIMEOUT);
 	LOG_INFO("sector %u took %" PRId64 " ms", sector, timeval_ms() - t0);
 	return retval;
@@ -339,8 +597,9 @@ static int jtagspi_erase(struct flash_bank *bank, unsigned int first,
 		}
 	}
 
-	if (first == 0 && last == (bank->num_sectors - 1)
-		&& info->dev->chip_erase_cmd != info->dev->erase_cmd) {
+	if (first == 0 && last == (bank->num_sectors - 1) &&
+		info->dev.chip_erase_cmd != 0x00 &&
+		info->dev.chip_erase_cmd != info->dev.erase_cmd) {
 		LOG_DEBUG("Trying bulk erase.");
 		retval = jtagspi_bulk_erase(bank);
 		if (retval == ERROR_OK)
@@ -349,7 +608,7 @@ static int jtagspi_erase(struct flash_bank *bank, unsigned int first,
 			LOG_WARNING("Bulk flash erase failed. Falling back to sector erase.");
 	}
 
-	if (info->dev->erase_cmd == 0x00)
+	if (info->dev.erase_cmd == 0x00)
 		return ERROR_FLASH_OPER_UNSUPPORTED;
 
 	for (unsigned int sector = first; sector <= last; sector++) {
@@ -374,49 +633,93 @@ static int jtagspi_protect(struct flash_bank *bank, int set, unsigned int first,
 static int jtagspi_read(struct flash_bank *bank, uint8_t *buffer, uint32_t offset, uint32_t count)
 {
 	struct jtagspi_flash_bank *info = bank->driver_priv;
+	uint32_t pagesize, currsize;
+	uint8_t addr[sizeof(uint32_t)];
+	int retval;
 
 	if (!(info->probed)) {
-		LOG_ERROR("Flash bank not yet probed.");
+		LOG_ERROR("Flash bank not probed.");
 		return ERROR_FLASH_BANK_NOT_PROBED;
 	}
 
-	jtagspi_cmd(bank, SPIFLASH_READ, &offset, buffer, -count*8);
+	/* if no sectorsize, use reasonable default */
+	pagesize = info->dev.sectorsize ? info->dev.sectorsize : info->dev.pagesize;
+	if (pagesize == 0)
+		pagesize = (info->dev.size_in_bytes <= SPIFLASH_DEF_PAGESIZE) ?
+			info->dev.size_in_bytes : SPIFLASH_DEF_PAGESIZE;
+
+	/* ATXP032/064/128 use always 4-byte addresses except for 0x03 read */
+	unsigned int addr_len = ((info->dev.read_cmd != 0x03) && info->always_4byte) ? 4 : info->addr_len;
+
+	while (count > 0) {
+		/* length up to end of current page */
+		currsize = ((offset + pagesize) & ~(pagesize - 1)) - offset;
+		/* but no more than remaining size */
+		currsize = (count < currsize) ? count : currsize;
+
+		retval = jtagspi_cmd(bank, info->dev.read_cmd, fill_addr(offset, addr_len, addr),
+			addr_len, buffer, -currsize);
+		if (retval != ERROR_OK) {
+			LOG_ERROR("page read error");
+			return retval;
+		}
+		LOG_DEBUG("read page at 0x%08" PRIx32, offset);
+		offset += currsize;
+		buffer += currsize;
+		count -= currsize;
+	}
 	return ERROR_OK;
 }
 
 static int jtagspi_page_write(struct flash_bank *bank, const uint8_t *buffer, uint32_t offset, uint32_t count)
 {
+	struct jtagspi_flash_bank *info = bank->driver_priv;
+	uint8_t addr[sizeof(uint32_t)];
 	int retval;
 
 	retval = jtagspi_write_enable(bank);
 	if (retval != ERROR_OK)
 		return retval;
-	jtagspi_cmd(bank, SPIFLASH_PAGE_PROGRAM, &offset, (uint8_t *) buffer, count*8);
+
+	/* ATXP032/064/128 use always 4-byte addresses except for 0x03 read */
+	unsigned int addr_len = ((info->dev.read_cmd != 0x03) && info->always_4byte) ? 4 : info->addr_len;
+
+	retval = jtagspi_cmd(bank, info->dev.pprog_cmd, fill_addr(offset, addr_len, addr),
+		addr_len, (uint8_t *) buffer, count);
+	if (retval != ERROR_OK)
+		return retval;
 	return jtagspi_wait(bank, JTAGSPI_MAX_TIMEOUT);
 }
 
 static int jtagspi_write(struct flash_bank *bank, const uint8_t *buffer, uint32_t offset, uint32_t count)
 {
 	struct jtagspi_flash_bank *info = bank->driver_priv;
+	uint32_t pagesize, currsize;
 	int retval;
-	uint32_t n, pagesize;
 
 	if (!(info->probed)) {
-		LOG_ERROR("Flash bank not yet probed.");
+		LOG_ERROR("Flash bank not probed.");
 		return ERROR_FLASH_BANK_NOT_PROBED;
 	}
 
 	/* if no write pagesize, use reasonable default */
-	pagesize = info->dev->pagesize ? info->dev->pagesize : SPIFLASH_DEF_PAGESIZE;
+	pagesize = info->dev.pagesize ? info->dev.pagesize : SPIFLASH_DEF_PAGESIZE;
+
+	while (count > 0) {
+		/* length up to end of current page */
+		currsize = ((offset + pagesize) & ~(pagesize - 1)) - offset;
+		/* but no more than remaining size */
+		currsize = (count < currsize) ? count : currsize;
 
-	for (n = 0; n < count; n += pagesize) {
-		retval = jtagspi_page_write(bank, buffer + n, offset + n,
-				MIN(count - n, pagesize));
+		retval = jtagspi_page_write(bank, buffer, offset, currsize);
 		if (retval != ERROR_OK) {
 			LOG_ERROR("page write error");
 			return retval;
 		}
-		LOG_DEBUG("wrote page at 0x%08" PRIx32, offset + n);
+		LOG_DEBUG("wrote page at 0x%08" PRIx32, offset);
+		offset += currsize;
+		buffer += currsize;
+		count -= currsize;
 	}
 	return ERROR_OK;
 }
@@ -430,22 +733,72 @@ static int jtagspi_info(struct flash_bank *bank, struct command_invocation *cmd)
 		return ERROR_OK;
 	}
 
-	command_print_sameline(cmd, "\nSPIFI flash information:\n"
-		"  Device \'%s\' (ID 0x%08" PRIx32 ")\n",
-		info->dev->name, info->dev->device_id);
+	command_print_sameline(cmd, "flash \'%s\', device id = 0x%06" PRIx32
+		", flash size = %" PRIu32 " %sbytes\n(page size = %" PRIu32
+		", read = 0x%02" PRIx8 ", qread = 0x%02" PRIx8
+		", pprog = 0x%02" PRIx8 ", mass_erase = 0x%02" PRIx8
+		", sector size = %" PRIu32 " %sbytes, sector_erase = 0x%02" PRIx8 ")",
+		info->dev.name, info->dev.device_id & 0xFFFFFF,
+		bank->size / 4096 ? bank->size / 1024 : bank->size,
+		bank->size / 4096 ? "k" : "", info->dev.pagesize,
+		info->dev.read_cmd, info->dev.qread_cmd,
+		info->dev.pprog_cmd, info->dev.chip_erase_cmd,
+		info->dev.sectorsize / 4096 ?
+		info->dev.sectorsize / 1024 : info->dev.sectorsize,
+		info->dev.sectorsize / 4096 ? "k" : "",
+		info->dev.erase_cmd);
 
 	return ERROR_OK;
 }
 
+static const struct command_registration jtagspi_exec_command_handlers[] = {
+	{
+		.name = "set",
+		.handler = jtagspi_handle_set,
+		.mode = COMMAND_EXEC,
+		.usage = "bank_id name chip_size page_size read_cmd unused pprg_cmd "
+			"[ mass_erase_cmd ] [ sector_size sector_erase_cmd ]",
+		.help = "Set device parameters if not autodetected.",
+	},
+	{
+		.name = "cmd",
+		.handler = jtagspi_handle_cmd,
+		.mode = COMMAND_EXEC,
+		.usage = "bank_id num_resp cmd_byte ...",
+		.help = "Send low-level command cmd_byte and following bytes, read num_bytes.",
+	},
+	{
+		.name = "always_4byte",
+		.handler = jtagspi_handle_always_4byte,
+		.mode = COMMAND_EXEC,
+		.usage = "bank_id [ on | off ]",
+		.help = "Use always 4-byte address except for basic 0x03.",
+	},
+
+	COMMAND_REGISTRATION_DONE
+};
+
+static const struct command_registration jtagspi_command_handlers[] = {
+	{
+		.name = "jtagspi",
+		.mode = COMMAND_ANY,
+		.help = "jtagspi command group",
+		.usage = "",
+		.chain = jtagspi_exec_command_handlers,
+	},
+	COMMAND_REGISTRATION_DONE
+};
+
 const struct flash_driver jtagspi_flash = {
 	.name = "jtagspi",
+	.commands = jtagspi_command_handlers,
 	.flash_bank_command = jtagspi_flash_bank_command,
 	.erase = jtagspi_erase,
 	.protect = jtagspi_protect,
 	.write = jtagspi_write,
 	.read = jtagspi_read,
 	.probe = jtagspi_probe,
-	.auto_probe = jtagspi_probe,
+	.auto_probe = jtagspi_auto_probe,
 	.erase_check = default_flash_blank_check,
 	.info = jtagspi_info,
 	.free_driver_priv = default_flash_free_driver_priv,