Source complete env in nix-shell with __structuredAttrs = true;

This is needed to push the adoption of structured attrs[1] forward. It's
now checked if a `__json` exists in the environment-map of the derivation
to be openend in a `nix-shell`.

Derivations with structured attributes enabled also make use of a file
named `.attrs.json` containing every environment variable represented as
JSON which is useful for e.g. `exportReferencesGraph`[2]. To
provide an environment similar to the build sandbox, `nix-shell` now
adds a `.attrs.json` to `cwd` (which is mostly equal to the one in the
build sandbox) and removes it using an exit hook when closing the shell.

Internally, the following things changed:

* The logic which generates both `.attrs.json` & `.attrs.sh` used to be
  tied to `LocalDerivationGoal`, the representation of a derivation's build
  on the local store. This was e.g. necessary to resolve references to
  input derivations for `exportReferencesGraph`[2].

* To make sure that the logic is actually usable for `nix-shell`, the
  behavior had to be moved to the parent-struct `DerivationGoal` to support
  `nix-shell --store`. For this reason (I assume), the worker doesn't
  return a `LocalDerivationGoal` by default when passing a `Derivation`
  to it.

* However the `inputRewrites` - used to provide placeholders even if the
  output path isn't known - aren't part of a (remote) `DerivationGoal`
  and thus input rewriting is optional in the structured-attrs generator
  of `DerivationGoal`. Currently, `LocalDerivationGoal` passes the hash
  rewrites to `DerivationGoal` (i.e. the behavior remains equal), but
  for a `nix-shell` this behavior is entirely omitted, from my
  understanding this isn't a problem since a `nix-shell` isn't actually
  supposed to be used for creating store-paths.

* The `nix-shell` now emulates the first step of a derivation's build
  using a `DerivationGoal` struct in order to retrieve both the raw JSON
  values for `.attrs.json` and the shell code for the rcfile.

  As this is practically a no-op for non-structured-attrs builds, this
  code-path is only reached if `__json` is part of the environment.

  In a quick, pretty unscientific performance measuring I observed
  a (reproducible) performance increase of 0.2s for opening a `nix-shell`
  which is probably related to the way this is implemented. While I'm
  open for suggestions for alternative implementations, I'd consider
  this to be OK as long as it's not really noticeable for an
  end-user.

[1] https://nixos.mayflower.consulting/blog/2020/01/20/structured-attrs/
[2] https://nixos.org/manual/nix/unstable/expressions/advanced-attributes.html#advanced-attributes

src/libstore/build/derivation-goal.cc

@@ -143,6 +143,115 @@ void DerivationGoal::work()
     (this->*state)();
 }
 
+static std::regex shVarName("[A-Za-z_][A-Za-z0-9_]*");
+std::optional<StructuredAttrsWithShellRC> DerivationGoal::generateStructuredAttrs(
+    std::optional<StringMap> inputRewrites)
+{
+    if (!parsedDrv) return std::nullopt;
+    auto structuredAttrs = parsedDrv->getStructuredAttrs();
+    if (!structuredAttrs) return std::nullopt;
+
+    auto json = *structuredAttrs;
+
+    /* Add an "outputs" object containing the output paths. */
+    nlohmann::json outputs;
+    for (auto & i : drv->outputs) {
+        if (inputRewrites) {
+            /* The placeholder must have a rewrite, so we use it to cover both the
+               cases where we know or don't know the output path ahead of time. */
+            outputs[i.first] = rewriteStrings(hashPlaceholder(i.first), inputRewrites.value());
+        } else {
+            /* This case is only relevant for the nix-shell */
+            outputs[i.first] = hashPlaceholder(i.first);
+        }
+    }
+    json["outputs"] = outputs;
+
+    /* Handle exportReferencesGraph. */
+    auto e = json.find("exportReferencesGraph");
+    if (e != json.end() && e->is_object()) {
+        for (auto i = e->begin(); i != e->end(); ++i) {
+            std::ostringstream str;
+            {
+                JSONPlaceholder jsonRoot(str, true);
+                StorePathSet storePaths;
+                for (auto & p : *i)
+                    storePaths.insert(worker.store.parseStorePath(p.get<std::string>()));
+                worker.store.pathInfoToJSON(jsonRoot,
+                    exportReferences(storePaths), false, true);
+            }
+            json[i.key()] = nlohmann::json::parse(str.str()); // urgh
+        }
+    }
+
+    /* As a convenience to bash scripts, write a shell file that
+       maps all attributes that are representable in bash -
+       namely, strings, integers, nulls, Booleans, and arrays and
+       objects consisting entirely of those values. (So nested
+       arrays or objects are not supported.) */
+
+    auto handleSimpleType = [](const nlohmann::json & value) -> std::optional<std::string> {
+        if (value.is_string())
+            return shellEscape(value);
+
+        if (value.is_number()) {
+            auto f = value.get<float>();
+            if (std::ceil(f) == f)
+                return std::to_string(value.get<int>());
+        }
+
+        if (value.is_null())
+            return std::string("''");
+
+        if (value.is_boolean())
+            return value.get<bool>() ? std::string("1") : std::string("");
+
+        return {};
+    };
+
+    std::string jsonSh;
+
+    for (auto i = json.begin(); i != json.end(); ++i) {
+
+        if (!std::regex_match(i.key(), shVarName)) continue;
+
+        auto & value = i.value();
+
+        auto s = handleSimpleType(value);
+        if (s)
+            jsonSh += fmt("declare %s=%s\n", i.key(), *s);
+
+        else if (value.is_array()) {
+            std::string s2;
+            bool good = true;
+
+            for (auto i = value.begin(); i != value.end(); ++i) {
+                auto s3 = handleSimpleType(i.value());
+                if (!s3) { good = false; break; }
+                s2 += *s3; s2 += ' ';
+            }
+
+            if (good)
+                jsonSh += fmt("declare -a %s=(%s)\n", i.key(), s2);
+        }
+
+        else if (value.is_object()) {
+            std::string s2;
+            bool good = true;
+
+            for (auto i = value.begin(); i != value.end(); ++i) {
+                auto s3 = handleSimpleType(i.value());
+                if (!s3) { good = false; break; }
+                s2 += fmt("[%s]=%s ", shellEscape(i.key()), *s3);
+            }
+
+            if (good)
+                jsonSh += fmt("declare -A %s=(%s)\n", i.key(), s2);
+        }
+    }
+
+    return std::make_pair(jsonSh, json);
+}
 
 void DerivationGoal::addWantedOutputs(const StringSet & outputs)
 {

src/libstore/build/derivation-goal.hh

@@ -8,6 +8,7 @@
 
 namespace nix {
 
+typedef std::pair<std::string, nlohmann::json> StructuredAttrsWithShellRC;
 using std::map;
 
 struct HookInstance;
@@ -118,6 +119,8 @@ struct DerivationGoal : public Goal
 
     BuildResult result;
 
+    std::optional<StructuredAttrsWithShellRC> generateStructuredAttrs(std::optional<StringMap> inputRewrites);
+
     /* The current round, if we're building multiple times. */
     size_t curRound = 1;
 

src/libstore/build/local-derivation-goal.cc

@@ -1083,113 +1083,18 @@ void LocalDerivationGoal::initEnv()
 }
 
 
-static std::regex shVarName("[A-Za-z_][A-Za-z0-9_]*");
-
-
 void LocalDerivationGoal::writeStructuredAttrs()
 {
-    auto structuredAttrs = parsedDrv->getStructuredAttrs();
-    if (!structuredAttrs) return;
-
-    auto json = *structuredAttrs;
-
-    /* Add an "outputs" object containing the output paths. */
-    nlohmann::json outputs;
-    for (auto & i : drv->outputs) {
-        /* The placeholder must have a rewrite, so we use it to cover both the
-           cases where we know or don't know the output path ahead of time. */
-        outputs[i.first] = rewriteStrings(hashPlaceholder(i.first), inputRewrites);
-    }
-    json["outputs"] = outputs;
-
-    /* Handle exportReferencesGraph. */
-    auto e = json.find("exportReferencesGraph");
-    if (e != json.end() && e->is_object()) {
-        for (auto i = e->begin(); i != e->end(); ++i) {
-            std::ostringstream str;
-            {
-                JSONPlaceholder jsonRoot(str, true);
-                StorePathSet storePaths;
-                for (auto & p : *i)
-                    storePaths.insert(worker.store.parseStorePath(p.get<std::string>()));
-                worker.store.pathInfoToJSON(jsonRoot,
-                    exportReferences(storePaths), false, true);
-            }
-            json[i.key()] = nlohmann::json::parse(str.str()); // urgh
-        }
+    if (auto structAttrs = generateStructuredAttrs(inputRewrites)) {
+        auto value = structAttrs.value();
+        auto jsonSh = value.first;
+        auto json = value.second;
+
+        writeFile(tmpDir + "/.attrs.sh", rewriteStrings(jsonSh, inputRewrites));
+        chownToBuilder(tmpDir + "/.attrs.sh");
+        writeFile(tmpDir + "/.attrs.json", rewriteStrings(json.dump(), inputRewrites));
+        chownToBuilder(tmpDir + "/.attrs.json");
     }
-
-    writeFile(tmpDir + "/.attrs.json", rewriteStrings(json.dump(), inputRewrites));
-    chownToBuilder(tmpDir + "/.attrs.json");
-
-    /* As a convenience to bash scripts, write a shell file that
-       maps all attributes that are representable in bash -
-       namely, strings, integers, nulls, Booleans, and arrays and
-       objects consisting entirely of those values. (So nested
-       arrays or objects are not supported.) */
-
-    auto handleSimpleType = [](const nlohmann::json & value) -> std::optional<std::string> {
-        if (value.is_string())
-            return shellEscape(value);
-
-        if (value.is_number()) {
-            auto f = value.get<float>();
-            if (std::ceil(f) == f)
-                return std::to_string(value.get<int>());
-        }
-
-        if (value.is_null())
-            return std::string("''");
-
-        if (value.is_boolean())
-            return value.get<bool>() ? std::string("1") : std::string("");
-
-        return {};
-    };
-
-    std::string jsonSh;
-
-    for (auto i = json.begin(); i != json.end(); ++i) {
-
-        if (!std::regex_match(i.key(), shVarName)) continue;
-
-        auto & value = i.value();
-
-        auto s = handleSimpleType(value);
-        if (s)
-            jsonSh += fmt("declare %s=%s\n", i.key(), *s);
-
-        else if (value.is_array()) {
-            std::string s2;
-            bool good = true;
-
-            for (auto i = value.begin(); i != value.end(); ++i) {
-                auto s3 = handleSimpleType(i.value());
-                if (!s3) { good = false; break; }
-                s2 += *s3; s2 += ' ';
-            }
-
-            if (good)
-                jsonSh += fmt("declare -a %s=(%s)\n", i.key(), s2);
-        }
-
-        else if (value.is_object()) {
-            std::string s2;
-            bool good = true;
-
-            for (auto i = value.begin(); i != value.end(); ++i) {
-                auto s3 = handleSimpleType(i.value());
-                if (!s3) { good = false; break; }
-                s2 += fmt("[%s]=%s ", shellEscape(i.key()), *s3);
-            }
-
-            if (good)
-                jsonSh += fmt("declare -A %s=(%s)\n", i.key(), s2);
-        }
-    }
-
-    writeFile(tmpDir + "/.attrs.sh", rewriteStrings(jsonSh, inputRewrites));
-    chownToBuilder(tmpDir + "/.attrs.sh");
 }
 
 

src/nix-build/nix-build.cc

@@ -1,10 +1,15 @@
 #include <cstring>
 #include <fstream>
 #include <iostream>
+#include <filesystem>
 #include <regex>
 #include <sstream>
 #include <vector>
+#include <map>
 
+#include <nlohmann/json.hpp>
+
+#include "parsed-derivations.hh"
 #include "store-api.hh"
 #include "local-fs-store.hh"
 #include "globals.hh"
@@ -20,6 +25,9 @@
 #include "attr-path.hh"
 #include "legacy.hh"
 
+#include "build/worker.hh"
+#include "build/derivation-goal.hh"
+
 using namespace nix;
 using namespace std::string_literals;
 
@@ -422,12 +430,38 @@ static void main_nix_build(int argc, char * * argv)
             } else
                 env[var.first] = var.second;
 
+        std::string structuredAttrsRC = "";
+        std::string exitCmd = "";
+
+        if (env.count("__json")) {
+            auto storePath = store->parseStorePath(drvInfo.queryDrvPath());
+            Worker worker(*store);
+            StringSet wantedOutputs;
+            for (auto const & output : drvInfo.queryOutputs()) {
+                wantedOutputs.insert(output.first);
+            }
+
+            // FIXME this means that a drv will be re-evaluated just to get all inputs.
+            auto drvGoal = worker.makeDerivationGoal(storePath, wantedOutputs, bmNormal);
+            drvGoal->getDerivation();
+            drvGoal->inputsRealised();
+            if (auto structAttrs = drvGoal->generateStructuredAttrs(std::nullopt)) {
+                auto val = structAttrs.value();
+                structuredAttrsRC = val.first;
+                auto attrsJSON = std::filesystem::current_path().string() + "/.attrs.json";
+                writeFile(attrsJSON, val.second.dump());
+                exitCmd = "\n_rm_attrs_json() { rm -f " + attrsJSON + "; }"
+                    + "\nexitHooks+=(_rm_attrs_json)"
+                    + "\nfailureHooks+=(_rm_attrs_json)\n";
+            }
+        }
+
         /* Run a shell using the derivation's environment.  For
            convenience, source $stdenv/setup to setup additional
            environment variables and shell functions.  Also don't
            lose the current $PATH directories. */
         auto rcfile = (Path) tmpDir + "/rc";
-        writeFile(rcfile, fmt(
+        std::string rc = fmt(
                 R"(_nix_shell_clean_tmpdir() { rm -rf %1%; }; )"s +
                 (keepTmp ?
                     "trap _nix_shell_clean_tmpdir EXIT; "
@@ -436,8 +470,9 @@ static void main_nix_build(int argc, char * * argv)
                     "_nix_shell_clean_tmpdir; ") +
                 (pure ? "" : "[ -n \"$PS1\" ] && [ -e ~/.bashrc ] && source ~/.bashrc;") +
                 "%2%"
-                "dontAddDisableDepTrack=1; "
-                "[ -e $stdenv/setup ] && source $stdenv/setup; "
+                "dontAddDisableDepTrack=1;\n"
+                + structuredAttrsRC + exitCmd +
+                "\n[ -e $stdenv/setup ] && source $stdenv/setup; "
                 "%3%"
                 "PATH=%4%:\"$PATH\"; "
                 "SHELL=%5%; "
@@ -455,7 +490,9 @@ static void main_nix_build(int argc, char * * argv)
                 shellEscape(dirOf(*shell)),
                 shellEscape(*shell),
                 (getenv("TZ") ? (string("export TZ=") + shellEscape(getenv("TZ")) + "; ") : ""),
-                envCommand));
+                envCommand);
+        vomit("Sourcing nix-shell with file %s and contents:\n%s", rcfile, rc);
+        writeFile(rcfile, rc);
 
         Strings envStrs;
         for (auto & i : env)

tests/structured-attrs-shell.nix

@@ -0,0 +1,19 @@
+with import ./config.nix;
+let
+  dep = mkDerivation {
+    name = "dep";
+    buildCommand = ''
+      mkdir $out; echo bla > $out/bla
+    '';
+  };
+in
+mkDerivation {
+  name = "structured2";
+  __structuredAttrs = true;
+  outputs = [ "out" "dev" ];
+  my.list = [ "a" "b" "c" ];
+  exportReferencesGraph.refs = [ dep ];
+  buildCommand = ''
+    touch ''${outputs[out]}; touch ''${outputs[dev]}
+  '';
+}

tests/structured-attrs.sh

@@ -8,3 +8,9 @@ nix-build structured-attrs.nix -A all -o $TEST_ROOT/result
 
 [[ $(cat $TEST_ROOT/result/foo) = bar ]]
 [[ $(cat $TEST_ROOT/result-dev/foo) = foo ]]
+
+export NIX_BUILD_SHELL=$SHELL
+[[ ! -e '.attrs.json' ]]
+env NIX_PATH=nixpkgs=shell.nix nix-shell structured-attrs-shell.nix \
+    --run 'test -e .attrs.json; test "3" = "$(jq ".my.list|length" < .attrs.json)"'
+[[ ! -e '.attrs.json' ]]