5 files changed, 67 insertions, 0 deletions

diff --git a/include/erebos/storage.h b/include/erebos/storage.h
index e894fee..15ee0bb 100644
--- a/include/erebos/storage.h
+++ b/include/erebos/storage.h
@@ -208,9 +208,11 @@ public:
 
 	vector<Ref> previous() const;
 	class Generation generation() const;
+	vector<Digest> roots() const;
 
 private:
 	class Generation generationLocked() const;
+	class vector<Digest> rootsLocked() const;
 
 protected:
 	Ref(const std::shared_ptr<const Priv> p): PartialRef(p) {}
@@ -412,6 +414,8 @@ public:
 	std::vector<Stored<T>> previous() const;
 	bool precedes(const Stored<T> &) const;
 
+	std::vector<Digest> roots() const { return p->ref.roots(); }
+
 	const Ref & ref() const { return p->ref; }
 
 private:
diff --git a/src/main.cpp b/src/main.cpp
index a3facea..4dc4582 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -149,6 +149,19 @@ void storedGeneration(const vector<string> & args)
 	printLine(ss.str());
 }
 
+void storedRoots(const vector<string> & args)
+{
+	auto ref = st.ref(Digest(args.at(0)));
+	if (!ref)
+		throw invalid_argument("ref " + args.at(0) + " not found");
+
+	ostringstream ss;
+	ss << "stored-roots " << string(ref->digest());
+	for (const auto & dgst : ref->roots())
+		ss << " " << string(dgst);
+	printLine(ss.str());
+}
+
 void createIdentity(const vector<string> & args)
 {
 	optional<Identity> identity;
@@ -339,6 +352,7 @@ void attachReject(const vector<string> & params)
 vector<Command> commands = {
 	{ "store", store },
 	{ "stored-generation", storedGeneration },
+	{ "stored-roots", storedRoots },
 	{ "create-identity", createIdentity },
 	{ "start-server", startServer },
 	{ "stop-server", stopServer },
diff --git a/src/storage.cpp b/src/storage.cpp
index 5784309..c2d097e 100644
--- a/src/storage.cpp
+++ b/src/storage.cpp
@@ -1029,6 +1029,34 @@ Generation Ref::generationLocked() const
 	return gen;
 }
 
+vector<Digest> Ref::roots() const
+{
+	scoped_lock lock(p->storage->p->rootsCacheLock);
+	return rootsLocked();
+}
+
+vector<Digest> Ref::rootsLocked() const
+{
+	auto it = p->storage->p->rootsCache.find(p->digest);
+	if (it != p->storage->p->rootsCache.end())
+		return it->second;
+
+	vector<Digest> roots;
+	auto prev = previous();
+
+	if (prev.empty()) {
+		roots.push_back(p->digest);
+	} else {
+		for (const auto & p : previous())
+			for (const auto & r : p.rootsLocked())
+				roots.push_back(r);
+
+		std::sort(roots.begin(), roots.end());
+		roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
+	}
+	return roots;
+}
+
 
 template<class S>
 RecordT<S>::Item::operator bool() const
diff --git a/src/storage.h b/src/storage.h
index ef335b8..30e4213 100644
--- a/src/storage.h
+++ b/src/storage.h
@@ -167,6 +167,9 @@ struct PartialStorage::Priv
 
 	mutable mutex generationCacheLock {};
 	mutable unordered_map<Digest, Generation> generationCache {};
+
+	mutable mutex rootsCacheLock {};
+	mutable unordered_map<Digest, vector<Digest>> rootsCache {};
 };
 
 struct PartialRef::Priv
diff --git a/test/storage.test b/test/storage.test
index 2ce87d5..10f5d2c 100644
--- a/test/storage.test
+++ b/test/storage.test
@@ -45,6 +45,18 @@ test:
 	send to p1 "stored-generation $r4"
 	expect from p1 /stored-generation $r4 2/
 
+	send to p1 "stored-roots $r1"
+	expect from p1 /stored-roots $r1 $r1/
+
+	send to p1 "stored-roots $r2"
+	expect from p1 /stored-roots $r2 $r1/
+
+	send to p1 "stored-roots $r3"
+	expect from p1 /stored-roots $r3 $r1/
+
+	send to p1 "stored-roots $r4"
+	expect from p1 /stored-roots $r4 $r1/
+
 	# Attach second root
 	send to p1:
 		"store rec"
@@ -73,3 +85,9 @@ test:
 
 	send to p1 "stored-generation $r2_2"
 	expect from p1 /stored-generation $r2_2 1/
+
+	send to p1 "stored-roots $r2_3"
+	expect from p1 /stored-roots $r2_3 $r2_1 $r1/
+
+	send to p1 "stored-roots $r2_2"
+	expect from p1 /stored-roots $r2_2 $r2_1/