1 files changed, 32 insertions, 23 deletions

diff --git a/drivers/mtd/mtdcore.c b/drivers/mtd/mtdcore.c
index 0b4ca0aa4132..0feacb9fbdac 100644
--- a/drivers/mtd/mtdcore.c
+++ b/drivers/mtd/mtdcore.c
@@ -28,6 +28,7 @@
 #include <linux/leds.h>
 #include <linux/debugfs.h>
 #include <linux/nvmem-provider.h>
+#include <linux/root_dev.h>
 
 #include <linux/mtd/mtd.h>
 #include <linux/mtd/partitions.h>
@@ -551,22 +552,22 @@ static void mtd_check_of_node(struct mtd_info *mtd)
 	struct device_node *partitions, *parent_dn, *mtd_dn = NULL;
 	const char *pname, *prefix = "partition-";
 	int plen, mtd_name_len, offset, prefix_len;
-	struct mtd_info *parent;
-	bool found = false;
 
 	/* Check if MTD already has a device node */
-	if (dev_of_node(&mtd->dev))
+	if (mtd_get_of_node(mtd))
 		return;
 
-	/* Check if a partitions node exist */
 	if (!mtd_is_partition(mtd))
 		return;
-	parent = mtd->parent;
-	parent_dn = of_node_get(dev_of_node(&parent->dev));
+
+	parent_dn = of_node_get(mtd_get_of_node(mtd->parent));
 	if (!parent_dn)
 		return;
 
-	partitions = of_get_child_by_name(parent_dn, "partitions");
+	if (mtd_is_partition(mtd->parent))
+		partitions = of_node_get(parent_dn);
+	else
+		partitions = of_get_child_by_name(parent_dn, "partitions");
 	if (!partitions)
 		goto exit_parent;
 
@@ -575,34 +576,26 @@ static void mtd_check_of_node(struct mtd_info *mtd)
 
 	/* Search if a partition is defined with the same name */
 	for_each_child_of_node(partitions, mtd_dn) {
-		offset = 0;
-
 		/* Skip partition with no/wrong prefix */
-		if (!of_node_name_prefix(mtd_dn, "partition-"))
+		if (!of_node_name_prefix(mtd_dn, prefix))
 			continue;
 
 		/* Label have priority. Check that first */
-		if (of_property_read_string(mtd_dn, "label", &pname)) {
-			of_property_read_string(mtd_dn, "name", &pname);
+		if (!of_property_read_string(mtd_dn, "label", &pname)) {
+			offset = 0;
+		} else {
+			pname = mtd_dn->name;
 			offset = prefix_len;
 		}
 
 		plen = strlen(pname) - offset;
 		if (plen == mtd_name_len &&
 		    !strncmp(mtd->name, pname + offset, plen)) {
-			found = true;
+			mtd_set_of_node(mtd, mtd_dn);
 			break;
 		}
 	}
 
-	if (!found)
-		goto exit_partitions;
-
-	/* Set of_node only for nvmem */
-	if (of_device_is_compatible(mtd_dn, "nvmem-cells"))
-		mtd_set_of_node(mtd, mtd_dn);
-
-exit_partitions:
 	of_node_put(partitions);
 exit_parent:
 	of_node_put(parent_dn);
@@ -723,8 +716,10 @@ int add_mtd_device(struct mtd_info *mtd)
 	mtd_check_of_node(mtd);
 	of_node_get(mtd_get_of_node(mtd));
 	error = device_register(&mtd->dev);
-	if (error)
+	if (error) {
+		put_device(&mtd->dev);
 		goto fail_added;
+	}
 
 	/* Add the nvmem provider */
 	error = mtd_nvmem_add(mtd);
@@ -743,6 +738,17 @@ int add_mtd_device(struct mtd_info *mtd)
 		not->add(mtd);
 
 	mutex_unlock(&mtd_table_mutex);
+
+	if (of_find_property(mtd_get_of_node(mtd), "linux,rootfs", NULL)) {
+		if (IS_BUILTIN(CONFIG_MTD)) {
+			pr_info("mtd: setting mtd%d (%s) as root device\n", mtd->index, mtd->name);
+			ROOT_DEV = MKDEV(MTD_BLOCK_MAJOR, mtd->index);
+		} else {
+			pr_warn("mtd: can't set mtd%d (%s) as root device - mtd must be builtin\n",
+				mtd->index, mtd->name);
+		}
+	}
+
 	/* We _know_ we aren't being removed, because
 	   our caller is still holding us here. So none
 	   of this try_ nonsense, and no bitching about it
@@ -774,6 +780,7 @@ int del_mtd_device(struct mtd_info *mtd)
 {
 	int ret;
 	struct mtd_notifier *not;
+	struct device_node *mtd_of_node;
 
 	mutex_lock(&mtd_table_mutex);
 
@@ -792,6 +799,7 @@ int del_mtd_device(struct mtd_info *mtd)
 		       mtd->index, mtd->name, mtd->usecount);
 		ret = -EBUSY;
 	} else {
+		mtd_of_node = mtd_get_of_node(mtd);
 		debugfs_remove_recursive(mtd->dbg.dfs_dir);
 
 		/* Try to remove the NVMEM provider */
@@ -803,7 +811,7 @@ int del_mtd_device(struct mtd_info *mtd)
 		memset(&mtd->dev, 0, sizeof(mtd->dev));
 
 		idr_remove(&mtd_idr, mtd->index);
-		of_node_put(mtd_get_of_node(mtd));
+		of_node_put(mtd_of_node);
 
 		module_put(THIS_MODULE);
 		ret = 0;
@@ -2483,6 +2491,7 @@ static int __init init_mtd(void)
 out_procfs:
 	if (proc_mtd)
 		remove_proc_entry("mtd", NULL);
+	bdi_unregister(mtd_bdi);
 	bdi_put(mtd_bdi);
 err_bdi:
 	class_unregister(&mtd_class);