1 files changed, 16 insertions, 3 deletions

diff --git a/drivers/mtd/nand/plat_nand.c b/drivers/mtd/nand/plat_nand.c
index 86e1d08eee0..4e16c6f5bdd 100644
--- a/drivers/mtd/nand/plat_nand.c
+++ b/drivers/mtd/nand/plat_nand.c
@@ -61,6 +61,8 @@ static int __devinit plat_nand_probe(struct platform_device *pdev)
         data->chip.cmd_ctrl = pdata->ctrl.cmd_ctrl;
         data->chip.dev_ready = pdata->ctrl.dev_ready;
         data->chip.select_chip = pdata->ctrl.select_chip;
+        data->chip.write_buf = pdata->ctrl.write_buf;
+        data->chip.read_buf = pdata->ctrl.read_buf;
         data->chip.chip_delay = pdata->chip.chip_delay;
         data->chip.options |= pdata->chip.options;
 
@@ -70,6 +72,13 @@ static int __devinit plat_nand_probe(struct platform_device *pdev)
 
         platform_set_drvdata(pdev, data);
 
+        /* Handle any platform specific setup */
+        if (pdata->ctrl.probe) {
+                res = pdata->ctrl.probe(pdev);
+                if (res)
+                        goto out;
+        }
+
         /* Scan to find existance of the device */
         if (nand_scan(&data->mtd, 1)) {
                 res = -ENXIO;
@@ -86,6 +95,8 @@ static int __devinit plat_nand_probe(struct platform_device *pdev)
                         return 0;
                 }
         }
+        if (pdata->chip.set_parts)
+                pdata->chip.set_parts(data->mtd.size, &pdata->chip);
         if (pdata->chip.partitions) {
                 data->parts = pdata->chip.partitions;
                 res = add_mtd_partitions(&data->mtd, data->parts,
@@ -99,6 +110,8 @@ static int __devinit plat_nand_probe(struct platform_device *pdev)
 
         nand_release(&data->mtd);
 out:
+        if (pdata->ctrl.remove)
+                pdata->ctrl.remove(pdev);
         platform_set_drvdata(pdev, NULL);
         iounmap(data->io_base);
         kfree(data);
@@ -111,15 +124,15 @@ out:
 static int __devexit plat_nand_remove(struct platform_device *pdev)
 {
         struct plat_nand_data *data = platform_get_drvdata(pdev);
-#ifdef CONFIG_MTD_PARTITIONS
         struct platform_nand_data *pdata = pdev->dev.platform_data;
-#endif
 
         nand_release(&data->mtd);
 #ifdef CONFIG_MTD_PARTITIONS
         if (data->parts && data->parts != pdata->chip.partitions)
                 kfree(data->parts);
 #endif
+        if (pdata->ctrl.remove)
+                pdata->ctrl.remove(pdev);
         iounmap(data->io_base);
         kfree(data);
 
@@ -128,7 +141,7 @@ static int __devexit plat_nand_remove(struct platform_device *pdev)
 
 static struct platform_driver plat_nand_driver = {
         .probe          = plat_nand_probe,
-        .remove         = plat_nand_remove,
+        .remove         = __devexit_p(plat_nand_remove),
         .driver         = {
                 .name   = "gen_nand",
                 .owner  = THIS_MODULE,