1 files changed, 29 insertions, 18 deletions

diff --git a/src/nix-env/nix-env.cc b/src/nix-env/nix-env.cc
index 4d8537705..b7d3f4337 100644
--- a/src/nix-env/nix-env.cc
+++ b/src/nix-env/nix-env.cc
@@ -13,7 +13,7 @@
 #include "json.hh"
 #include "value-to-json.hh"
 #include "xml-writer.hh"
-#include "legacy.hh"
+#include "../nix/legacy.hh"
 
 #include <cerrno>
 #include <ctime>
@@ -718,28 +718,39 @@ static void uninstallDerivations(Globals & globals, Strings & selectors,
     while (true) {
         string lockToken = optimisticLockProfile(profile);
 
-        DrvInfos installedElems = queryInstalled(*globals.state, profile);
-        DrvInfos newElems;
+        DrvInfos workingElems = queryInstalled(*globals.state, profile);
 
-        for (auto & i : installedElems) {
-            DrvName drvName(i.queryName());
-            bool found = false;
-            for (auto & j : selectors)
-                /* !!! the repeated calls to followLinksToStorePath()
-                   are expensive, should pre-compute them. */
-                if ((isPath(j) && globals.state->store->parseStorePath(i.queryOutPath()) == globals.state->store->followLinksToStorePath(j))
-                    || DrvName(j).matches(drvName))
-                {
-                    printInfo("uninstalling '%s'", i.queryName());
-                    found = true;
-                    break;
-                }
-            if (!found) newElems.push_back(i);
+        for (auto & selector : selectors) {
+            DrvInfos::iterator split = workingElems.begin();
+            if (isPath(selector)) {
+                StorePath selectorStorePath = globals.state->store->followLinksToStorePath(selector);
+                split = std::partition(
+                    workingElems.begin(), workingElems.end(),
+                    [&selectorStorePath, globals](auto &elem) {
+                        return selectorStorePath != globals.state->store->parseStorePath(elem.queryOutPath());
+                    }
+                );
+            } else {
+                DrvName selectorName(selector);
+                split = std::partition(
+                    workingElems.begin(), workingElems.end(),
+                    [&selectorName](auto &elem){
+                        DrvName elemName(elem.queryName());
+                        return !selectorName.matches(elemName);
+                    }
+                );
+            }
+            if (split == workingElems.end())
+                warn("selector '%s' matched no installed derivations", selector);
+            for (auto removedElem = split; removedElem != workingElems.end(); removedElem++) {
+                printInfo("uninstalling '%s'", removedElem->queryName());
+            }
+            workingElems.erase(split, workingElems.end());
         }
 
         if (globals.dryRun) return;
 
-        if (createUserEnv(*globals.state, newElems,
+        if (createUserEnv(*globals.state, workingElems,
                 profile, settings.envKeepDerivations, lockToken)) break;
     }
 }