#include <linux/module.h>
 #include <linux/slab.h>
 #include <linux/bitops.h>
-#include <linux/mtd/mtd.h>
-#include <linux/mtd/rawnand.h>
 #include <linux/mtd/nand.h>
 #include <linux/mtd/nand-ecc-sw-bch.h>
-#include <linux/bch.h>
-
-/**
- * struct nand_bch_control - private NAND BCH control structure
- * @bch:       BCH control structure
- * @errloc:    error location array
- * @eccmask:   XOR ecc mask, allows erased pages to be decoded as valid
- */
-struct nand_bch_control {
-       struct bch_control   *bch;
-       unsigned int         *errloc;
-       unsigned char        *eccmask;
-};
 
 /**
  * nand_ecc_sw_bch_calculate - Calculate the ECC corresponding to a data block
 int nand_ecc_sw_bch_calculate(struct nand_device *nand,
                              const unsigned char *buf, unsigned char *code)
 {
-       struct nand_chip *chip = mtd_to_nand(nanddev_to_mtd(nand));
-       struct nand_bch_control *nbc = chip->ecc.priv;
+       struct nand_ecc_sw_bch_conf *engine_conf = nand->ecc.ctx.priv;
        unsigned int i;
 
-       memset(code, 0, chip->ecc.bytes);
-       bch_encode(nbc->bch, buf, chip->ecc.size, code);
+       memset(code, 0, engine_conf->code_size);
+       bch_encode(engine_conf->bch, buf, nand->ecc.ctx.conf.step_size, code);
 
        /* apply mask so that an erased page is a valid codeword */
-       for (i = 0; i < chip->ecc.bytes; i++)
-               code[i] ^= nbc->eccmask[i];
+       for (i = 0; i < engine_conf->code_size; i++)
+               code[i] ^= engine_conf->eccmask[i];
 
        return 0;
 }
 int nand_ecc_sw_bch_correct(struct nand_device *nand, unsigned char *buf,
                            unsigned char *read_ecc, unsigned char *calc_ecc)
 {
-       struct nand_chip *chip = mtd_to_nand(nanddev_to_mtd(nand));
-       struct nand_bch_control *nbc = chip->ecc.priv;
-       unsigned int *errloc = nbc->errloc;
+       struct nand_ecc_sw_bch_conf *engine_conf = nand->ecc.ctx.priv;
+       unsigned int step_size = nand->ecc.ctx.conf.step_size;
+       unsigned int *errloc = engine_conf->errloc;
        int i, count;
 
-       count = bch_decode(nbc->bch, NULL, chip->ecc.size, read_ecc, calc_ecc,
-                          NULL, errloc);
+       count = bch_decode(engine_conf->bch, NULL, step_size, read_ecc,
+                          calc_ecc, NULL, errloc);
        if (count > 0) {
                for (i = 0; i < count; i++) {
-                       if (errloc[i] < (chip->ecc.size * 8))
+                       if (errloc[i] < (step_size * 8))
                                /* The error is in the data area: correct it */
                                buf[errloc[i] >> 3] ^= (1 << (errloc[i] & 7));
 
  *
  * Returns: a pointer to a new NAND BCH control structure, or NULL upon failure
  *
- * Initialize NAND BCH error correction. Parameters @eccsize and @eccbytes
- * are used to compute the following BCH parameters:
+ * Initialize NAND BCH error correction. @nand.ecc parameters 'step_size' and
+ * 'bytes' are used to compute the following BCH parameters:
  *     m, the Galois field order
  *     t, the error correction capability
- * @eccbytes should be equal to the number of bytes required to store m * t
+ * 'bytes' should be equal to the number of bytes required to store m * t
  * bits, where m is such that 2^m - 1 > step_size * 8.
  *
  * Example: to configure 4 bit correction per 512 bytes, you should pass
- * @eccsize = 512 (thus, m = 13 is the smallest integer such that 2^m - 1 > 512 * 8)
- * @eccbytes = 7 (7 bytes are required to store m * t = 13 * 4 = 52 bits)
+ * step_size = 512 (thus, m = 13 is the smallest integer such that 2^m - 1 > 512 * 8)
+ * bytes = 7 (7 bytes are required to store m * t = 13 * 4 = 52 bits)
  */
 int nand_ecc_sw_bch_init(struct nand_device *nand)
 {
        struct mtd_info *mtd = nanddev_to_mtd(nand);
-       struct nand_chip *chip = mtd_to_nand(mtd);
        unsigned int m, t, eccsteps, i;
-       struct nand_bch_control *nbc = NULL;
+       struct nand_ecc_sw_bch_conf *engine_conf = nand->ecc.ctx.priv;
        unsigned char *erased_page;
-       unsigned int eccsize = chip->ecc.size;
-       unsigned int eccbytes = chip->ecc.bytes;
-       unsigned int eccstrength = chip->ecc.strength;
+       unsigned int eccsize = nand->ecc.ctx.conf.step_size;
+       unsigned int eccbytes = engine_conf->code_size;
+       unsigned int eccstrength = nand->ecc.ctx.conf.strength;
 
        if (!eccbytes && eccstrength) {
                eccbytes = DIV_ROUND_UP(eccstrength * fls(8 * eccsize), 8);
-               chip->ecc.bytes = eccbytes;
+               engine_conf->code_size = eccbytes;
        }
 
        if (!eccsize || !eccbytes) {
        m = fls(1+8*eccsize);
        t = (eccbytes*8)/m;
 
-       nbc = kzalloc(sizeof(*nbc), GFP_KERNEL);
-       if (!nbc)
-               return -ENOMEM;
-
-       chip->ecc.priv = nbc;
-
-       nbc->bch = bch_init(m, t, 0, false);
-       if (!nbc->bch)
-               goto fail;
+       engine_conf->bch = bch_init(m, t, 0, false);
+       if (!engine_conf->bch)
+               return -EINVAL;
 
        /* verify that eccbytes has the expected value */
-       if (nbc->bch->ecc_bytes != eccbytes) {
+       if (engine_conf->bch->ecc_bytes != eccbytes) {
                pr_warn("invalid eccbytes %u, should be %u\n",
-                       eccbytes, nbc->bch->ecc_bytes);
+                       eccbytes, engine_conf->bch->ecc_bytes);
                goto fail;
        }
 
                goto fail;
        }
 
-       /*
-        * ecc->steps and ecc->total might be used by mtd->ooblayout->ecc(),
-        * which is called by mtd_ooblayout_count_eccbytes().
-        * Make sure they are properly initialized before calling
-        * mtd_ooblayout_count_eccbytes().
-        * FIXME: we should probably rework the sequencing in nand_scan_tail()
-        * to avoid setting those fields twice.
-        */
-       chip->ecc.steps = eccsteps;
-       chip->ecc.total = eccsteps * eccbytes;
-       nand->base.ecc.ctx.total = chip->ecc.total;
        if (mtd_ooblayout_count_eccbytes(mtd) != (eccsteps*eccbytes)) {
                pr_warn("invalid ecc layout\n");
                goto fail;
        }
 
-       nbc->eccmask = kzalloc(eccbytes, GFP_KERNEL);
-       nbc->errloc = kmalloc_array(t, sizeof(*nbc->errloc), GFP_KERNEL);
-       if (!nbc->eccmask || !nbc->errloc)
+       engine_conf->eccmask = kzalloc(eccbytes, GFP_KERNEL);
+       engine_conf->errloc = kmalloc_array(t, sizeof(*engine_conf->errloc),
+                                           GFP_KERNEL);
+       if (!engine_conf->eccmask || !engine_conf->errloc)
                goto fail;
 
        /*
                goto fail;
 
        memset(erased_page, 0xff, eccsize);
-       bch_encode(nbc->bch, erased_page, eccsize, nbc->eccmask);
+       bch_encode(engine_conf->bch, erased_page, eccsize,
+                  engine_conf->eccmask);
        kfree(erased_page);
 
        for (i = 0; i < eccbytes; i++)
-               nbc->eccmask[i] ^= 0xff;
+               engine_conf->eccmask[i] ^= 0xff;
 
        if (!eccstrength)
-               chip->ecc.strength = (eccbytes * 8) / fls(8 * eccsize);
+               nand->ecc.ctx.conf.strength = (eccbytes * 8) / fls(8 * eccsize);
 
        return 0;
 
  */
 void nand_ecc_sw_bch_cleanup(struct nand_device *nand)
 {
-       struct nand_chip *chip = mtd_to_nand(nanddev_to_mtd(nand));
-       struct nand_bch_control *nbc = chip->ecc.priv;
-
-       if (nbc) {
-               bch_free(nbc->bch);
-               kfree(nbc->errloc);
-               kfree(nbc->eccmask);
-               kfree(nbc);
+       struct nand_ecc_sw_bch_conf *engine_conf = nand->ecc.ctx.priv;
+
+       if (engine_conf) {
+               bch_free(engine_conf->bch);
+               kfree(engine_conf->errloc);
+               kfree(engine_conf->eccmask);
        }
 }
 EXPORT_SYMBOL(nand_ecc_sw_bch_cleanup);
 
 int rawnand_sw_bch_init(struct nand_chip *chip)
 {
        struct nand_device *base = &chip->base;
+       struct nand_ecc_sw_bch_conf *engine_conf;
+       int ret;
+
+       base->ecc.user_conf.engine_type = NAND_ECC_ENGINE_TYPE_SOFT;
+       base->ecc.user_conf.algo = NAND_ECC_ALGO_BCH;
+       base->ecc.user_conf.step_size = chip->ecc.size;
+       base->ecc.user_conf.strength = chip->ecc.strength;
+
+       engine_conf = kzalloc(sizeof(*engine_conf), GFP_KERNEL);
+       if (!engine_conf)
+               return -ENOMEM;
+
+       engine_conf->code_size = chip->ecc.bytes;
+
+       base->ecc.ctx.priv = engine_conf;
 
-       return nand_ecc_sw_bch_init(base);
+       ret = nand_ecc_sw_bch_init(base);
+       if (ret)
+               kfree(base->ecc.ctx.priv);
+
+       chip->ecc.size = base->ecc.ctx.conf.step_size;
+       chip->ecc.strength = base->ecc.ctx.conf.strength;
+       chip->ecc.total = base->ecc.ctx.total;
+       chip->ecc.steps = engine_conf->nsteps;
+       chip->ecc.bytes = engine_conf->code_size;
+
+       return ret;
 }
 EXPORT_SYMBOL(rawnand_sw_bch_init);
 
 
        nand_ecc_sw_bch_cleanup(base);
 
-       chip->ecc.priv = NULL;
+       kfree(base->ecc.ctx.priv);
 }
 EXPORT_SYMBOL(rawnand_sw_bch_cleanup);
 
         * Set the number of read / write steps for one page depending on ECC
         * mode.
         */
-       ecc->steps = mtd->writesize / ecc->size;
+       if (!ecc->steps)
+               ecc->steps = mtd->writesize / ecc->size;
        if (ecc->steps * ecc->size != mtd->writesize) {
                WARN(1, "Invalid ECC parameters\n");
                ret = -EINVAL;
                goto err_nand_manuf_cleanup;
        }
 
-       ecc->total = ecc->steps * ecc->bytes;
-       chip->base.ecc.ctx.total = ecc->total;
+       if (!ecc->total) {
+               ecc->total = ecc->steps * ecc->bytes;
+               chip->base.ecc.ctx.total = ecc->total;
+       }
 
        if (ecc->total > mtd->oobsize) {
                WARN(1, "Total number of ECC bytes exceeded oobsize\n");