/**
  * struct pending_free - information about data structures about to be freed
  * @zapped: Head of a list with struct lock_class elements.
+ * @lock_chains_being_freed: Bitmap that indicates which lock_chains[] elements
+ *     are about to be freed.
  */
 struct pending_free {
        struct list_head zapped;
+       DECLARE_BITMAP(lock_chains_being_freed, MAX_LOCKDEP_CHAINS);
 };
 
 /**
        return 0;
 }
 
-static unsigned long nr_lock_chains;
 struct lock_chain lock_chains[MAX_LOCKDEP_CHAINS];
+static DECLARE_BITMAP(lock_chains_in_use, MAX_LOCKDEP_CHAINS);
 int nr_chain_hlocks;
 static u16 chain_hlocks[MAX_LOCKDEP_CHAIN_HLOCKS];
 
  */
 long lockdep_next_lockchain(long i)
 {
-       return i + 1 < nr_lock_chains ? i + 1 : -2;
+       i = find_next_bit(lock_chains_in_use, ARRAY_SIZE(lock_chains), i + 1);
+       return i < ARRAY_SIZE(lock_chains) ? i : -2;
 }
 
 unsigned long lock_chain_count(void)
 {
-       return nr_lock_chains;
+       return bitmap_weight(lock_chains_in_use, ARRAY_SIZE(lock_chains));
+}
+
+/* Must be called with the graph lock held. */
+static struct lock_chain *alloc_lock_chain(void)
+{
+       int idx = find_first_zero_bit(lock_chains_in_use,
+                                     ARRAY_SIZE(lock_chains));
+
+       if (unlikely(idx >= ARRAY_SIZE(lock_chains)))
+               return NULL;
+       __set_bit(idx, lock_chains_in_use);
+       return lock_chains + idx;
 }
 
 /*
        struct lock_chain *chain;
        int i, j;
 
-       /*
-        * Allocate a new chain entry from the static array, and add
-        * it to the hash:
-        */
-
        /*
         * The caller must hold the graph lock, ensure we've got IRQs
         * disabled to make this an IRQ-safe lock.. for recursion reasons
        if (DEBUG_LOCKS_WARN_ON(!irqs_disabled()))
                return 0;
 
-       if (unlikely(nr_lock_chains >= MAX_LOCKDEP_CHAINS)) {
+       chain = alloc_lock_chain();
+       if (!chain) {
                if (!debug_locks_off_graph_unlock())
                        return 0;
 
                dump_stack();
                return 0;
        }
-       chain = lock_chains + nr_lock_chains++;
        chain->chain_key = chain_key;
        chain->irq_context = hlock->irq_context;
        i = get_first_held_lock(curr, hlock);
 }
 
 /* Remove a class from a lock chain. Must be called with the graph lock held. */
-static void remove_class_from_lock_chain(struct lock_chain *chain,
+static void remove_class_from_lock_chain(struct pending_free *pf,
+                                        struct lock_chain *chain,
                                         struct lock_class *class)
 {
 #ifdef CONFIG_PROVE_LOCKING
         * hlist_for_each_entry_rcu() loop is safe.
         */
        hlist_del_rcu(&chain->entry);
+       __set_bit(chain - lock_chains, pf->lock_chains_being_freed);
        if (chain->depth == 0)
                return;
        /*
         */
        if (lookup_chain_cache(chain_key))
                return;
-       if (WARN_ON_ONCE(nr_lock_chains >= MAX_LOCKDEP_CHAINS)) {
+       new_chain = alloc_lock_chain();
+       if (WARN_ON_ONCE(!new_chain)) {
                debug_locks_off();
                return;
        }
-       /*
-        * Leak *chain because it is not safe to reinsert it before an RCU
-        * grace period has expired.
-        */
-       new_chain = lock_chains + nr_lock_chains++;
        *new_chain = *chain;
        hlist_add_head_rcu(&new_chain->entry, chainhashentry(chain_key));
 #endif
 }
 
 /* Must be called with the graph lock held. */
-static void remove_class_from_lock_chains(struct lock_class *class)
+static void remove_class_from_lock_chains(struct pending_free *pf,
+                                         struct lock_class *class)
 {
        struct lock_chain *chain;
        struct hlist_head *head;
        for (i = 0; i < ARRAY_SIZE(chainhash_table); i++) {
                head = chainhash_table + i;
                hlist_for_each_entry_rcu(chain, head, entry) {
-                       remove_class_from_lock_chain(chain, class);
+                       remove_class_from_lock_chain(pf, chain, class);
                }
        }
 }
                          class->name);
        }
 
-       remove_class_from_lock_chains(class);
+       remove_class_from_lock_chains(pf, class);
 }
 
 static void reinit_class(struct lock_class *class)
                reinit_class(class);
 
        list_splice_init(&pf->zapped, &free_lock_classes);
+
+#ifdef CONFIG_PROVE_LOCKING
+       bitmap_andnot(lock_chains_in_use, lock_chains_in_use,
+                     pf->lock_chains_being_freed, ARRAY_SIZE(lock_chains));
+       bitmap_clear(pf->lock_chains_being_freed, 0, ARRAY_SIZE(lock_chains));
+#endif
 }
 
 static void free_zapped_rcu(struct rcu_head *ch)
 #ifdef CONFIG_PROVE_LOCKING
                + sizeof(lock_cq)
                + sizeof(lock_chains)
+               + sizeof(lock_chains_in_use)
                + sizeof(chain_hlocks)
 #endif
                ) / 1024