rb_insert_color(&e->u.rb, root);
 }
 
+/**
+ * wl_tree_destroy - destroy a wear-leveling entry.
+ * @ubi: UBI device description object
+ * @e: the wear-leveling entry to add
+ *
+ * This function destroys a wear leveling entry and removes
+ * the reference from the lookup table.
+ */
+static void wl_entry_destroy(struct ubi_device *ubi, struct ubi_wl_entry *e)
+{
+       ubi->lookuptbl[e->pnum] = NULL;
+       kmem_cache_free(ubi_wl_entry_slab, e);
+}
+
 /**
  * do_work - do one pending work.
  * @ubi: UBI device description object
        err = do_sync_erase(ubi, e1, vol_id, lnum, 0);
        if (err) {
                if (e2)
-                       kmem_cache_free(ubi_wl_entry_slab, e2);
+                       wl_entry_destroy(ubi, e2);
                goto out_ro;
        }
 
        spin_unlock(&ubi->wl_lock);
 
        ubi_free_vid_hdr(ubi, vid_hdr);
-       kmem_cache_free(ubi_wl_entry_slab, e1);
-       kmem_cache_free(ubi_wl_entry_slab, e2);
+       wl_entry_destroy(ubi, e1);
+       wl_entry_destroy(ubi, e2);
 
 out_ro:
        ubi_ro_mode(ubi);
        if (shutdown) {
                dbg_wl("cancel erasure of PEB %d EC %d", pnum, e->ec);
                kfree(wl_wrk);
-               kmem_cache_free(ubi_wl_entry_slab, e);
+               wl_entry_destroy(ubi, e);
                return 0;
        }
 
                return err;
        }
 
-       kmem_cache_free(ubi_wl_entry_slab, e);
+       wl_entry_destroy(ubi, e);
        if (err != -EIO)
                /*
                 * If this is not %-EIO, we have no idea what to do. Scheduling
 
 /**
  * tree_destroy - destroy an RB-tree.
+ * @ubi: UBI device description object
  * @root: the root of the tree to destroy
  */
-static void tree_destroy(struct rb_root *root)
+static void tree_destroy(struct ubi_device *ubi, struct rb_root *root)
 {
        struct rb_node *rb;
        struct ubi_wl_entry *e;
                                        rb->rb_right = NULL;
                        }
 
-                       kmem_cache_free(ubi_wl_entry_slab, e);
+                       wl_entry_destroy(ubi, e);
                }
        }
 }
                ubi_assert(!ubi_is_fm_block(ubi, e->pnum));
                ubi->lookuptbl[e->pnum] = e;
                if (schedule_erase(ubi, e, aeb->vol_id, aeb->lnum, 0)) {
-                       kmem_cache_free(ubi_wl_entry_slab, e);
+                       wl_entry_destroy(ubi, e);
                        goto out_free;
                }
 
 
 out_free:
        shutdown_work(ubi);
-       tree_destroy(&ubi->used);
-       tree_destroy(&ubi->free);
-       tree_destroy(&ubi->scrub);
+       tree_destroy(ubi, &ubi->used);
+       tree_destroy(ubi, &ubi->free);
+       tree_destroy(ubi, &ubi->scrub);
        kfree(ubi->lookuptbl);
        return err;
 }
        for (i = 0; i < UBI_PROT_QUEUE_LEN; ++i) {
                list_for_each_entry_safe(e, tmp, &ubi->pq[i], u.list) {
                        list_del(&e->u.list);
-                       kmem_cache_free(ubi_wl_entry_slab, e);
+                       wl_entry_destroy(ubi, e);
                }
        }
 }
        ubi_fastmap_close(ubi);
        shutdown_work(ubi);
        protection_queue_destroy(ubi);
-       tree_destroy(&ubi->used);
-       tree_destroy(&ubi->erroneous);
-       tree_destroy(&ubi->free);
-       tree_destroy(&ubi->scrub);
+       tree_destroy(ubi, &ubi->used);
+       tree_destroy(ubi, &ubi->erroneous);
+       tree_destroy(ubi, &ubi->free);
+       tree_destroy(ubi, &ubi->scrub);
        kfree(ubi->lookuptbl);
 }