[ledger] Add AppendReferences methods to TreeNode and Object

This will allow a unified handling of object references for local and
downloaded pieces.

Change-Id: I0718935b16079ebeb916a4b20ba3f8a432f7c868
Bug: LE-697
Tested: fx run-test ledger_tests -t ledger_unittests
diff --git a/src/ledger/bin/storage/fake/fake_object.cc b/src/ledger/bin/storage/fake/fake_object.cc
index b927ee1..08f0e9f 100644
--- a/src/ledger/bin/storage/fake/fake_object.cc
+++ b/src/ledger/bin/storage/fake/fake_object.cc
@@ -35,5 +35,10 @@
   return Status::OK;
 }
 
+Status FakeObject::AppendReferences(
+    ObjectReferencesAndPriority* references) const {
+  return Status::OK;
+}
+
 }  // namespace fake
 }  // namespace storage
diff --git a/src/ledger/bin/storage/fake/fake_object.h b/src/ledger/bin/storage/fake/fake_object.h
index 1b074ac..5d503e6 100644
--- a/src/ledger/bin/storage/fake/fake_object.h
+++ b/src/ledger/bin/storage/fake/fake_object.h
@@ -31,9 +31,10 @@
   explicit FakeObject(std::unique_ptr<const Piece> piece);
   explicit FakeObject(ObjectIdentifier identifier, fxl::StringView content);
 
-  // Object:
   ObjectIdentifier GetIdentifier() const override;
   Status GetData(fxl::StringView* data) const override;
+  Status AppendReferences(
+      ObjectReferencesAndPriority* references) const override;
 
  private:
   std::unique_ptr<const Piece> piece_;
diff --git a/src/ledger/bin/storage/impl/btree/tree_node.cc b/src/ledger/bin/storage/impl/btree/tree_node.cc
index e6a4f9b..4b86f77 100644
--- a/src/ledger/bin/storage/impl/btree/tree_node.cc
+++ b/src/ledger/bin/storage/impl/btree/tree_node.cc
@@ -23,6 +23,30 @@
 
 namespace storage {
 namespace btree {
+namespace {
+
+// Extracts |children| and |entries| references to non-inline values into
+// |references|.
+void ExtractReferences(const std::vector<Entry>& entries,
+                       const std::map<size_t, ObjectIdentifier>& children,
+                       ObjectReferencesAndPriority* references) {
+  for (const auto& [size, identifier] : children) {
+    const auto& digest = identifier.object_digest();
+    FXL_DCHECK(storage::IsDigestValid(digest));
+    if (!GetObjectDigestInfo(digest).is_inlined()) {
+      // Node-node references are always treated as eager.
+      references->emplace(digest, KeyPriority::EAGER);
+    }
+  }
+  for (const auto& entry : entries) {
+    const auto& digest = entry.object_identifier.object_digest();
+    if (!GetObjectDigestInfo(digest).is_inlined()) {
+      references->emplace(digest, entry.priority);
+    }
+  }
+}
+
+}  // namespace
 
 TreeNode::TreeNode(ObjectIdentifier identifier, uint8_t level,
                    std::vector<Entry> entries,
@@ -83,20 +107,7 @@
     fit::function<void(Status, ObjectIdentifier)> callback) {
   FXL_DCHECK(children.empty() || children.rbegin()->first <= entries.size());
   ObjectReferencesAndPriority tree_references;
-  for (const auto& [size, identifier] : children) {
-    const auto& digest = identifier.object_digest();
-    FXL_DCHECK(storage::IsDigestValid(digest));
-    if (!GetObjectDigestInfo(digest).is_inlined()) {
-      // Node-node references are always treated as eager.
-      tree_references.emplace(digest, KeyPriority::EAGER);
-    }
-  }
-  for (const auto& entry : entries) {
-    const auto& digest = entry.object_identifier.object_digest();
-    if (!GetObjectDigestInfo(digest).is_inlined()) {
-      tree_references.emplace(digest, entry.priority);
-    }
-  }
+  ExtractReferences(entries, children, &tree_references);
   std::string encoding = EncodeNode(level, entries, children);
   page_storage->AddObjectFromLocal(
       ObjectType::TREE_NODE, storage::DataSource::Create(std::move(encoding)),
@@ -111,6 +122,10 @@
   return Status::OK;
 }
 
+void TreeNode::AppendReferences(ObjectReferencesAndPriority* references) const {
+  ExtractReferences(entries_, children_, references);
+}
+
 Status TreeNode::FindKeyOrChild(convert::ExtendedStringView key,
                                 int* index) const {
   if (key.empty()) {
diff --git a/src/ledger/bin/storage/impl/btree/tree_node.h b/src/ledger/bin/storage/impl/btree/tree_node.h
index 24a58bc..5c6e85c 100644
--- a/src/ledger/bin/storage/impl/btree/tree_node.h
+++ b/src/ledger/bin/storage/impl/btree/tree_node.h
@@ -57,6 +57,10 @@
   // to be in [0, GetKeyCount() - 1].
   Status GetEntry(int index, Entry* entry) const;
 
+  // Adds to |references| the references from this node to its non-inlined
+  // children and values.
+  void AppendReferences(ObjectReferencesAndPriority* references) const;
+
   // Searches for the given |key| in this node. If it is found, |OK| is
   // returned and index contains the index of the entry. If not,
   // |INTERNAL_NOT_FOUND| is returned and index stores the index of the child
diff --git a/src/ledger/bin/storage/impl/btree/tree_node_unittest.cc b/src/ledger/bin/storage/impl/btree/tree_node_unittest.cc
index 02be759..8be9e54 100644
--- a/src/ledger/bin/storage/impl/btree/tree_node_unittest.cc
+++ b/src/ledger/bin/storage/impl/btree/tree_node_unittest.cc
@@ -228,6 +228,45 @@
 
   const ObjectDigest digest0 = object0->GetIdentifier().object_digest();
   const ObjectDigest digest1 = object1->GetIdentifier().object_digest();
+
+  // Check that references returned by each TreeNode are correct.
+  ObjectReferencesAndPriority references;
+  root->AppendReferences(&references);
+  EXPECT_THAT(
+      references,
+      UnorderedElementsAre(
+          // Keys
+          Pair(digest1, KeyPriority::LAZY),   // key03
+          Pair(digest1, KeyPriority::EAGER),  // key07
+          // Children
+          Pair(child0->GetIdentifier().object_digest(), KeyPriority::EAGER),
+          Pair(child1->GetIdentifier().object_digest(), KeyPriority::EAGER),
+          Pair(child2->GetIdentifier().object_digest(), KeyPriority::EAGER)));
+  references.clear();
+  child0->AppendReferences(&references);
+  EXPECT_THAT(references,
+              UnorderedElementsAre(Pair(digest0, KeyPriority::LAZY),   // key00
+                                   Pair(digest1, KeyPriority::EAGER),  // key01
+                                   Pair(digest0, KeyPriority::EAGER)   // key02
+                                   ));
+  references.clear();
+  child1->AppendReferences(&references);
+  EXPECT_THAT(
+      references,
+      UnorderedElementsAre(Pair(digest0, KeyPriority::LAZY),  // key04 and key06
+                           Pair(digest1, KeyPriority::EAGER)  // key05
+                           ));
+  references.clear();
+  child2->AppendReferences(&references);
+  EXPECT_THAT(references,
+              UnorderedElementsAre(
+                  Pair(digest0, KeyPriority::EAGER),  // key08 and key10
+                  Pair(digest1, KeyPriority::LAZY)    // key09
+                  // No reference to key11 (points to inline object02)
+                  ));
+
+  // Check that references have been correctly added to PageStorage during
+  // object creation.
   EXPECT_THAT(
       fake_storage_.GetReferences(),
       // All the pieces are small enough not to get split so we know all objects
diff --git a/src/ledger/bin/storage/impl/object_impl.cc b/src/ledger/bin/storage/impl/object_impl.cc
index a981454..547904e 100644
--- a/src/ledger/bin/storage/impl/object_impl.cc
+++ b/src/ledger/bin/storage/impl/object_impl.cc
@@ -8,6 +8,7 @@
 
 #include <utility>
 
+#include "src/ledger/bin/storage/impl/btree/tree_node.h"
 #include "src/ledger/bin/storage/impl/file_index.h"
 #include "src/ledger/bin/storage/impl/file_index_generated.h"
 #include "src/ledger/bin/storage/impl/object_digest.h"
@@ -79,6 +80,25 @@
 
 ObjectIdentifier LevelDBPiece::GetIdentifier() const { return identifier_; }
 
+Status BaseObject::AppendReferences(
+    ObjectReferencesAndPriority* references) const {
+  FXL_DCHECK(references);
+  // Blobs have no references.
+  const auto digest_info = GetObjectDigestInfo(GetIdentifier().object_digest());
+  if (digest_info.object_type == ObjectType::BLOB) {
+    return Status::OK;
+  }
+  FXL_DCHECK(digest_info.object_type == ObjectType::TREE_NODE);
+  // Parse the object into a TreeNode.
+  std::unique_ptr<const btree::TreeNode> node;
+  Status status = btree::TreeNode::FromObject(*this, &node);
+  if (status != Status::OK) {
+    return status;
+  }
+  node->AppendReferences(references);
+  return Status::OK;
+}
+
 ChunkObject::ChunkObject(std::unique_ptr<const Piece> piece)
     : piece_(std::move(piece)) {
   FXL_DCHECK(
diff --git a/src/ledger/bin/storage/impl/object_impl.h b/src/ledger/bin/storage/impl/object_impl.h
index d5a06b4..fb52860 100644
--- a/src/ledger/bin/storage/impl/object_impl.h
+++ b/src/ledger/bin/storage/impl/object_impl.h
@@ -68,8 +68,15 @@
   std::unique_ptr<leveldb::Iterator> iterator_;
 };
 
+// Common methods shared by all object implementations.
+class BaseObject : public Object {
+ public:
+  Status AppendReferences(
+      ObjectReferencesAndPriority* references) const override;
+};
+
 // Object whose data is backed by a single chunk piece.
-class ChunkObject : public Object {
+class ChunkObject : public BaseObject {
  public:
   // |piece| must be of type CHUNK; index pieces cannot be turned into objects
   // automatically.
@@ -84,7 +91,7 @@
 };
 
 // Object whose data is backed by a VMO.
-class VmoObject : public Object {
+class VmoObject : public BaseObject {
  public:
   VmoObject(ObjectIdentifier identifier, fsl::SizedVmo vmo);
   ~VmoObject() override;
diff --git a/src/ledger/bin/storage/impl/object_impl_unittest.cc b/src/ledger/bin/storage/impl/object_impl_unittest.cc
index 4f2c788..c53a983 100644
--- a/src/ledger/bin/storage/impl/object_impl_unittest.cc
+++ b/src/ledger/bin/storage/impl/object_impl_unittest.cc
@@ -12,6 +12,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "peridot/lib/scoped_tmpfs/scoped_tmpfs.h"
+#include "src/ledger/bin/storage/impl/btree/encoding.h"
 #include "src/ledger/bin/storage/impl/constants.h"
 #include "src/ledger/bin/storage/impl/file_index.h"
 #include "src/ledger/bin/storage/impl/object_digest.h"
@@ -25,6 +26,7 @@
 
 namespace storage {
 namespace {
+using ::testing::IsEmpty;
 using ::testing::Pair;
 using ::testing::UnorderedElementsAre;
 
@@ -220,5 +222,76 @@
   EXPECT_TRUE(CheckObjectValue(object, identifier, data));
 }
 
+TEST_F(ObjectImplTest, ObjectReferences) {
+  // Create various types of identifiers for the object children and values.
+  constexpr size_t kInlineSize = kStorageHashSize;
+  const ObjectIdentifier inline_blob = CreateObjectIdentifier(
+      ComputeObjectDigest(PieceType::CHUNK, ObjectType::BLOB,
+                          RandomString(environment_.random(), kInlineSize)));
+  ASSERT_TRUE(GetObjectDigestInfo(inline_blob.object_digest()).is_inlined());
+
+  const ObjectIdentifier inline_treenode = CreateObjectIdentifier(
+      ComputeObjectDigest(PieceType::CHUNK, ObjectType::TREE_NODE,
+                          RandomString(environment_.random(), kInlineSize)));
+  ASSERT_TRUE(
+      GetObjectDigestInfo(inline_treenode.object_digest()).is_inlined());
+
+  constexpr size_t kNoInlineSize = kStorageHashSize + 1;
+  const ObjectIdentifier noinline_blob = CreateObjectIdentifier(
+      ComputeObjectDigest(PieceType::CHUNK, ObjectType::BLOB,
+                          RandomString(environment_.random(), kNoInlineSize)));
+  ASSERT_FALSE(GetObjectDigestInfo(noinline_blob.object_digest()).is_inlined());
+
+  const ObjectIdentifier noinline_treenode = CreateObjectIdentifier(
+      ComputeObjectDigest(PieceType::CHUNK, ObjectType::TREE_NODE,
+                          RandomString(environment_.random(), kNoInlineSize)));
+  ASSERT_FALSE(
+      GetObjectDigestInfo(noinline_treenode.object_digest()).is_inlined());
+
+  // Create a TreeNode object.
+  const std::string data =
+      btree::EncodeNode(/*level=*/0,
+                        {
+                            Entry{"key01", inline_blob, KeyPriority::EAGER},
+                            Entry{"key02", noinline_blob, KeyPriority::EAGER},
+                            Entry{"key03", inline_blob, KeyPriority::LAZY},
+                            Entry{"key04", noinline_blob, KeyPriority::LAZY},
+                        },
+                        {
+                            {0, inline_treenode},
+                            {1, noinline_treenode},
+                        });
+  const ChunkObject tree_object(std::make_unique<DataChunkPiece>(
+      CreateObjectIdentifier(
+          ComputeObjectDigest(PieceType::CHUNK, ObjectType::TREE_NODE, data)),
+      DataSource::DataChunk::Create(data)));
+
+  // Check that inline children are not included in the references.
+  ObjectReferencesAndPriority references;
+  ASSERT_EQ(Status::OK, tree_object.AppendReferences(&references));
+  EXPECT_THAT(references,
+              UnorderedElementsAre(
+                  Pair(noinline_blob.object_digest(), KeyPriority::EAGER),
+                  Pair(noinline_blob.object_digest(), KeyPriority::LAZY),
+                  Pair(noinline_treenode.object_digest(), KeyPriority::EAGER)));
+
+  // Create a blob object with the exact same data and check that it does not
+  // return any reference.
+  const ChunkObject blob_object(std::make_unique<DataChunkPiece>(
+      CreateObjectIdentifier(
+          ComputeObjectDigest(PieceType::CHUNK, ObjectType::BLOB, data)),
+      DataSource::DataChunk::Create(data)));
+  references.clear();
+  ASSERT_EQ(Status::OK, blob_object.AppendReferences(&references));
+  EXPECT_THAT(references, IsEmpty());
+
+  // Create an invalid tree node object and check that we return an error.
+  const ChunkObject invalid_object(std::make_unique<DataChunkPiece>(
+      CreateObjectIdentifier(
+          ComputeObjectDigest(PieceType::CHUNK, ObjectType::TREE_NODE, "")),
+      DataSource::DataChunk::Create("")));
+  ASSERT_EQ(Status::FORMAT_ERROR, invalid_object.AppendReferences(&references));
+}
+
 }  // namespace
 }  // namespace storage
diff --git a/src/ledger/bin/storage/public/object.h b/src/ledger/bin/storage/public/object.h
index d6dbd5e..87260ca 100644
--- a/src/ledger/bin/storage/public/object.h
+++ b/src/ledger/bin/storage/public/object.h
@@ -31,6 +31,12 @@
   // Returns a vmo containing the data.
   virtual Status GetVmo(fsl::SizedVmo* vmo) const;
 
+  // Adds tree-level references from this object to other objects into
+  // |references|. Does not clear |references|. Does not add piece-level
+  // references (use |Piece::AppendReferences| instead).
+  virtual Status AppendReferences(
+      ObjectReferencesAndPriority* references) const = 0;
+
  private:
   FXL_DISALLOW_COPY_AND_ASSIGN(Object);
 };
@@ -50,10 +56,9 @@
   // piece is not deleted.
   virtual fxl::StringView GetData() const = 0;
 
-  // If this piece references other pieces (eg. an index referencing some
-  // chunks), adds them to |references|. Does not clear |references|. Does not
-  // add |Object| references even if the piece represents a full object of its
-  // own.
+  // Adds piece-level references from this piece to other pieces into
+  // |references|. Does not clear |references|. Does not add tree-level
+  // references (use |Object::AppendReferences| instead).
   virtual Status AppendReferences(
       ObjectReferencesAndPriority* references) const = 0;
 
diff --git a/src/ledger/bin/storage/public/object_unittest.cc b/src/ledger/bin/storage/public/object_unittest.cc
index 115b24c..fb017e4 100644
--- a/src/ledger/bin/storage/public/object_unittest.cc
+++ b/src/ledger/bin/storage/public/object_unittest.cc
@@ -25,6 +25,11 @@
     return Status::OK;
   }
 
+  Status AppendReferences(
+      ObjectReferencesAndPriority* references) const override {
+    return Status::OK;
+  }
+
  private:
   std::string value_;
 };