#ifdef CONFIG_MEMORY_HOTPLUG
 /*
- * Validate the node associated with the memory section we are
- * trying to add.
- */
-int valid_hot_add_scn(int *nid, unsigned long start, u32 lmb_size,
-                     unsigned long scn_addr)
-{
-       nodemask_t nodes;
-
-       if (*nid < 0 || !node_online(*nid))
-               *nid = any_online_node(NODE_MASK_ALL);
-
-       if ((scn_addr >= start) && (scn_addr < (start + lmb_size))) {
-               nodes_setall(nodes);
-               while (NODE_DATA(*nid)->node_spanned_pages == 0) {
-                       node_clear(*nid, nodes);
-                       *nid = any_online_node(nodes);
-               }
-
-               return 1;
-       }
-
-       return 0;
-}
-
-/*
- * Find the node associated with a hot added memory section represented
- * by the ibm,dynamic-reconfiguration-memory node.
+ * Find the node associated with a hot added memory section for
+ * memory represented in the device tree by the property
+ * ibm,dynamic-reconfiguration-memory/ibm,dynamic-memory.
  */
 static int hot_add_drconf_scn_to_nid(struct device_node *memory,
                                     unsigned long scn_addr)
 {
        const u32 *dm;
-       unsigned int n, rc;
+       unsigned int drconf_cell_cnt, rc;
        unsigned long lmb_size;
-       int default_nid = any_online_node(NODE_MASK_ALL);
-       int nid;
        struct assoc_arrays aa;
+       int nid = -1;
 
-       n = of_get_drconf_memory(memory, &dm);
-       if (!n)
-               return default_nid;;
+       drconf_cell_cnt = of_get_drconf_memory(memory, &dm);
+       if (!drconf_cell_cnt)
+               return -1;
 
        lmb_size = of_get_lmb_size(memory);
        if (!lmb_size)
-               return default_nid;
+               return -1;
 
        rc = of_get_assoc_arrays(memory, &aa);
        if (rc)
-               return default_nid;
+               return -1;
 
-       for (; n != 0; --n) {
+       for (; drconf_cell_cnt != 0; --drconf_cell_cnt) {
                struct of_drconf_cell drmem;
 
                read_drconf_cell(&drmem, &dm);
                    || !(drmem.flags & DRCONF_MEM_ASSIGNED))
                        continue;
 
+               if ((scn_addr < drmem.base_addr)
+                   || (scn_addr >= (drmem.base_addr + lmb_size)))
+                       continue;
+
                nid = of_drconf_to_nid_single(&drmem, &aa);
+               break;
+       }
+
+       return nid;
+}
+
+/*
+ * Find the node associated with a hot added memory section for memory
+ * represented in the device tree as a node (i.e. memory@XXXX) for
+ * each lmb.
+ */
+int hot_add_node_scn_to_nid(unsigned long scn_addr)
+{
+       struct device_node *memory = NULL;
+       int nid = -1;
+
+       while ((memory = of_find_node_by_type(memory, "memory")) != NULL) {
+               unsigned long start, size;
+               int ranges;
+               const unsigned int *memcell_buf;
+               unsigned int len;
+
+               memcell_buf = of_get_property(memory, "reg", &len);
+               if (!memcell_buf || len <= 0)
+                       continue;
+
+               /* ranges in cell */
+               ranges = (len >> 2) / (n_mem_addr_cells + n_mem_size_cells);
+
+               while (ranges--) {
+                       start = read_n_cells(n_mem_addr_cells, &memcell_buf);
+                       size = read_n_cells(n_mem_size_cells, &memcell_buf);
+
+                       if ((scn_addr < start) || (scn_addr >= (start + size)))
+                               continue;
+
+                       nid = of_node_to_nid_single(memory);
+                       break;
+               }
 
-               if (valid_hot_add_scn(&nid, drmem.base_addr, lmb_size,
-                                     scn_addr))
-                       return nid;
+               of_node_put(memory);
+               if (nid >= 0)
+                       break;
        }
 
-       BUG();  /* section address should be found above */
-       return 0;
+       return nid;
 }
 
 /*
 int hot_add_scn_to_nid(unsigned long scn_addr)
 {
        struct device_node *memory = NULL;
-       int nid;
+       int nid, found = 0;
 
        if (!numa_enabled || (min_common_depth < 0))
                return any_online_node(NODE_MASK_ALL);
        if (memory) {
                nid = hot_add_drconf_scn_to_nid(memory, scn_addr);
                of_node_put(memory);
-               return nid;
+       } else {
+               nid = hot_add_node_scn_to_nid(scn_addr);
        }
 
-       while ((memory = of_find_node_by_type(memory, "memory")) != NULL) {
-               unsigned long start, size;
-               int ranges;
-               const unsigned int *memcell_buf;
-               unsigned int len;
-
-               memcell_buf = of_get_property(memory, "reg", &len);
-               if (!memcell_buf || len <= 0)
-                       continue;
+       if (nid < 0 || !node_online(nid))
+               nid = any_online_node(NODE_MASK_ALL);
 
-               /* ranges in cell */
-               ranges = (len >> 2) / (n_mem_addr_cells + n_mem_size_cells);
-ha_new_range:
-               start = read_n_cells(n_mem_addr_cells, &memcell_buf);
-               size = read_n_cells(n_mem_size_cells, &memcell_buf);
-               nid = of_node_to_nid_single(memory);
+       if (NODE_DATA(nid)->node_spanned_pages)
+               return nid;
 
-               if (valid_hot_add_scn(&nid, start, size, scn_addr)) {
-                       of_node_put(memory);
-                       return nid;
+       for_each_online_node(nid) {
+               if (NODE_DATA(nid)->node_spanned_pages) {
+                       found = 1;
+                       break;
                }
-
-               if (--ranges)           /* process all ranges in cell */
-                       goto ha_new_range;
        }
-       BUG();  /* section address should be found above */
-       return 0;
+
+       BUG_ON(!found);
+       return nid;
 }
+
 #endif /* CONFIG_MEMORY_HOTPLUG */