From cf2a89184f2f7c8c27bb6b9014dc4060daffa89f Mon Sep 17 00:00:00 2001
From: Rebecca Turner <rbt@sent.as>
Date: Sat, 9 Mar 2024 14:01:24 -0800
Subject: [PATCH] Implement `repl-init-file`

Closes #9940
---
 src/libcmd/repl.cc | 79 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/libcmd/repl.cc b/src/libcmd/repl.cc
index 75f20d635842..b0fd5fcb3cda 100644
--- a/src/libcmd/repl.cc
+++ b/src/libcmd/repl.cc
@@ -43,7 +43,7 @@ extern "C" {
 #include "finally.hh"
 #include "markdown.hh"
 #include "local-fs-store.hh"
-#include "print.hh"
+#include "gc-small-vector.hh"
 
 #if HAVE_BOEHMGC
 #define GC_INCLUDE_NEW
@@ -114,6 +114,26 @@ struct NixRepl
     void evalString(std::string s, Value & v);
     void loadDebugTraceEnv(DebugTrace & dt);
 
+    /**
+     * Load the `repl-init-file` and add the resulting AttrSet to the top-level
+     * bindings.
+     */
+    void loadReplInitFile();
+    /**
+     * Get the `info` AttrSet that's passed as the first argument to the
+     * `repl-init-file`.
+     */
+    Value * replInitInfo();
+
+    /**
+     * Get the current top-level bindings as an AttrSet.
+     */
+    Value * bindingsToAttrs();
+    /**
+     * Parse a file, evaluate its result, and force the resulting value.
+     */
+    Value * evalFile(SourcePath & path);
+
     void printValue(std::ostream & str,
                               Value & v,
                               unsigned int maxDepth = std::numeric_limits<unsigned int>::max())
@@ -887,6 +907,42 @@ void NixRepl::loadFiles()
         notice("Loading installable '%1%'...", what);
         addAttrsToScope(*i);
     }
+
+    loadReplInitFile();
+}
+
+void NixRepl::loadReplInitFile()
+{
+    if (!evalSettings.replInitFile) {
+        return;
+    }
+
+    SourcePath replInitFile = evalSettings.replInitFile;
+    notice("Loading '%1%'...", replInitFile);
+    auto replInit = evalFile(replInitFile);
+
+    auto info = replInitInfo();
+
+    auto currentAttrs = bindingsToAttrs();
+
+    Value &newAttrs(*state->allocValue());
+    SmallValueVector<2> args = {info, currentAttrs};
+    state->callFunction(*replInit, args.size(), args.data(), newAttrs, replInit->determinePos(noPos));
+
+    addAttrsToScope(newAttrs);
+}
+
+Value * NixRepl::replInitInfo()
+{
+    auto builder = state->buildBindings(1);
+
+    Value * currentSystem(state->allocValue());
+    currentSystem->mkString(evalSettings.getCurrentSystem());
+    builder.insert(state->symbols.create("currentSystem"), currentSystem);
+
+    Value * info(state->allocValue());
+    info->mkAttrs(builder.finish());
+    return info;
 }
 
 
@@ -919,6 +975,18 @@ void NixRepl::addVarToScope(const Symbol name, Value & v)
     varNames.emplace(state->symbols[name]);
 }
 
+Value * NixRepl::bindingsToAttrs()
+{
+    auto builder = state->buildBindings(staticEnv->vars.size());
+    for (auto & [symbol, displacement] : staticEnv->vars) {
+        builder.insert(symbol, env->values[displacement]);
+    }
+
+    Value * attrs(state->allocValue());
+    attrs->mkAttrs(builder.finish());
+    return attrs;
+}
+
 
 Expr * NixRepl::parseString(std::string s)
 {
@@ -933,6 +1001,15 @@ void NixRepl::evalString(std::string s, Value & v)
     state->forceValue(v, v.determinePos(noPos));
 }
 
+Value * NixRepl::evalFile(SourcePath & path)
+{
+    auto expr = state->parseExprFromFile(evalSettings.replInitFile, staticEnv);
+    Value * result(state->allocValue());
+    expr->eval(*state, *env, *result);
+    state->forceValue(*result, result->determinePos(noPos));
+    return result;
+}
+
 
 std::unique_ptr<AbstractNixRepl> AbstractNixRepl::create(
    const SearchPath & searchPath, nix::ref<Store> store, ref<EvalState> state,