diff --git a/mon/lp-optimizer.js b/mon/lp-optimizer.js
index 1ea455d4..241e57e0 100644
--- a/mon/lp-optimizer.js
+++ b/mon/lp-optimizer.js
@@ -502,44 +502,6 @@ function put_aligned_pgs(aligned_pgs, int_pgs, prev_int_pgs, keygen)
     }
 }
 
-// Convert multi-level osd_tree = { level: number|string, id?: string, size?: number, children?: osd_tree }[]
-// levels = { string: number }
-// to a two-level osd_tree suitable for all_combinations()
-// FIXME: Replace with extract_tree_levels({ level: -Infinity, children: osd_tree }, [ failure_domain_level, osd_level ], levels)
-function flatten_tree(osd_tree, levels, failure_domain_level, osd_level, domains = {}, i = { i: 1 })
-{
-    osd_level = levels[osd_level] || osd_level;
-    failure_domain_level = levels[failure_domain_level] || failure_domain_level;
-    for (const node of osd_tree)
-    {
-        if ((levels[node.level] || node.level) < failure_domain_level)
-        {
-            flatten_tree(node.children||[], levels, failure_domain_level, osd_level, domains, i);
-        }
-        else
-        {
-            domains['dom'+(i.i++)] = extract_osds([ node ], levels, osd_level);
-        }
-    }
-    return domains;
-}
-
-function extract_osds(osd_tree, levels, osd_level, osds = {})
-{
-    for (const node of osd_tree)
-    {
-        if ((levels[node.level] || node.level) >= osd_level)
-        {
-            osds[node.id] = node.size;
-        }
-        else
-        {
-            extract_osds(node.children||[], levels, osd_level, osds);
-        }
-    }
-    return osds;
-}
-
 // Convert multi-level tree_node = { level: number|string, id?: string, size?: number, children?: tree_node[] }
 // levels = { string: number }
 // to a multi-level OSD tree suitable for random_hier_combinations()
@@ -896,7 +858,6 @@ module.exports = {
     pg_weights_space_efficiency,
     pg_list_space_efficiency,
     pg_per_osd_space_efficiency,
-    flatten_tree,
     extract_tree_levels,
 
     lp_solve,
diff --git a/mon/mon.js b/mon/mon.js
index 02ced4b3..fcbfde9e 100644
--- a/mon/mon.js
+++ b/mon/mon.js
@@ -1112,27 +1112,23 @@ class Mon
     filter_osds_by_tags(orig_tree, flat_tree, tags)
     {
         if (!tags)
-        {
-            return;
-        }
+            return 1;
         for (const tag of (tags instanceof Array ? tags : [ tags ]))
         {
-            for (const host in flat_tree)
+            for (const item in flat_tree)
             {
-                let found = 0;
-                for (const osd in flat_tree[host])
+                if (flat_tree[item] instanceof Object)
                 {
-                    if (!orig_tree[osd].tags || !orig_tree[osd].tags[tag])
-                        delete flat_tree[host][osd];
-                    else
-                        found++;
-                }
-                if (!found)
-                {
-                    delete flat_tree[host];
+                    if (!filter_osds_by_tags(orig_tree, flat_tree[item], tags))
+                        delete flat_tree[item];
                 }
+                else if (!orig_tree[item].tags || !orig_tree[item].tags[tag])
+                    delete flat_tree[item];
             }
         }
+        for (const item in flat_tree)
+            return 1;
+        return 0;
     }
 
     get_affinity_osds(pool_cfg, up_osds, osd_tree)
@@ -1191,9 +1187,8 @@ class Mon
                 {
                     continue;
                 }
-                let pool_tree = osd_tree[pool_cfg.root_node || ''];
-                pool_tree = pool_tree ? pool_tree.children : [];
-                pool_tree = LPOptimizer.flatten_tree(pool_tree, levels, pool_cfg.failure_domain, 'osd');
+                let pool_tree = osd_tree[pool_cfg.root_node || ''] || {};
+                pool_tree = LPOptimizer.extract_tree_levels(pool_tree, [ pool_cfg.failure_domain, 'osd' ], levels);
                 this.filter_osds_by_tags(osd_tree, pool_tree, pool_cfg.osd_tags);
                 // These are for the purpose of building history.osd_sets
                 const real_prev_pgs = [];
diff --git a/mon/test-optimize-undersized.js b/mon/test-optimize-undersized.js
index d1a766e6..7571478c 100644
--- a/mon/test-optimize-undersized.js
+++ b/mon/test-optimize-undersized.js
@@ -36,7 +36,7 @@ const crush_tree = [
     ] },
 ];
 
-const osd_tree = LPOptimizer.flatten_tree(crush_tree, {}, 1, 3);
+const osd_tree = LPOptimizer.extract_tree_levels({ level: -Infinity, children: crush_tree }, [ 1, 3 ], {});
 console.log(osd_tree);
 
 async function run()
@@ -47,32 +47,32 @@ async function run()
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 0);
     console.log('\nAdding 1st failure domain:');
-    cur_tree['dom1'] = osd_tree['dom1'];
+    cur_tree['l1_1'] = osd_tree['l1_1'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 12 && res.total_space == 12);
     console.log('\nAdding 2nd failure domain:');
-    cur_tree['dom2'] = osd_tree['dom2'];
+    cur_tree['l1_2'] = osd_tree['l1_2'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 24 && res.total_space == 24);
     console.log('\nAdding 3rd failure domain:');
-    cur_tree['dom3'] = osd_tree['dom3'];
+    cur_tree['l1_3'] = osd_tree['l1_3'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 36 && res.total_space == 36);
     console.log('\nRemoving 3rd failure domain:');
-    delete cur_tree['dom3'];
+    delete cur_tree['l1_3'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 24 && res.total_space == 24);
     console.log('\nRemoving 2nd failure domain:');
-    delete cur_tree['dom2'];
+    delete cur_tree['l1_2'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 12 && res.total_space == 12);
     console.log('\nRemoving 1st failure domain:');
-    delete cur_tree['dom1'];
+    delete cur_tree['l1_1'];
     res = await LPOptimizer.optimize_change({ prev_pgs: res.int_pgs, osd_tree: cur_tree, pg_size: 3 });
     LPOptimizer.print_change_stats(res, false);
     assert(res.space == 0);
diff --git a/mon/test-optimize.js b/mon/test-optimize.js
index 309650b3..ee1c3d99 100644
--- a/mon/test-optimize.js
+++ b/mon/test-optimize.js
@@ -108,7 +108,11 @@ async function run()
     LPOptimizer.print_change_stats(res, false);
 
     console.log('\n256 PGs, size=3, failure domain=rack');
-    res = await LPOptimizer.optimize_initial({ osd_tree: LPOptimizer.flatten_tree(crush_tree, {}, 1, 3), pg_size: 3, pg_count: 256 });
+    res = await LPOptimizer.optimize_initial({
+        osd_tree: LPOptimizer.extract_tree_levels({ level: -Infinity, children: crush_tree }, [ 1, 3 ], {}),
+        pg_size: 3,
+        pg_count: 256,
+    });
     LPOptimizer.print_change_stats(res, false);
 }