From 8d95e1f29930551a68bc4e08fac4670a9e39280a Mon Sep 17 00:00:00 2001
From: regnat <rg@regnat.ovh>
Date: Thu, 3 Jun 2021 06:59:19 +0200
Subject: [PATCH] Set the cache when opening a flake

---
 src/libexpr/eval.cc        | 22 +++++++++++++++++++---
 src/libexpr/eval.hh        |  3 +++
 src/libexpr/flake/flake.cc |  5 +++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/libexpr/eval.cc b/src/libexpr/eval.cc
index 5ca4ecfe7..77301b75d 100644
--- a/src/libexpr/eval.cc
+++ b/src/libexpr/eval.cc
@@ -1173,10 +1173,11 @@ bool EvalState::getAttrField(Value & attrs, const std::vector<Symbol> & selector
         auto [ cacheResult, resultingCursor ] = eval_cache.getValue(*this, selector, dest);
         switch (cacheResult.returnCode) {
             case ValueCache::CacheHit:
+                if (cacheResult.lastQueriedSymbolIfMissing)
+                    return false;
                 return true;
-            case ValueCache::CacheMiss:
-                return false;
-            case ValueCache::Forward: // Fixme: Handle properly
+            case ValueCache::CacheMiss: // FIXME: Handle properly
+            case ValueCache::Forward: // FIXME: Handle properly
             case ValueCache::NoCacheKey:
             case ValueCache::UnCacheable:
                 ;
@@ -2162,6 +2163,21 @@ Strings EvalSettings::getDefaultNixPath()
     return res;
 }
 
+std::shared_ptr<tree_cache::Cache> EvalState::openTreeCache(Hash cacheKey)
+{
+    if (auto iter = evalCache.find(cacheKey); iter != evalCache.end())
+        return iter->second;
+
+    if (!(evalSettings.useEvalCache && evalSettings.pureEval))
+        return nullptr;
+    auto thisCache = tree_cache::Cache::tryCreate(
+            cacheKey,
+            symbols
+    );
+    evalCache.insert({cacheKey, thisCache});
+    return thisCache;
+}
+
 EvalSettings evalSettings;
 
 static GlobalConfig::Register rEvalSettings(&evalSettings);
diff --git a/src/libexpr/eval.hh b/src/libexpr/eval.hh
index b3ed03233..28c153a2d 100644
--- a/src/libexpr/eval.hh
+++ b/src/libexpr/eval.hh
@@ -116,6 +116,7 @@ private:
     typedef std::map<Path, Value> FileEvalCache;
 #endif
     FileEvalCache fileEvalCache;
+    std::map<Hash, std::shared_ptr<tree_cache::Cache>> evalCache;
 
     SearchPath searchPath;
 
@@ -132,6 +133,8 @@ public:
     EvalState(const Strings & _searchPath, ref<Store> store);
     ~EvalState();
 
+    std::shared_ptr<tree_cache::Cache> openTreeCache(Hash);
+
     void addToSearchPath(const string & s);
 
     SearchPath getSearchPath() { return searchPath; }
diff --git a/src/libexpr/flake/flake.cc b/src/libexpr/flake/flake.cc
index 2e94490d4..217bb1e8e 100644
--- a/src/libexpr/flake/flake.cc
+++ b/src/libexpr/flake/flake.cc
@@ -620,6 +620,11 @@ void callFlake(EvalState & state,
     state.callFunction(**vCallFlake, *vLocks, *vTmp1, noPos);
     state.callFunction(*vTmp1, *vRootSrc, *vTmp2, noPos);
     state.callFunction(*vTmp2, *vRootSubdir, vRes, noPos);
+    state.forceAttrs(vRes); // FIXME: Make more robust and lazy
+    auto fingerprint = lockedFlake.getFingerprint();
+    auto evalCache = state.openTreeCache(fingerprint);
+    auto cacheRoot = evalCache ? evalCache->getRoot() : nullptr;
+    vRes.attrs->eval_cache = ValueCache(cacheRoot);
 }
 
 static void prim_getFlake(EvalState & state, const Pos & pos, Value * * args, Value & v)