4 files changed, 40 insertions, 14 deletions

diff --git a/src/fix.cc b/src/fix.cc
index 10f0e4413012..9b0d95912f10 100644
--- a/src/fix.cc
+++ b/src/fix.cc
@@ -213,7 +213,7 @@ static Expr evalExpr2(EvalState & state, Expr e)
 
             if (ATmatch(value, "FSId(<str>)", &s1)) {
                 FSId id = parseHash(s1);
-                Strings paths = fstatePaths(id, false);
+                Strings paths = fstatePaths(id);
                 if (paths.size() != 1) abort();
                 string path = *(paths.begin());
                 fs.derive.inputs.push_back(id);
diff --git a/src/nix.cc b/src/nix.cc
index 6dc5776e25db..e9f04ff59d74 100644
--- a/src/nix.cc
+++ b/src/nix.cc
@@ -80,17 +80,24 @@ static void opAdd(Strings opFlags, Strings opArgs)
 }
 
 
-string dotQuote(const string & s)
+static string dotQuote(const string & s)
 {
     return "\"" + s + "\"";
 }
 
 
+FSId maybeNormalise(const FSId & id, bool normalise)
+{
+    return normalise ? normaliseFState(id) : id;
+}
+
+
 /* Perform various sorts of queries. */
 static void opQuery(Strings opFlags, Strings opArgs)
 {
     enum { qList, qRefs, qGenerators, qExpansion, qGraph 
     } query = qList;
+    bool normalise = false;
 
     for (Strings::iterator i = opFlags.begin();
          i != opFlags.end(); i++)
@@ -99,6 +106,7 @@ static void opQuery(Strings opFlags, Strings opArgs)
         else if (*i == "--generators" || *i == "-g") query = qGenerators;
         else if (*i == "--expansion" || *i == "-e") query = qExpansion;
         else if (*i == "--graph") query = qGraph;
+        else if (*i == "--normalise" || *i == "-n") normalise = true;
         else throw UsageError(format("unknown flag `%1%'") % *i);
 
     switch (query) {
@@ -108,7 +116,8 @@ static void opQuery(Strings opFlags, Strings opArgs)
             for (Strings::iterator i = opArgs.begin();
                  i != opArgs.end(); i++)
             {
-                Strings paths2 = fstatePaths(argToId(*i), true);
+                Strings paths2 = fstatePaths(
+                    maybeNormalise(argToId(*i), normalise));
                 paths.insert(paths2.begin(), paths2.end());
             }
             for (StringSet::iterator i = paths.begin(); 
@@ -122,7 +131,8 @@ static void opQuery(Strings opFlags, Strings opArgs)
             for (Strings::iterator i = opArgs.begin();
                  i != opArgs.end(); i++)
             {
-                Strings paths2 = fstateRefs(argToId(*i));
+                Strings paths2 = fstateRefs(
+                    maybeNormalise(argToId(*i), normalise));
                 paths.insert(paths2.begin(), paths2.end());
             }
             for (StringSet::iterator i = paths.begin(); 
diff --git a/src/normalise.cc b/src/normalise.cc
index 95d96351af54..eefb790b6f98 100644
--- a/src/normalise.cc
+++ b/src/normalise.cc
@@ -217,12 +217,11 @@ void realiseSlice(const FSId & id, FSIdSet pending)
 }
 
 
-Strings fstatePaths(const FSId & id, bool normalise)
+Strings fstatePaths(const FSId & id)
 {
     Strings paths;
 
-    FState fs = parseFState(termFromId(
-        normalise ? normaliseFState(id) : id));
+    FState fs = parseFState(termFromId(id));
 
     if (fs.type == FState::fsSlice) {
         /* !!! fix complexity */
@@ -245,14 +244,31 @@ Strings fstatePaths(const FSId & id, bool normalise)
 }
 
 
+static void fstateRefsSet(const FSId & id, StringSet & paths)
+{
+    FState fs = parseFState(termFromId(id));
+
+    if (fs.type == FState::fsSlice) {
+        for (SliceElems::iterator i = fs.slice.elems.begin();
+             i != fs.slice.elems.end(); i++)
+            paths.insert(i->path);
+    }
+    
+    else if (fs.type == FState::fsDerive) {
+        for (FSIds::iterator i = fs.derive.inputs.begin();
+             i != fs.derive.inputs.end(); i++)
+            fstateRefsSet(*i, paths);
+    }
+
+    else abort();
+}
+
+
 Strings fstateRefs(const FSId & id)
 {
-    Strings paths;
-    FState fs = parseFState(termFromId(normaliseFState(id)));
-    for (SliceElems::const_iterator i = fs.slice.elems.begin();
-         i != fs.slice.elems.end(); i++)
-        paths.push_back(i->path);
-    return paths;
+    StringSet paths;
+    fstateRefsSet(id, paths);
+    return Strings(paths.begin(), paths.end());
 }
 
 
diff --git a/src/normalise.hh b/src/normalise.hh
index 619fef8bf623..a5b45c861926 100644
--- a/src/normalise.hh
+++ b/src/normalise.hh
@@ -13,7 +13,7 @@ void realiseSlice(const FSId & id, FSIdSet pending = FSIdSet());
 
 /* Get the list of root (output) paths of the given
    fstate-expression. */
-Strings fstatePaths(const FSId & id, bool normalise);
+Strings fstatePaths(const FSId & id);
 
 /* Get the list of paths referenced by the given fstate-expression. */
 Strings fstateRefs(const FSId & id);