From eff120d1b93b99d6ae61b20e18c31a5324a2be4f Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <eelco.dolstra@logicblox.com>
Date: Mon, 22 Sep 2014 14:46:42 +0200
Subject: [PATCH] =?UTF-8?q?Add=20a=20function=20=E2=80=98valueSize?=
 =?UTF-8?q?=E2=80=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

It returns the size of value, including all other values and
environments reachable from it. It is intended for debugging memory
consumption issues.
---
 src/libexpr/eval.cc    | 80 ++++++++++++++++++++++++++++++++++++++++++
 src/libexpr/eval.hh    |  5 +--
 src/libexpr/primops.cc | 10 ++++++
 src/libexpr/value.hh   |  6 ++++
 4 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/src/libexpr/eval.cc b/src/libexpr/eval.cc
index cf8aafa8f..bd49cec4f 100644
--- a/src/libexpr/eval.cc
+++ b/src/libexpr/eval.cc
@@ -404,9 +404,12 @@ Value * EvalState::allocValue()
 
 Env & EvalState::allocEnv(unsigned int size)
 {
+    assert(size <= std::numeric_limits<decltype(Env::size)>::max());
+
     nrEnvs++;
     nrValuesInEnvs += size;
     Env * env = (Env *) GC_MALLOC(sizeof(Env) + size * sizeof(Value *));
+    env->size = size;
 
     /* Clear the values because maybeThunk() and lookupVar fromWith expects this. */
     for (unsigned i = 0; i < size; ++i)
@@ -1488,4 +1491,81 @@ void EvalState::printCanaries()
 }
 
 
+size_t valueSize(Value & v)
+{
+    std::set<const void *> seen;
+
+    auto doString = [&](const char * s) -> size_t {
+        if (seen.find(s) != seen.end()) return 0;
+        seen.insert(s);
+        return strlen(s) + 1;
+    };
+
+    std::function<size_t(Value & v)> doValue;
+    std::function<size_t(Env & v)> doEnv;
+
+    doValue = [&](Value & v) -> size_t {
+        if (seen.find(&v) != seen.end()) return 0;
+        seen.insert(&v);
+
+        size_t sz = sizeof(Value);
+
+        switch (v.type) {
+        case tString:
+            sz += doString(v.string.s);
+            if (v.string.context)
+                for (const char * * p = v.string.context; *p; ++p)
+                    sz += doString(*p);
+            break;
+        case tPath:
+            sz += doString(v.path);
+            break;
+        case tAttrs:
+            for (auto & i : *v.attrs)
+                sz += doValue(*i.value);
+            break;
+        case tList:
+            for (unsigned int n = 0; n < v.list.length; ++n)
+                sz += doValue(*v.list.elems[n]);
+            break;
+        case tThunk:
+            sz += doEnv(*v.thunk.env);
+            break;
+        case tApp:
+            sz += doValue(*v.app.left);
+            sz += doValue(*v.app.right);
+            break;
+        case tLambda:
+            sz += doEnv(*v.lambda.env);
+            break;
+        case tPrimOpApp:
+            sz += doValue(*v.primOpApp.left);
+            sz += doValue(*v.primOpApp.right);
+            break;
+        default:
+            ;
+        }
+
+        return sz;
+    };
+
+    doEnv = [&](Env & env) -> size_t {
+        if (seen.find(&env) != seen.end()) return 0;
+        seen.insert(&env);
+
+        size_t sz = sizeof(Env);
+
+        for (unsigned int i = 0; i < env.size; ++i)
+            if (env.values[i])
+                sz += doValue(*env.values[i]);
+
+        if (env.up) sz += doEnv(*env.up);
+
+        return sz;
+    };
+
+    return doValue(v);
+}
+
+
 }
diff --git a/src/libexpr/eval.hh b/src/libexpr/eval.hh
index dcd6209e3..d8ea0f0ce 100644
--- a/src/libexpr/eval.hh
+++ b/src/libexpr/eval.hh
@@ -95,8 +95,9 @@ struct PrimOp
 struct Env
 {
     Env * up;
-    unsigned short prevWith; // nr of levels up to next `with' environment
-    bool haveWithAttrs;
+    unsigned short size; // used by ‘valueSize’
+    unsigned short prevWith:15; // nr of levels up to next `with' environment
+    unsigned short haveWithAttrs:1;
     Value * values[0];
 };
 
diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index 66321c769..c721a5681 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -423,6 +423,13 @@ void prim_gcCanary(EvalState & state, const Pos & pos, Value * * args, Value & v
 }
 
 
+void prim_valueSize(EvalState & state, const Pos & pos, Value * * args, Value & v)
+{
+    /* We're not forcing the argument on purpose. */
+    mkInt(v, valueSize(*args[0]));
+}
+
+
 /*************************************************************
  * Derivations
  *************************************************************/
@@ -1416,8 +1423,11 @@ void EvalState::createBaseEnv()
     addPrimOp("__addErrorContext", 2, prim_addErrorContext);
     addPrimOp("__tryEval", 1, prim_tryEval);
     addPrimOp("__getEnv", 1, prim_getEnv);
+
+    // Debugging
     addPrimOp("__trace", 2, prim_trace);
     addPrimOp("__gcCanary", 1, prim_gcCanary);
+    addPrimOp("__valueSize", 1, prim_valueSize);
 
     // Paths
     addPrimOp("__toPath", 1, prim_toPath);
diff --git a/src/libexpr/value.hh b/src/libexpr/value.hh
index 2feb2f949..227f5e173 100644
--- a/src/libexpr/value.hh
+++ b/src/libexpr/value.hh
@@ -159,4 +159,10 @@ static inline void mkPathNoCopy(Value & v, const char * s)
 void mkPath(Value & v, const char * s);
 
 
+/* Compute the size in bytes of the given value, including all values
+   and environments reachable from it. Static expressions (Exprs) are
+   not included. */
+size_t valueSize(Value & v);
+
+
 }