Store StructuredAttrs directly in Derivation

Instead of parsing a structured attrs at some later point, we parsed it
right away when parsing the A-Term format, and likewise serialize it to
`__json = <JSON dump>` when serializing a derivation to A-Term.

The JSON format can directly contain the JSON structured attrs without
so encoding it, so we just do that.

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

@@ -32,12 +32,9 @@ DerivationBuildingGoal::DerivationBuildingGoal(
 {
     drv = std::make_unique<Derivation>(drv_);
 
-    if (auto parsedOpt = StructuredAttrs::tryParse(drv->env)) {
-        parsedDrv = std::make_unique<StructuredAttrs>(*parsedOpt);
-    }
     try {
         drvOptions =
-            std::make_unique<DerivationOptions>(DerivationOptions::fromStructuredAttrs(drv->env, parsedDrv.get()));
+            std::make_unique<DerivationOptions>(DerivationOptions::fromStructuredAttrs(drv->env, drv->structuredAttrs));
     } catch (Error & e) {
         e.addTrace({}, "while parsing derivation '%s'", worker.store.printStorePath(drvPath));
         throw;
@@ -661,7 +658,6 @@ Goal::Co DerivationBuildingGoal::tryToBuild()
                     buildMode,
                     buildResult,
                     *drv,
-                    parsedDrv.get(),
                     *drvOptions,
                     inputPaths,
                     initialOutputs,

src/libstore/build/derivation-goal.cc

@@ -63,9 +63,8 @@ Goal::Co DerivationGoal::haveDerivation()
     trace("have derivation");
 
     auto drvOptions = [&]() -> DerivationOptions {
-        auto parsedOpt = StructuredAttrs::tryParse(drv->env);
         try {
-            return DerivationOptions::fromStructuredAttrs(drv->env, parsedOpt ? &*parsedOpt : nullptr);
+            return DerivationOptions::fromStructuredAttrs(drv->env, drv->structuredAttrs);
         } catch (Error & e) {
             e.addTrace({}, "while parsing derivation '%s'", worker.store.printStorePath(drvPath));
             throw;

src/libstore/derivation-options.cc

@@ -92,6 +92,12 @@ using OutputChecks = DerivationOptions::OutputChecks;
 
 using OutputChecksVariant = std::variant<OutputChecks, std::map<std::string, OutputChecks>>;
 
+DerivationOptions DerivationOptions::fromStructuredAttrs(
+    const StringMap & env, const std::optional<StructuredAttrs> & parsed, bool shouldWarn)
+{
+    return fromStructuredAttrs(env, parsed ? &*parsed : nullptr);
+}
+
 DerivationOptions
 DerivationOptions::fromStructuredAttrs(const StringMap & env, const StructuredAttrs * parsed, bool shouldWarn)
 {

src/libstore/derivations.cc

@@ -452,6 +452,7 @@ Derivation parseDerivation(
         expect(str, ")");
         drv.env.insert_or_assign(std::move(name), std::move(value));
     }
+    drv.structuredAttrs = StructuredAttrs::tryExtract(drv.env);
 
     expect(str, ")");
     return drv;
@@ -685,16 +686,28 @@ std::string Derivation::unparse(
 
     s += ",[";
     first = true;
-    for (auto & i : env) {
-        if (first)
-            first = false;
-        else
+
+    auto unparseEnv = [&](const StringPairs atermEnv) {
+        for (auto & i : atermEnv) {
+            if (first)
+                first = false;
+            else
+                s += ',';
+            s += '(';
+            printString(s, i.first);
             s += ',';
-        s += '(';
-        printString(s, i.first);
-        s += ',';
-        printString(s, maskOutputs && outputs.count(i.first) ? "" : i.second);
-        s += ')';
+            printString(s, maskOutputs && outputs.count(i.first) ? "" : i.second);
+            s += ')';
+        }
+    };
+
+    StructuredAttrs::checkKeyNotInUse(env);
+    if (structuredAttrs) {
+        StringPairs scratch = env;
+        scratch.insert(structuredAttrs->unparse());
+        unparseEnv(scratch);
+    } else {
+        unparseEnv(env);
     }
 
     s += "])";
@@ -948,6 +961,7 @@ Source & readDerivation(Source & in, const StoreDirConfig & store, BasicDerivati
         auto value = readString(in);
         drv.env[key] = value;
     }
+    drv.structuredAttrs = StructuredAttrs::tryExtract(drv.env);
 
     return in;
 }
@@ -983,9 +997,21 @@ void writeDerivation(Sink & out, const StoreDirConfig & store, const BasicDeriva
     }
     CommonProto::write(store, CommonProto::WriteConn{.to = out}, drv.inputSrcs);
     out << drv.platform << drv.builder << drv.args;
-    out << drv.env.size();
-    for (auto & i : drv.env)
-        out << i.first << i.second;
+
+    auto writeEnv = [&](const StringPairs atermEnv) {
+        out << atermEnv.size();
+        for (auto & [k, v] : atermEnv)
+            out << k << v;
+    };
+
+    StructuredAttrs::checkKeyNotInUse(drv.env);
+    if (drv.structuredAttrs) {
+        StringPairs scratch = drv.env;
+        scratch.insert(drv.structuredAttrs->unparse());
+        writeEnv(scratch);
+    } else {
+        writeEnv(drv.env);
+    }
 }
 
 std::string hashPlaceholder(const OutputNameView outputName)
@@ -1017,6 +1043,17 @@ void BasicDerivation::applyRewrites(const StringMap & rewrites)
         newEnv.emplace(envName, envValue);
     }
     env = std::move(newEnv);
+
+    if (structuredAttrs) {
+        // TODO rewrite the JSON AST properly, rather than dump parse round trip.
+        auto [k, jsonS] = structuredAttrs->unparse();
+        jsonS = rewriteStrings(jsonS, rewrites);
+        StringPairs newEnv;
+        newEnv.insert(std::pair{k, std::move(jsonS)});
+        auto newStructuredAttrs = StructuredAttrs::tryExtract(newEnv);
+        assert(newStructuredAttrs);
+        structuredAttrs = std::move(*newStructuredAttrs);
+    }
 }
 
 static void rewriteDerivation(Store & store, BasicDerivation & drv, const StringMap & rewrites)
@@ -1338,10 +1375,8 @@ nlohmann::json Derivation::toJSON(const StoreDirConfig & store) const
     res["args"] = args;
     res["env"] = env;
 
-    if (auto it = env.find("__json"); it != env.end()) {
-        res["env"].erase("__json");
-        res["structuredAttrs"] = nlohmann::json::parse(it->second);
-    }
+    if (structuredAttrs)
+        res["structuredAttrs"] = structuredAttrs->structuredAttrs;
 
     return res;
 }
@@ -1411,7 +1446,7 @@ Derivation Derivation::fromJSON(
     }
 
     if (auto structuredAttrs = get(json, "structuredAttrs"))
-        res.env.insert_or_assign("__json", structuredAttrs->dump());
+        res.structuredAttrs = StructuredAttrs{*structuredAttrs};
 
     return res;
 }

src/libstore/include/nix/store/build/derivation-building-goal.hh

@@ -1,8 +1,8 @@
 #pragma once
 ///@file
 
-#include "nix/store/parsed-derivations.hh"
 #include "nix/store/derivations.hh"
+#include "nix/store/parsed-derivations.hh"
 #include "nix/store/derivation-options.hh"
 #include "nix/store/build/derivation-building-misc.hh"
 #include "nix/store/outputs-spec.hh"
@@ -39,7 +39,6 @@ struct DerivationBuildingGoal : public Goal
      */
     std::unique_ptr<Derivation> drv;
 
-    std::unique_ptr<StructuredAttrs> parsedDrv;
     std::unique_ptr<DerivationOptions> drvOptions;
 
     /**

src/libstore/include/nix/store/derivation-options.hh

@@ -176,6 +176,9 @@ struct DerivationOptions
     static DerivationOptions
     fromStructuredAttrs(const StringMap & env, const StructuredAttrs * parsed, bool shouldWarn = true);
 
+    static DerivationOptions
+    fromStructuredAttrs(const StringMap & env, const std::optional<StructuredAttrs> & parsed, bool shouldWarn = true);
+
     /**
      * @param drv Must be the same derivation we parsed this from. In
      * the future we'll flip things around so a `BasicDerivation` has

src/libstore/include/nix/store/derivations.hh

@@ -7,6 +7,7 @@
 #include "nix/store/content-address.hh"
 #include "nix/util/repair-flag.hh"
 #include "nix/store/derived-path-map.hh"
+#include "nix/store/parsed-derivations.hh"
 #include "nix/util/sync.hh"
 #include "nix/util/variant-wrapper.hh"
 
@@ -286,7 +287,12 @@ struct BasicDerivation
     std::string platform;
     Path builder;
     Strings args;
+    /**
+     * Must not contain the key `__json`, at least in order to serialize to A-Term.
+     */
     StringPairs env;
+    std::optional<StructuredAttrs> structuredAttrs;
+
     std::string name;
 
     BasicDerivation() = default;

src/libstore/include/nix/store/parsed-derivations.hh

@@ -18,7 +18,30 @@ struct StructuredAttrs
 {
     nlohmann::json structuredAttrs;
 
-    static std::optional<StructuredAttrs> tryParse(const StringPairs & env);
+    bool operator==(const StructuredAttrs &) const = default;
+
+    /**
+     * Unconditionally parse from a JSON string. Used by `tryExtract`.
+     */
+    static StructuredAttrs parse(const std::string & encoded);
+
+    /**
+     * Like `tryParse`, but removes the env var which encoded the structured
+     * attrs from the map if one is found.
+     */
+    static std::optional<StructuredAttrs> tryExtract(StringPairs & env);
+
+    /**
+     * Opposite of `tryParse`, at least if one makes a map from this
+     * single key-value PR.
+     */
+    std::pair<std::string_view, std::string> unparse() const;
+
+    /**
+     * Ensures that the structured attrs "env var" is not in used, so we
+     * are free to use it instead.
+     */
+    static void checkKeyNotInUse(const StringPairs & env);
 
     nlohmann::json prepareStructuredAttrs(
         Store & store,

src/libstore/misc.cc

@@ -224,13 +224,11 @@ MissingPaths Store::queryMissing(const std::vector<DerivedPath> & targets)
                         return;
 
                     auto drv = make_ref<Derivation>(derivationFromPath(drvPath));
-                    auto parsedDrv = StructuredAttrs::tryParse(drv->env);
                     DerivationOptions drvOptions;
                     try {
                         // FIXME: this is a lot of work just to get the value
                         // of `allowSubstitutes`.
-                        drvOptions =
-                            DerivationOptions::fromStructuredAttrs(drv->env, parsedDrv ? &*parsedDrv : nullptr);
+                        drvOptions = DerivationOptions::fromStructuredAttrs(drv->env, drv->structuredAttrs);
                     } catch (Error & e) {
                         e.addTrace({}, "while parsing derivation '%s'", printStorePath(drvPath));
                         throw;

src/libstore/parsed-derivations.cc

@@ -8,20 +8,41 @@
 
 namespace nix {
 
-std::optional<StructuredAttrs> StructuredAttrs::tryParse(const StringPairs & env)
+static constexpr std::string_view envVarName = "__json";
+
+StructuredAttrs StructuredAttrs::parse(const std::string & encoded)
+{
+    try {
+        return StructuredAttrs{
+            .structuredAttrs = nlohmann::json::parse(encoded),
+        };
+    } catch (std::exception & e) {
+        throw Error("cannot process __json attribute: %s", e.what());
+    }
+}
+
+std::optional<StructuredAttrs> StructuredAttrs::tryExtract(StringPairs & env)
 {
     /* Parse the __json attribute, if any. */
-    auto jsonAttr = env.find("__json");
+    auto jsonAttr = env.find(envVarName);
     if (jsonAttr != env.end()) {
-        try {
-            return StructuredAttrs{
-                .structuredAttrs = nlohmann::json::parse(jsonAttr->second),
-            };
-        } catch (std::exception & e) {
-            throw Error("cannot process __json attribute: %s", e.what());
-        }
-    }
-    return {};
+        auto encoded = std::move(jsonAttr->second);
+        env.erase(jsonAttr);
+        return parse(encoded);
+    } else
+        return {};
+}
+
+std::pair<std::string_view, std::string> StructuredAttrs::unparse() const
+{
+    return {envVarName, structuredAttrs.dump()};
+}
+
+void StructuredAttrs::checkKeyNotInUse(const StringPairs & env)
+{
+    if (env.count(envVarName))
+        throw Error(
+            "Cannot have an environment variable named '__json'. This key is reserved for encoding structured attrs");
 }
 
 static std::regex shVarName("[A-Za-z_][A-Za-z0-9_]*");
@@ -175,4 +196,5 @@ std::string StructuredAttrs::writeShell(const nlohmann::json & json)
 
     return jsonSh;
 }
+
 } // namespace nix

src/libstore/unix/build/derivation-builder.cc

@@ -809,7 +809,7 @@ void DerivationBuilderImpl::startBuilder()
     writeStructuredAttrs();
 
     /* Handle exportReferencesGraph(), if set. */
-    if (!parsedDrv) {
+    if (!drv.structuredAttrs) {
         for (auto & [fileName, ss] : drvOptions.exportReferencesGraph) {
             StorePathSet storePathSet;
             for (auto & storePathS : ss) {
@@ -1081,7 +1081,7 @@ void DerivationBuilderImpl::initEnv()
     /* In non-structured mode, set all bindings either directory in the
        environment or via a file, as specified by
        `DerivationOptions::passAsFile`. */
-    if (!parsedDrv) {
+    if (!drv.structuredAttrs) {
         for (auto & i : drv.env) {
             if (drvOptions.passAsFile.find(i.first) == drvOptions.passAsFile.end()) {
                 env[i.first] = i.second;
@@ -1149,8 +1149,8 @@ void DerivationBuilderImpl::initEnv()
 
 void DerivationBuilderImpl::writeStructuredAttrs()
 {
-    if (parsedDrv) {
-        auto json = parsedDrv->prepareStructuredAttrs(store, drvOptions, inputPaths, drv.outputs);
+    if (drv.structuredAttrs) {
+        auto json = drv.structuredAttrs->prepareStructuredAttrs(store, drvOptions, inputPaths, drv.outputs);
         nlohmann::json rewritten;
         for (auto & [i, v] : json["outputs"].get<nlohmann::json::object_t>()) {
             /* The placeholder must have a rewrite, so we use it to cover both the

src/libstore/unix/include/nix/store/build/derivation-builder.hh

@@ -27,15 +27,6 @@ struct DerivationBuilderParams
      */
     const Derivation & drv;
 
-    /**
-     * The "structured attrs" of `drv`, if it has them.
-     *
-     * @todo this should be part of `Derivation`.
-     *
-     * @todo this should be renamed from `parsedDrv`.
-     */
-    const StructuredAttrs * parsedDrv;
-
     /**
      * The derivation options of `drv`.
      *
@@ -63,14 +54,12 @@ struct DerivationBuilderParams
         const BuildMode & buildMode,
         BuildResult & buildResult,
         const Derivation & drv,
-        const StructuredAttrs * parsedDrv,
         const DerivationOptions & drvOptions,
         const StorePathSet & inputPaths,
         std::map<std::string, InitialOutput> & initialOutputs)
         : drvPath{drvPath}
         , buildResult{buildResult}
         , drv{drv}
-        , parsedDrv{parsedDrv}
         , drvOptions{drvOptions}
         , inputPaths{inputPaths}
         , initialOutputs{initialOutputs}