drivers/flash/flash_mspi_nor.c - third_party/github/zephyrproject-rtos/zephyr - Git at Google

 /*
  * Copyright (c) 2024 Nordic Semiconductor ASA
  *
  * SPDX-License-Identifier: Apache-2.0
  */

 #define DT_DRV_COMPAT jedec_mspi_nor

 #include <zephyr/drivers/flash.h>
 #include <zephyr/drivers/gpio.h>
 #include <zephyr/drivers/mspi.h>
 #include <zephyr/logging/log.h>
 #include <zephyr/pm/device.h>
 #include <zephyr/pm/device_runtime.h>

 #include "jesd216.h"
 #include "spi_nor.h"

 LOG_MODULE_REGISTER(flash_mspi_nor, CONFIG_FLASH_LOG_LEVEL);

 #if DT_ANY_INST_HAS_PROP_STATUS_OKAY(reset_gpios)
 #define WITH_RESET_GPIO 1
 #endif

 struct flash_mspi_nor_data {
 	struct k_sem acquired;
 	struct mspi_xfer_packet packet;
 	struct mspi_xfer xfer;
 };

 struct flash_mspi_nor_config {
 	const struct device *bus;
 	uint32_t flash_size;
 	struct mspi_dev_id mspi_id;
 	struct mspi_dev_cfg mspi_cfg;
 	enum mspi_dev_cfg_mask mspi_cfg_mask;
 #if defined(CONFIG_MSPI_XIP)
 	struct mspi_xip_cfg xip_cfg;
 #endif
 #if defined(WITH_RESET_GPIO)
 	struct gpio_dt_spec reset;
 	uint32_t reset_pulse_us;
 	uint32_t reset_recovery_us;
 #endif
 #if defined(CONFIG_FLASH_PAGE_LAYOUT)
 	struct flash_pages_layout layout;
 #endif
 	uint8_t jedec_id[SPI_NOR_MAX_ID_LEN];
 };

 static int acquire(const struct device *dev)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	int rc;

 	k_sem_take(&dev_data->acquired, K_FOREVER);

 	rc = pm_device_runtime_get(dev_config->bus);
 	if (rc < 0) {
 		LOG_ERR("pm_device_runtime_get() failed: %d", rc);
 	} else {
 		/* This acquires the MSPI controller and reconfigures it
 		 * if needed for the flash device.
 		 */
 		rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
 				     dev_config->mspi_cfg_mask,
 				     &dev_config->mspi_cfg);
 		if (rc < 0) {
 			LOG_ERR("mspi_dev_config() failed: %d", rc);
 		} else {
 			return 0;
 		}

 		(void)pm_device_runtime_put(dev_config->bus);
 	}

 	k_sem_give(&dev_data->acquired);
 	return rc;
 }

 static void release(const struct device *dev)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;

 	/* This releases the MSPI controller. */
 	(void)mspi_get_channel_status(dev_config->bus, 0);

 	(void)pm_device_runtime_put(dev_config->bus);

 	k_sem_give(&dev_data->acquired);
 }

 static inline uint32_t dev_flash_size(const struct device *dev)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;

 	return dev_config->flash_size;
 }

 static inline uint16_t dev_page_size(const struct device *dev)
 {
 	return SPI_NOR_PAGE_SIZE;
 }

 static int api_read(const struct device *dev, off_t addr, void *dest,
 		    size_t size)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	const uint32_t flash_size = dev_flash_size(dev);
 	int rc;

 	if (size == 0) {
 		return 0;
 	}

 	if ((addr < 0) || ((addr + size) > flash_size)) {
 		return -EINVAL;
 	}

 	rc = acquire(dev);
 	if (rc < 0) {
 		return rc;
 	}

 	/* TODO: get rid of all these hard-coded values for MX25Ux chips */
 	dev_data->xfer.cmd_length  = 2;
 	dev_data->xfer.addr_length = 4;
 	dev_data->xfer.rx_dummy    = 20;
 	dev_data->packet.dir       = MSPI_RX;
 	dev_data->packet.cmd       = SPI_NOR_OCMD_RD;
 	dev_data->packet.address   = addr;
 	dev_data->packet.data_buf  = dest;
 	dev_data->packet.num_bytes = size;
 	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 			     &dev_data->xfer);

 	release(dev);

 	if (rc < 0) {
 		LOG_ERR("SPI_NOR_OCMD_RD xfer failed: %d", rc);
 		return rc;
 	}

 	return 0;
 }

 static int wait_until_ready(const struct device *dev, k_timeout_t poll_period)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	uint8_t status_reg;
 	int rc;

 	while (true) {
 		dev_data->xfer.cmd_length  = 2;
 		dev_data->xfer.addr_length = 4;
 		dev_data->xfer.rx_dummy    = 4;
 		dev_data->packet.dir       = MSPI_RX;
 		dev_data->packet.cmd       = SPI_NOR_OCMD_RDSR;
 		dev_data->packet.address   = 0;
 		dev_data->packet.data_buf  = &status_reg;
 		dev_data->packet.num_bytes = sizeof(status_reg);
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("SPI_NOR_OCMD_RDSR xfer failed: %d", rc);
 			return rc;
 		}
 		if (!(status_reg & SPI_NOR_WIP_BIT)) {
 			break;
 		}

 		k_sleep(poll_period);
 	}

 	return 0;
 }

 static int api_write(const struct device *dev, off_t addr, const void *src,
 		     size_t size)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	const uint32_t flash_size = dev_flash_size(dev);
 	const uint16_t page_size = dev_page_size(dev);
 	int rc;

 	if (size == 0) {
 		return 0;
 	}

 	if ((addr < 0) || ((addr + size) > flash_size)) {
 		return -EINVAL;
 	}

 	rc = acquire(dev);
 	if (rc < 0) {
 		return rc;
 	}

 	while (size > 0) {
 		/* Split write into parts, each within one page only. */
 		uint16_t page_offset = (uint16_t)(addr % page_size);
 		uint16_t page_left = page_size - page_offset;
 		uint16_t to_write = (uint16_t)MIN(size, page_left);

 		dev_data->xfer.cmd_length  = 2;
 		dev_data->xfer.tx_dummy    = 0;
 		dev_data->packet.dir       = MSPI_TX;

 		dev_data->xfer.addr_length = 0;
 		dev_data->packet.cmd       = SPI_NOR_OCMD_WREN;
 		dev_data->packet.num_bytes = 0;
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("SPI_NOR_OCMD_WREN xfer failed: %d", rc);
 			break;
 		}

 		dev_data->xfer.addr_length = 4;
 		dev_data->packet.cmd       = SPI_NOR_OCMD_PAGE_PRG;
 		dev_data->packet.address   = addr;
 		dev_data->packet.data_buf  = (uint8_t *)src;
 		dev_data->packet.num_bytes = to_write;
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("SPI_NOR_OCMD_PAGE_PRG xfer failed: %d", rc);
 			break;
 		}

 		addr += to_write;
 		src   = (const uint8_t *)src + to_write;
 		size -= to_write;

 		rc = wait_until_ready(dev, K_MSEC(1));
 		if (rc < 0) {
 			break;
 		}
 	}

 	release(dev);

 	return rc;
 }

 static int api_erase(const struct device *dev, off_t addr, size_t size)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	const uint32_t flash_size = dev_flash_size(dev);
 	int rc = 0;

 	if ((addr < 0) || ((addr + size) > flash_size)) {
 		return -EINVAL;
 	}

 	if (!SPI_NOR_IS_SECTOR_ALIGNED(addr)) {
 		return -EINVAL;
 	}

 	if ((size % SPI_NOR_SECTOR_SIZE) != 0) {
 		return -EINVAL;
 	}

 	rc = acquire(dev);
 	if (rc < 0) {
 		return rc;
 	}

 	while (size > 0) {
 		dev_data->xfer.cmd_length  = 2;
 		dev_data->xfer.tx_dummy    = 0;
 		dev_data->packet.dir       = MSPI_TX;
 		dev_data->packet.num_bytes = 0;

 		dev_data->xfer.addr_length = 0;
 		dev_data->packet.cmd       = SPI_NOR_OCMD_WREN;
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("SPI_NOR_OCMD_WREN xfer failed: %d", rc);
 			break;
 		}

 		if (size == flash_size) {
 			/* Chip erase. */
 			dev_data->xfer.addr_length = 0;
 			dev_data->packet.cmd       = SPI_NOR_OCMD_CE;

 			size -= flash_size;
 		} else {
 			/* Sector erase. */
 			dev_data->xfer.addr_length = 4;
 			dev_data->packet.cmd       = SPI_NOR_OCMD_SE;
 			dev_data->packet.address   = addr;

 			addr += SPI_NOR_SECTOR_SIZE;
 			size -= SPI_NOR_SECTOR_SIZE;
 		}
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("Erase command 0x%02x xfer failed: %d",
 				dev_data->packet.cmd, rc);
 			break;
 		}

 		rc = wait_until_ready(dev, K_MSEC(1));
 		if (rc < 0) {
 			break;
 		}
 	}

 	release(dev);

 	return rc;
 }

 static const
 struct flash_parameters *api_get_parameters(const struct device *dev)
 {
 	ARG_UNUSED(dev);

 	static const struct flash_parameters parameters = {
 		.write_block_size = 1,
 		.erase_value = 0xff,
 	};

 	return &parameters;
 }

 static int read_jedec_id(const struct device *dev, uint8_t *id)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	int rc;

 	dev_data->xfer.cmd_length  = 2;
 	dev_data->xfer.addr_length = 4;
 	dev_data->xfer.rx_dummy    = 4;
 	dev_data->packet.dir       = MSPI_RX;
 	dev_data->packet.cmd       = JESD216_OCMD_READ_ID;
 	dev_data->packet.address   = 0;
 	dev_data->packet.data_buf  = id;
 	dev_data->packet.num_bytes = JESD216_READ_ID_LEN;
 	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 			     &dev_data->xfer);
 	if (rc < 0) {
 		printk("mspi_transceive() failed: %d\n", rc);
 		return rc;
 	}

 	return rc;
 }

 #if defined(CONFIG_FLASH_PAGE_LAYOUT)
 static void api_page_layout(const struct device *dev,
 			     const struct flash_pages_layout **layout,
 			     size_t *layout_size)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;

 	*layout = &dev_config->layout;
 	*layout_size = 1;
 }
 #endif /* CONFIG_FLASH_PAGE_LAYOUT */

 #if defined(CONFIG_FLASH_JESD216_API)
 static int api_sfdp_read(const struct device *dev, off_t addr, void *dest,
 			 size_t size)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	int rc;

 	if (size == 0) {
 		return 0;
 	}

 	rc = acquire(dev);
 	if (rc < 0) {
 		return rc;
 	}

 	dev_data->xfer.cmd_length  = 2;
 	dev_data->xfer.addr_length = 4;
 	dev_data->xfer.rx_dummy    = 20;
 	dev_data->packet.dir       = MSPI_RX;
 	dev_data->packet.cmd       = JESD216_OCMD_READ_SFDP;
 	dev_data->packet.address   = addr;
 	dev_data->packet.data_buf  = dest;
 	dev_data->packet.num_bytes = size;
 	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 			     &dev_data->xfer);
 	if (rc < 0) {
 		printk("JESD216_OCMD_READ_SFDP xfer failed: %d\n", rc);
 		return rc;
 	}

 	release(dev);

 	return rc;
 }

 static int api_read_jedec_id(const struct device *dev, uint8_t *id)
 {
 	int rc = 0;

 	rc = acquire(dev);
 	if (rc < 0) {
 		return rc;
 	}

 	rc = read_jedec_id(dev, id);

 	release(dev);

 	return rc;
 }
 #endif /* CONFIG_FLASH_JESD216_API  */

 static int dev_pm_action_cb(const struct device *dev,
 			    enum pm_device_action action)
 {
 	switch (action) {
 	case PM_DEVICE_ACTION_SUSPEND:
 		break;
 	case PM_DEVICE_ACTION_RESUME:
 		break;
 	default:
 		return -ENOTSUP;
 	}

 	return 0;
 }

 static int flash_chip_init(const struct device *dev)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	struct mspi_dev_cfg init_dev_cfg = dev_config->mspi_cfg;
 	uint8_t id[JESD216_READ_ID_LEN] = {0};
 	int rc;

 	init_dev_cfg.freq = MHZ(1);
 	init_dev_cfg.io_mode = MSPI_IO_MODE_SINGLE;

 	rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
 			     MSPI_DEVICE_CONFIG_ALL, &init_dev_cfg);
 	if (rc < 0) {
 		LOG_ERR("Failed to set initial device config: %d", rc);
 		return rc;
 	}

 	dev_data->xfer.xfer_mode  = MSPI_PIO;
 	dev_data->xfer.packets    = &dev_data->packet;
 	dev_data->xfer.num_packet = 1;
 	dev_data->xfer.timeout    = 10;

 	dev_data->xfer.cmd_length  = 1;
 	dev_data->xfer.addr_length = 0;
 	dev_data->xfer.tx_dummy    = 0;
 	dev_data->xfer.rx_dummy    = 0;

 	dev_data->packet.dir       = MSPI_RX;
 	dev_data->packet.cmd       = JESD216_CMD_READ_ID;
 	dev_data->packet.data_buf  = id;
 	dev_data->packet.num_bytes = sizeof(id);
 	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 			     &dev_data->xfer);
 	if (rc < 0) {
 		LOG_ERR("Failed to read JEDEC ID in single line mode: %d", rc);
 		return rc;
 	}

 	/*
 	 * If the read ID does not match the one from DTS, assume the flash
 	 * is already in the Octa I/O mode, so switching it is not needed.
 	 */
 	if (memcmp(id, dev_config->jedec_id, sizeof(id)) == 0) {
 		static const uint8_t enable_sopi[] = { 0x01 };

 		dev_data->packet.dir       = MSPI_TX;
 		dev_data->packet.cmd       = SPI_NOR_CMD_WREN;
 		dev_data->packet.num_bytes = 0;
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			LOG_ERR("SPI_NOR_CMD_WREN xfer failed: %d", rc);
 			return rc;
 		}

 		dev_data->xfer.addr_length = 4;
 		dev_data->packet.cmd       = SPI_NOR_CMD_WR_CFGREG2;
 		dev_data->packet.address   = 0;
 		dev_data->packet.data_buf  = (uint8_t *)&enable_sopi;
 		dev_data->packet.num_bytes = sizeof(enable_sopi);
 		rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
 				     &dev_data->xfer);
 		if (rc < 0) {
 			printk("SPI_NOR_CMD_WR_CFGREG2 xfer failed: %d\n", rc);
 			return rc;
 		}
 	}

 	rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
 			     MSPI_DEVICE_CONFIG_ALL, &dev_config->mspi_cfg);
 	if (rc < 0) {
 		LOG_ERR("Failed to set device config: %d", rc);
 		return rc;
 	}

 	rc = read_jedec_id(dev, id);
 	if (rc < 0) {
 		return rc;
 	}

 	if (memcmp(id, dev_config->jedec_id, sizeof(id)) != 0) {
 		LOG_ERR("JEDEC ID mismatch, read: %02x %02x %02x, "
 			"expected: %02x %02x %02x",
 			id[0], id[1], id[2],
 			dev_config->jedec_id[0],
 			dev_config->jedec_id[1],
 			dev_config->jedec_id[2]);
 		return -ENODEV;
 	}

 #if defined(CONFIG_MSPI_XIP)
 	/* Enable XIP access for this chip if specified so in DT. */
 	if (dev_config->xip_cfg.enable) {
 		rc = mspi_xip_config(dev_config->bus, &dev_config->mspi_id,
 				     &dev_config->xip_cfg);
 		if (rc < 0) {
 			return rc;
 		}
 	}
 #endif

 	return 0;
 }

 static int drv_init(const struct device *dev)
 {
 	const struct flash_mspi_nor_config *dev_config = dev->config;
 	struct flash_mspi_nor_data *dev_data = dev->data;
 	int rc;

 	if (!device_is_ready(dev_config->bus)) {
 		LOG_ERR("Device %s is not ready", dev_config->bus->name);
 		return -ENODEV;
 	}

 #if defined(WITH_RESET_GPIO)
 	if (dev_config->reset.port) {
 		if (!gpio_is_ready_dt(&dev_config->reset)) {
 			LOG_ERR("Device %s is not ready",
 				dev_config->reset.port->name);
 			return -ENODEV;
 		}

 		rc = gpio_pin_configure_dt(&dev_config->reset,
 					   GPIO_OUTPUT_ACTIVE);
 		if (rc < 0) {
 			LOG_ERR("Failed to activate RESET: %d", rc);
 			return -EIO;
 		}

 		if (dev_config->reset_pulse_us != 0) {
 			k_busy_wait(dev_config->reset_pulse_us);
 		}

 		rc = gpio_pin_set_dt(&dev_config->reset, 0);
 		if (rc < 0) {
 			LOG_ERR("Failed to deactivate RESET: %d", rc);
 			return -EIO;
 		}

 		if (dev_config->reset_recovery_us != 0) {
 			k_busy_wait(dev_config->reset_recovery_us);
 		}
 	}
 #endif

 	rc = pm_device_runtime_get(dev_config->bus);
 	if (rc < 0) {
 		LOG_ERR("pm_device_runtime_get() failed: %d", rc);
 		return rc;
 	}

 	rc = flash_chip_init(dev);

 	/* Release the MSPI controller - it was acquired by the call to
 	 * mspi_dev_config() in flash_chip_init().
 	 */
 	(void)mspi_get_channel_status(dev_config->bus, 0);

 	(void)pm_device_runtime_put(dev_config->bus);

 	if (rc < 0) {
 		return rc;
 	}

 	k_sem_init(&dev_data->acquired, 1, K_SEM_MAX_LIMIT);

 	return pm_device_driver_init(dev, dev_pm_action_cb);
 }

 static DEVICE_API(flash, drv_api) = {
 	.read = api_read,
 	.write = api_write,
 	.erase = api_erase,
 	.get_parameters = api_get_parameters,
 #if defined(CONFIG_FLASH_PAGE_LAYOUT)
 	.page_layout = api_page_layout,
 #endif
 #if defined(CONFIG_FLASH_JESD216_API)
 	.sfdp_read = api_sfdp_read,
 	.read_jedec_id = api_read_jedec_id,
 #endif
 };

 #define FLASH_SIZE_INST(inst) (DT_INST_PROP(inst, size) / 8)

 #if defined(CONFIG_FLASH_PAGE_LAYOUT)
 BUILD_ASSERT((CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE % 4096) == 0,
 	"MSPI_NOR_FLASH_LAYOUT_PAGE_SIZE must be multiple of 4096");
 #define FLASH_PAGE_LAYOUT_DEFINE(inst) \
 	.layout = { \
 		.pages_size = CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE, \
 		.pages_count = FLASH_SIZE_INST(inst) \
 			     / CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE, \
 	},
 #define FLASH_PAGE_LAYOUT_CHECK(inst) \
 BUILD_ASSERT((FLASH_SIZE_INST(inst) % CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE) == 0, \
 	"MSPI_NOR_FLASH_LAYOUT_PAGE_SIZE incompatible with flash size, instance " #inst);
 #else
 #define FLASH_PAGE_LAYOUT_DEFINE(inst)
 #define FLASH_PAGE_LAYOUT_CHECK(inst)
 #endif

 /* MSPI bus must be initialized before this device. */
 #if (CONFIG_MSPI_INIT_PRIORITY < CONFIG_FLASH_INIT_PRIORITY)
 #define INIT_PRIORITY CONFIG_FLASH_INIT_PRIORITY
 #else
 #define INIT_PRIORITY UTIL_INC(CONFIG_MSPI_INIT_PRIORITY)
 #endif

 #define FLASH_MSPI_NOR_INST(inst)						\
 	BUILD_ASSERT(DT_INST_ENUM_IDX(inst, mspi_io_mode) ==			\
 		     MSPI_IO_MODE_OCTAL,					\
 		"Only Octal I/O mode is supported for now");			\
 	PM_DEVICE_DT_INST_DEFINE(inst, dev_pm_action_cb);			\
 	static struct flash_mspi_nor_data dev##inst##_data;			\
 	static const struct flash_mspi_nor_config dev##inst##_config = {	\
 		.bus = DEVICE_DT_GET(DT_INST_BUS(inst)),			\
 		.flash_size = FLASH_SIZE_INST(inst),				\
 		.mspi_id = MSPI_DEVICE_ID_DT_INST(inst),			\
 		.mspi_cfg = MSPI_DEVICE_CONFIG_DT_INST(inst),			\
 		.mspi_cfg_mask = DT_PROP(DT_INST_BUS(inst),			\
 					 software_multiperipheral)		\
 			       ? MSPI_DEVICE_CONFIG_ALL				\
 			       : MSPI_DEVICE_CONFIG_NONE,			\
 	IF_ENABLED(CONFIG_MSPI_XIP,						\
 		(.xip_cfg = MSPI_XIP_CONFIG_DT_INST(inst),))			\
 	IF_ENABLED(WITH_RESET_GPIO,						\
 		(.reset = GPIO_DT_SPEC_INST_GET_OR(inst, reset_gpios, {0}),	\
 		.reset_pulse_us = DT_INST_PROP_OR(inst, t_reset_pulse, 0)	\
 				/ 1000,						\
 		.reset_recovery_us = DT_INST_PROP_OR(inst, t_reset_recovery, 0)	\
 				   / 1000,))					\
 		FLASH_PAGE_LAYOUT_DEFINE(inst)					\
 		.jedec_id = DT_INST_PROP(inst, jedec_id),			\
 	};									\
 	FLASH_PAGE_LAYOUT_CHECK(inst)						\
 	DEVICE_DT_INST_DEFINE(inst,						\
 		drv_init, PM_DEVICE_DT_INST_GET(inst),				\
 		&dev##inst##_data, &dev##inst##_config,				\
 		POST_KERNEL, INIT_PRIORITY,					\
 		&drv_api);

 DT_INST_FOREACH_STATUS_OKAY(FLASH_MSPI_NOR_INST)
	/*
	* Copyright (c) 2024 Nordic Semiconductor ASA
	*
	* SPDX-License-Identifier: Apache-2.0
	*/

	#define DT_DRV_COMPAT jedec_mspi_nor

	#include <zephyr/drivers/flash.h>
	#include <zephyr/drivers/gpio.h>
	#include <zephyr/drivers/mspi.h>
	#include <zephyr/logging/log.h>
	#include <zephyr/pm/device.h>
	#include <zephyr/pm/device_runtime.h>

	#include "jesd216.h"
	#include "spi_nor.h"

	LOG_MODULE_REGISTER(flash_mspi_nor, CONFIG_FLASH_LOG_LEVEL);

	#if DT_ANY_INST_HAS_PROP_STATUS_OKAY(reset_gpios)
	#define WITH_RESET_GPIO 1
	#endif

	struct flash_mspi_nor_data {
	struct k_sem acquired;
	struct mspi_xfer_packet packet;
	struct mspi_xfer xfer;
	};

	struct flash_mspi_nor_config {
	const struct device *bus;
	uint32_t flash_size;
	struct mspi_dev_id mspi_id;
	struct mspi_dev_cfg mspi_cfg;
	enum mspi_dev_cfg_mask mspi_cfg_mask;
	#if defined(CONFIG_MSPI_XIP)
	struct mspi_xip_cfg xip_cfg;
	#endif
	#if defined(WITH_RESET_GPIO)
	struct gpio_dt_spec reset;
	uint32_t reset_pulse_us;
	uint32_t reset_recovery_us;
	#endif
	#if defined(CONFIG_FLASH_PAGE_LAYOUT)
	struct flash_pages_layout layout;
	#endif
	uint8_t jedec_id[SPI_NOR_MAX_ID_LEN];
	};

	static int acquire(const struct device *dev)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	int rc;

	k_sem_take(&dev_data->acquired, K_FOREVER);

	rc = pm_device_runtime_get(dev_config->bus);
	if (rc < 0) {
	LOG_ERR("pm_device_runtime_get() failed: %d", rc);
	} else {
	/* This acquires the MSPI controller and reconfigures it
	* if needed for the flash device.
	*/
	rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
	dev_config->mspi_cfg_mask,
	&dev_config->mspi_cfg);
	if (rc < 0) {
	LOG_ERR("mspi_dev_config() failed: %d", rc);
	} else {
	return 0;
	}

	(void)pm_device_runtime_put(dev_config->bus);
	}

	k_sem_give(&dev_data->acquired);
	return rc;
	}

	static void release(const struct device *dev)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;

	/* This releases the MSPI controller. */
	(void)mspi_get_channel_status(dev_config->bus, 0);

	(void)pm_device_runtime_put(dev_config->bus);

	k_sem_give(&dev_data->acquired);
	}

	static inline uint32_t dev_flash_size(const struct device *dev)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;

	return dev_config->flash_size;
	}

	static inline uint16_t dev_page_size(const struct device *dev)
	{
	return SPI_NOR_PAGE_SIZE;
	}

	static int api_read(const struct device dev, off_t addr, void dest,
	size_t size)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	const uint32_t flash_size = dev_flash_size(dev);
	int rc;

	if (size == 0) {
	return 0;
	}

	if ((addr < 0) \|\| ((addr + size) > flash_size)) {
	return -EINVAL;
	}

	rc = acquire(dev);
	if (rc < 0) {
	return rc;
	}

	/* TODO: get rid of all these hard-coded values for MX25Ux chips */
	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.addr_length = 4;
	dev_data->xfer.rx_dummy = 20;
	dev_data->packet.dir = MSPI_RX;
	dev_data->packet.cmd = SPI_NOR_OCMD_RD;
	dev_data->packet.address = addr;
	dev_data->packet.data_buf = dest;
	dev_data->packet.num_bytes = size;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);

	release(dev);

	if (rc < 0) {
	LOG_ERR("SPI_NOR_OCMD_RD xfer failed: %d", rc);
	return rc;
	}

	return 0;
	}

	static int wait_until_ready(const struct device *dev, k_timeout_t poll_period)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	uint8_t status_reg;
	int rc;

	while (true) {
	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.addr_length = 4;
	dev_data->xfer.rx_dummy = 4;
	dev_data->packet.dir = MSPI_RX;
	dev_data->packet.cmd = SPI_NOR_OCMD_RDSR;
	dev_data->packet.address = 0;
	dev_data->packet.data_buf = &status_reg;
	dev_data->packet.num_bytes = sizeof(status_reg);
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("SPI_NOR_OCMD_RDSR xfer failed: %d", rc);
	return rc;
	}
	if (!(status_reg & SPI_NOR_WIP_BIT)) {
	break;
	}

	k_sleep(poll_period);
	}

	return 0;
	}

	static int api_write(const struct device dev, off_t addr, const void src,
	size_t size)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	const uint32_t flash_size = dev_flash_size(dev);
	const uint16_t page_size = dev_page_size(dev);
	int rc;

	if (size == 0) {
	return 0;
	}

	if ((addr < 0) \|\| ((addr + size) > flash_size)) {
	return -EINVAL;
	}

	rc = acquire(dev);
	if (rc < 0) {
	return rc;
	}

	while (size > 0) {
	/* Split write into parts, each within one page only. */
	uint16_t page_offset = (uint16_t)(addr % page_size);
	uint16_t page_left = page_size - page_offset;
	uint16_t to_write = (uint16_t)MIN(size, page_left);

	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.tx_dummy = 0;
	dev_data->packet.dir = MSPI_TX;

	dev_data->xfer.addr_length = 0;
	dev_data->packet.cmd = SPI_NOR_OCMD_WREN;
	dev_data->packet.num_bytes = 0;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("SPI_NOR_OCMD_WREN xfer failed: %d", rc);
	break;
	}

	dev_data->xfer.addr_length = 4;
	dev_data->packet.cmd = SPI_NOR_OCMD_PAGE_PRG;
	dev_data->packet.address = addr;
	dev_data->packet.data_buf = (uint8_t *)src;
	dev_data->packet.num_bytes = to_write;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("SPI_NOR_OCMD_PAGE_PRG xfer failed: %d", rc);
	break;
	}

	addr += to_write;
	src = (const uint8_t *)src + to_write;
	size -= to_write;

	rc = wait_until_ready(dev, K_MSEC(1));
	if (rc < 0) {
	break;
	}
	}

	release(dev);

	return rc;
	}

	static int api_erase(const struct device *dev, off_t addr, size_t size)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	const uint32_t flash_size = dev_flash_size(dev);
	int rc = 0;

	if ((addr < 0) \|\| ((addr + size) > flash_size)) {
	return -EINVAL;
	}

	if (!SPI_NOR_IS_SECTOR_ALIGNED(addr)) {
	return -EINVAL;
	}

	if ((size % SPI_NOR_SECTOR_SIZE) != 0) {
	return -EINVAL;
	}

	rc = acquire(dev);
	if (rc < 0) {
	return rc;
	}

	while (size > 0) {
	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.tx_dummy = 0;
	dev_data->packet.dir = MSPI_TX;
	dev_data->packet.num_bytes = 0;

	dev_data->xfer.addr_length = 0;
	dev_data->packet.cmd = SPI_NOR_OCMD_WREN;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("SPI_NOR_OCMD_WREN xfer failed: %d", rc);
	break;
	}

	if (size == flash_size) {
	/* Chip erase. */
	dev_data->xfer.addr_length = 0;
	dev_data->packet.cmd = SPI_NOR_OCMD_CE;

	size -= flash_size;
	} else {
	/* Sector erase. */
	dev_data->xfer.addr_length = 4;
	dev_data->packet.cmd = SPI_NOR_OCMD_SE;
	dev_data->packet.address = addr;

	addr += SPI_NOR_SECTOR_SIZE;
	size -= SPI_NOR_SECTOR_SIZE;
	}
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("Erase command 0x%02x xfer failed: %d",
	dev_data->packet.cmd, rc);
	break;
	}

	rc = wait_until_ready(dev, K_MSEC(1));
	if (rc < 0) {
	break;
	}
	}

	release(dev);

	return rc;
	}

	static const
	struct flash_parameters api_get_parameters(const struct device dev)
	{
	ARG_UNUSED(dev);

	static const struct flash_parameters parameters = {
	.write_block_size = 1,
	.erase_value = 0xff,
	};

	return &parameters;
	}

	static int read_jedec_id(const struct device dev, uint8_t id)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	int rc;

	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.addr_length = 4;
	dev_data->xfer.rx_dummy = 4;
	dev_data->packet.dir = MSPI_RX;
	dev_data->packet.cmd = JESD216_OCMD_READ_ID;
	dev_data->packet.address = 0;
	dev_data->packet.data_buf = id;
	dev_data->packet.num_bytes = JESD216_READ_ID_LEN;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	printk("mspi_transceive() failed: %d\n", rc);
	return rc;
	}

	return rc;
	}

	#if defined(CONFIG_FLASH_PAGE_LAYOUT)
	static void api_page_layout(const struct device *dev,
	const struct flash_pages_layout **layout,
	size_t *layout_size)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;

	*layout = &dev_config->layout;
	*layout_size = 1;
	}
	#endif /* CONFIG_FLASH_PAGE_LAYOUT */

	#if defined(CONFIG_FLASH_JESD216_API)
	static int api_sfdp_read(const struct device dev, off_t addr, void dest,
	size_t size)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	int rc;

	if (size == 0) {
	return 0;
	}

	rc = acquire(dev);
	if (rc < 0) {
	return rc;
	}

	dev_data->xfer.cmd_length = 2;
	dev_data->xfer.addr_length = 4;
	dev_data->xfer.rx_dummy = 20;
	dev_data->packet.dir = MSPI_RX;
	dev_data->packet.cmd = JESD216_OCMD_READ_SFDP;
	dev_data->packet.address = addr;
	dev_data->packet.data_buf = dest;
	dev_data->packet.num_bytes = size;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	printk("JESD216_OCMD_READ_SFDP xfer failed: %d\n", rc);
	return rc;
	}

	release(dev);

	return rc;
	}

	static int api_read_jedec_id(const struct device dev, uint8_t id)
	{
	int rc = 0;

	rc = acquire(dev);
	if (rc < 0) {
	return rc;
	}

	rc = read_jedec_id(dev, id);

	release(dev);

	return rc;
	}
	#endif /* CONFIG_FLASH_JESD216_API */

	static int dev_pm_action_cb(const struct device *dev,
	enum pm_device_action action)
	{
	switch (action) {
	case PM_DEVICE_ACTION_SUSPEND:
	break;
	case PM_DEVICE_ACTION_RESUME:
	break;
	default:
	return -ENOTSUP;
	}

	return 0;
	}

	static int flash_chip_init(const struct device *dev)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	struct mspi_dev_cfg init_dev_cfg = dev_config->mspi_cfg;
	uint8_t id[JESD216_READ_ID_LEN] = {0};
	int rc;

	init_dev_cfg.freq = MHZ(1);
	init_dev_cfg.io_mode = MSPI_IO_MODE_SINGLE;

	rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
	MSPI_DEVICE_CONFIG_ALL, &init_dev_cfg);
	if (rc < 0) {
	LOG_ERR("Failed to set initial device config: %d", rc);
	return rc;
	}

	dev_data->xfer.xfer_mode = MSPI_PIO;
	dev_data->xfer.packets = &dev_data->packet;
	dev_data->xfer.num_packet = 1;
	dev_data->xfer.timeout = 10;

	dev_data->xfer.cmd_length = 1;
	dev_data->xfer.addr_length = 0;
	dev_data->xfer.tx_dummy = 0;
	dev_data->xfer.rx_dummy = 0;

	dev_data->packet.dir = MSPI_RX;
	dev_data->packet.cmd = JESD216_CMD_READ_ID;
	dev_data->packet.data_buf = id;
	dev_data->packet.num_bytes = sizeof(id);
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("Failed to read JEDEC ID in single line mode: %d", rc);
	return rc;
	}

	/*
	* If the read ID does not match the one from DTS, assume the flash
	* is already in the Octa I/O mode, so switching it is not needed.
	*/
	if (memcmp(id, dev_config->jedec_id, sizeof(id)) == 0) {
	static const uint8_t enable_sopi[] = { 0x01 };

	dev_data->packet.dir = MSPI_TX;
	dev_data->packet.cmd = SPI_NOR_CMD_WREN;
	dev_data->packet.num_bytes = 0;
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	LOG_ERR("SPI_NOR_CMD_WREN xfer failed: %d", rc);
	return rc;
	}

	dev_data->xfer.addr_length = 4;
	dev_data->packet.cmd = SPI_NOR_CMD_WR_CFGREG2;
	dev_data->packet.address = 0;
	dev_data->packet.data_buf = (uint8_t *)&enable_sopi;
	dev_data->packet.num_bytes = sizeof(enable_sopi);
	rc = mspi_transceive(dev_config->bus, &dev_config->mspi_id,
	&dev_data->xfer);
	if (rc < 0) {
	printk("SPI_NOR_CMD_WR_CFGREG2 xfer failed: %d\n", rc);
	return rc;
	}
	}

	rc = mspi_dev_config(dev_config->bus, &dev_config->mspi_id,
	MSPI_DEVICE_CONFIG_ALL, &dev_config->mspi_cfg);
	if (rc < 0) {
	LOG_ERR("Failed to set device config: %d", rc);
	return rc;
	}

	rc = read_jedec_id(dev, id);
	if (rc < 0) {
	return rc;
	}

	if (memcmp(id, dev_config->jedec_id, sizeof(id)) != 0) {
	LOG_ERR("JEDEC ID mismatch, read: %02x %02x %02x, "
	"expected: %02x %02x %02x",
	id[0], id[1], id[2],
	dev_config->jedec_id[0],
	dev_config->jedec_id[1],
	dev_config->jedec_id[2]);
	return -ENODEV;
	}

	#if defined(CONFIG_MSPI_XIP)
	/* Enable XIP access for this chip if specified so in DT. */
	if (dev_config->xip_cfg.enable) {
	rc = mspi_xip_config(dev_config->bus, &dev_config->mspi_id,
	&dev_config->xip_cfg);
	if (rc < 0) {
	return rc;
	}
	}
	#endif

	return 0;
	}

	static int drv_init(const struct device *dev)
	{
	const struct flash_mspi_nor_config *dev_config = dev->config;
	struct flash_mspi_nor_data *dev_data = dev->data;
	int rc;

	if (!device_is_ready(dev_config->bus)) {
	LOG_ERR("Device %s is not ready", dev_config->bus->name);
	return -ENODEV;
	}

	#if defined(WITH_RESET_GPIO)
	if (dev_config->reset.port) {
	if (!gpio_is_ready_dt(&dev_config->reset)) {
	LOG_ERR("Device %s is not ready",
	dev_config->reset.port->name);
	return -ENODEV;
	}

	rc = gpio_pin_configure_dt(&dev_config->reset,
	GPIO_OUTPUT_ACTIVE);
	if (rc < 0) {
	LOG_ERR("Failed to activate RESET: %d", rc);
	return -EIO;
	}

	if (dev_config->reset_pulse_us != 0) {
	k_busy_wait(dev_config->reset_pulse_us);
	}

	rc = gpio_pin_set_dt(&dev_config->reset, 0);
	if (rc < 0) {
	LOG_ERR("Failed to deactivate RESET: %d", rc);
	return -EIO;
	}

	if (dev_config->reset_recovery_us != 0) {
	k_busy_wait(dev_config->reset_recovery_us);
	}
	}
	#endif

	rc = pm_device_runtime_get(dev_config->bus);
	if (rc < 0) {
	LOG_ERR("pm_device_runtime_get() failed: %d", rc);
	return rc;
	}

	rc = flash_chip_init(dev);

	/* Release the MSPI controller - it was acquired by the call to
	* mspi_dev_config() in flash_chip_init().
	*/
	(void)mspi_get_channel_status(dev_config->bus, 0);

	(void)pm_device_runtime_put(dev_config->bus);

	if (rc < 0) {
	return rc;
	}

	k_sem_init(&dev_data->acquired, 1, K_SEM_MAX_LIMIT);

	return pm_device_driver_init(dev, dev_pm_action_cb);
	}

	static DEVICE_API(flash, drv_api) = {
	.read = api_read,
	.write = api_write,
	.erase = api_erase,
	.get_parameters = api_get_parameters,
	#if defined(CONFIG_FLASH_PAGE_LAYOUT)
	.page_layout = api_page_layout,
	#endif
	#if defined(CONFIG_FLASH_JESD216_API)
	.sfdp_read = api_sfdp_read,
	.read_jedec_id = api_read_jedec_id,
	#endif
	};

	#define FLASH_SIZE_INST(inst) (DT_INST_PROP(inst, size) / 8)

	#if defined(CONFIG_FLASH_PAGE_LAYOUT)
	BUILD_ASSERT((CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE % 4096) == 0,
	"MSPI_NOR_FLASH_LAYOUT_PAGE_SIZE must be multiple of 4096");
	#define FLASH_PAGE_LAYOUT_DEFINE(inst) \
	.layout = { \
	.pages_size = CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE, \
	.pages_count = FLASH_SIZE_INST(inst) \
	/ CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE, \
	},
	#define FLASH_PAGE_LAYOUT_CHECK(inst) \
	BUILD_ASSERT((FLASH_SIZE_INST(inst) % CONFIG_FLASH_MSPI_NOR_LAYOUT_PAGE_SIZE) == 0, \
	"MSPI_NOR_FLASH_LAYOUT_PAGE_SIZE incompatible with flash size, instance " #inst);
	#else
	#define FLASH_PAGE_LAYOUT_DEFINE(inst)
	#define FLASH_PAGE_LAYOUT_CHECK(inst)
	#endif

	/* MSPI bus must be initialized before this device. */
	#if (CONFIG_MSPI_INIT_PRIORITY < CONFIG_FLASH_INIT_PRIORITY)
	#define INIT_PRIORITY CONFIG_FLASH_INIT_PRIORITY
	#else
	#define INIT_PRIORITY UTIL_INC(CONFIG_MSPI_INIT_PRIORITY)
	#endif

	#define FLASH_MSPI_NOR_INST(inst) \
	BUILD_ASSERT(DT_INST_ENUM_IDX(inst, mspi_io_mode) == \
	MSPI_IO_MODE_OCTAL, \
	"Only Octal I/O mode is supported for now"); \
	PM_DEVICE_DT_INST_DEFINE(inst, dev_pm_action_cb); \
	static struct flash_mspi_nor_data dev##inst##_data; \
	static const struct flash_mspi_nor_config dev##inst##_config = { \
	.bus = DEVICE_DT_GET(DT_INST_BUS(inst)), \
	.flash_size = FLASH_SIZE_INST(inst), \
	.mspi_id = MSPI_DEVICE_ID_DT_INST(inst), \
	.mspi_cfg = MSPI_DEVICE_CONFIG_DT_INST(inst), \
	.mspi_cfg_mask = DT_PROP(DT_INST_BUS(inst), \
	software_multiperipheral) \
	? MSPI_DEVICE_CONFIG_ALL \
	: MSPI_DEVICE_CONFIG_NONE, \
	IF_ENABLED(CONFIG_MSPI_XIP, \
	(.xip_cfg = MSPI_XIP_CONFIG_DT_INST(inst),)) \
	IF_ENABLED(WITH_RESET_GPIO, \
	(.reset = GPIO_DT_SPEC_INST_GET_OR(inst, reset_gpios, {0}), \
	.reset_pulse_us = DT_INST_PROP_OR(inst, t_reset_pulse, 0) \
	/ 1000, \
	.reset_recovery_us = DT_INST_PROP_OR(inst, t_reset_recovery, 0) \
	/ 1000,)) \
	FLASH_PAGE_LAYOUT_DEFINE(inst) \
	.jedec_id = DT_INST_PROP(inst, jedec_id), \
	}; \
	FLASH_PAGE_LAYOUT_CHECK(inst) \
	DEVICE_DT_INST_DEFINE(inst, \
	drv_init, PM_DEVICE_DT_INST_GET(inst), \
	&dev##inst##_data, &dev##inst##_config, \
	POST_KERNEL, INIT_PRIORITY, \
	&drv_api);

	DT_INST_FOREACH_STATUS_OKAY(FLASH_MSPI_NOR_INST)