[lldb] Add MemoryTagMap class

The tag map holds a sparse set of memory tags and allows
you to query ranges for tags.

Granules that do not have tags will be set to llvm::None.
to keep the ordering intact. If there are no tags for the
requested range we'll just return an empty result so that
callers don't need to check that all values are llvm::None.

This will be combined with MemoryTagManager's MakeTaggedRanges:
* MakeTaggedRanges
* Read from all those ranges
* Insert the results into the tag map
* Give the tag map to whatever needs to print tags

Which in this case will be "memory read"/DumpDataExtractor.

Reviewed By: JDevlieghere

Differential Revision: https://reviews.llvm.org/D112825

lldb/include/lldb/Target/MemoryTagMap.h

@@ -0,0 +1,98 @@
+//===-- MemoryTagMap.h ------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLDB_TARGET_MEMORYTAGMAP_H
+#define LLDB_TARGET_MEMORYTAGMAP_H
+
+#include "lldb/Target/MemoryTagManager.h"
+#include "lldb/lldb-private.h"
+#include "llvm/ADT/Optional.h"
+#include <map>
+
+namespace lldb_private {
+
+/// MemoryTagMap provides a way to give a sparse read result
+/// when reading memory tags for a range. This is useful when
+/// you want to annotate some large memory dump that might include
+/// tagged memory but you don't know that it is all tagged.
+class MemoryTagMap {
+public:
+  /// Init an empty tag map
+  ///
+  /// \param [in] manager
+  ///     Non-null pointer to a memory tag manager.
+  MemoryTagMap(const MemoryTagManager *manager);
+
+  /// Insert tags into the map starting from addr.
+  ///
+  /// \param [in] addr
+  ///     Start address of the range to insert tags for.
+  ///     This address should be granule aligned and have had
+  ///     any non address bits removed.
+  ///     (ideally you would use the base of the range you used
+  ///     to read the tags in the first place)
+  ///
+  /// \param [in] tags
+  ///     Vector of tags to insert. The first tag will be inserted
+  ///     at addr, the next at addr+granule size and so on until
+  ///     all tags have been inserted.
+  void InsertTags(lldb::addr_t addr, const std::vector<lldb::addr_t> tags);
+
+  bool Empty() const;
+
+  /// Lookup memory tags for a range of memory from addr to addr+len.
+  ///
+  /// \param [in] addr
+  ///    The start of the range. This may include non address bits and
+  ///    does not have to be granule aligned.
+  ///
+  /// \param [in] len
+  ///    The length in bytes of the range to read tags for. This does
+  ///    not need to be multiple of the granule size.
+  ///
+  /// \return
+  ///    A vector containing the tags found for the granules in the
+  ///    range. (which is the result of granule aligning the given range)
+  ///
+  ///    Each item in the vector is an optional tag. Meaning that if
+  ///    it is valid then the granule had a tag and if not, it didn't.
+  ///
+  ///    If the range had no tags at all, the vector will be empty.
+  ///    If some of the range was tagged it will have items and some
+  ///    of them may be llvm::None.
+  ///    (this saves the caller checking whether all items are llvm::None)
+  std::vector<llvm::Optional<lldb::addr_t>> GetTags(lldb::addr_t addr,
+                                                    size_t len) const;
+
+private:
+  /// Lookup the tag for address
+  ///
+  /// \param [in] address
+  ///     The address to lookup a tag for. This should be aligned
+  ///     to a granule boundary.
+  ///
+  /// \return
+  ///     The tag for the granule that address refers to, or llvm::None
+  ///     if it has no memory tag.
+  llvm::Optional<lldb::addr_t> GetTag(lldb::addr_t addr) const;
+
+  // A map of granule aligned addresses to their memory tag
+  std::map<lldb::addr_t, lldb::addr_t> m_addr_to_tag;
+
+  // Memory tag manager used to align addresses and get granule size.
+  // Ideally this would be a const& but only certain architectures will
+  // have a memory tag manager class to provide here. So for a method
+  // returning a MemoryTagMap, Optional<MemoryTagMap> allows it to handle
+  // architectures without memory tagging. Optionals cannot hold references
+  // so we go with a pointer that we assume will be not be null.
+  const MemoryTagManager *m_manager;
+};
+
+} // namespace lldb_private
+
+#endif // LLDB_TARGET_MEMORYTAGMAP_H

lldb/source/Target/CMakeLists.txt

@@ -20,6 +20,7 @@ add_lldb_library(lldbTarget
   Memory.cpp
   MemoryHistory.cpp
   MemoryRegionInfo.cpp
+  MemoryTagMap.cpp
   ModuleCache.cpp
   OperatingSystem.cpp
   PathMappingList.cpp

lldb/source/Target/MemoryTagMap.cpp

@@ -0,0 +1,64 @@
+//===-- MemoryTagMap.cpp --------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "lldb/Target/MemoryTagMap.h"
+
+using namespace lldb_private;
+
+MemoryTagMap::MemoryTagMap(const MemoryTagManager *manager)
+    : m_manager(manager) {
+  assert(m_manager && "valid tag manager required to construct a MemoryTagMap");
+}
+
+void MemoryTagMap::InsertTags(lldb::addr_t addr,
+                              const std::vector<lldb::addr_t> tags) {
+  // We're assuming that addr has no non address bits and is granule aligned.
+  size_t granule_size = m_manager->GetGranuleSize();
+  for (auto tag : tags) {
+    m_addr_to_tag[addr] = tag;
+    addr += granule_size;
+  }
+}
+
+bool MemoryTagMap::Empty() const { return m_addr_to_tag.empty(); }
+
+std::vector<llvm::Optional<lldb::addr_t>>
+MemoryTagMap::GetTags(lldb::addr_t addr, size_t len) const {
+  // Addr and len might be unaligned
+  addr = m_manager->RemoveTagBits(addr);
+  MemoryTagManager::TagRange range(addr, len);
+  range = m_manager->ExpandToGranule(range);
+
+  std::vector<llvm::Optional<lldb::addr_t>> tags;
+  lldb::addr_t end_addr = range.GetRangeEnd();
+  addr = range.GetRangeBase();
+  bool got_valid_tags = false;
+  size_t granule_size = m_manager->GetGranuleSize();
+
+  for (; addr < end_addr; addr += granule_size) {
+    llvm::Optional<lldb::addr_t> tag = GetTag(addr);
+    tags.push_back(tag);
+    if (tag)
+      got_valid_tags = true;
+  }
+
+  // To save the caller checking if every item is llvm::None,
+  // we return an empty vector if we got no tags at all.
+  if (got_valid_tags)
+    return tags;
+  return {};
+}
+
+llvm::Optional<lldb::addr_t> MemoryTagMap::GetTag(lldb::addr_t addr) const {
+  // Here we assume that addr is granule aligned, just like when the tags
+  // were inserted.
+  auto found = m_addr_to_tag.find(addr);
+  if (found == m_addr_to_tag.end())
+    return llvm::None;
+  return found->second;
+}

lldb/unittests/Target/CMakeLists.txt

@@ -3,6 +3,7 @@ add_lldb_unittest(TargetTests
   DynamicRegisterInfoTest.cpp
   ExecutionContextTest.cpp
   MemoryRegionInfoTest.cpp
+  MemoryTagMapTest.cpp
   ModuleCacheTest.cpp
   PathMappingListTest.cpp
   RemoteAwarePlatformTest.cpp

lldb/unittests/Target/MemoryTagMapTest.cpp

@@ -0,0 +1,81 @@
+//===-- MemoryTagMapTest.cpp ----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "lldb/Target/MemoryTagMap.h"
+#include "Plugins/Process/Utility/MemoryTagManagerAArch64MTE.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace lldb_private;
+using namespace lldb;
+
+// In these tests we use the AArch64 MTE tag manager because it is the only
+// implementation of a memory tag manager. MemoryTagMap itself is generic.
+
+TEST(MemoryTagMapTest, EmptyTagMap) {
+  MemoryTagManagerAArch64MTE manager;
+  MemoryTagMap tag_map(&manager);
+
+  tag_map.InsertTags(0, {});
+  ASSERT_TRUE(tag_map.Empty());
+  tag_map.InsertTags(0, {0});
+  ASSERT_FALSE(tag_map.Empty());
+}
+
+TEST(MemoryTagMapTest, GetTags) {
+  using TagsVec = std::vector<llvm::Optional<lldb::addr_t>>;
+
+  MemoryTagManagerAArch64MTE manager;
+  MemoryTagMap tag_map(&manager);
+
+  // No tags for an address not in the map
+  ASSERT_TRUE(tag_map.GetTags(0, 16).empty());
+
+  tag_map.InsertTags(0, {0, 1});
+
+  // No tags if you read zero length
+  ASSERT_TRUE(tag_map.GetTags(0, 0).empty());
+
+  EXPECT_THAT(tag_map.GetTags(0, 16), ::testing::ContainerEq(TagsVec{0}));
+
+  EXPECT_THAT(tag_map.GetTags(0, 32), ::testing::ContainerEq(TagsVec{0, 1}));
+
+  // Last granule of the range is not tagged
+  EXPECT_THAT(tag_map.GetTags(0, 48),
+              ::testing::ContainerEq(TagsVec{0, 1, llvm::None}));
+
+  EXPECT_THAT(tag_map.GetTags(16, 32),
+              ::testing::ContainerEq(TagsVec{1, llvm::None}));
+
+  // Reading beyond that address gives you no tags at all
+  EXPECT_THAT(tag_map.GetTags(32, 16), ::testing::ContainerEq(TagsVec{}));
+
+  // Address is granule aligned for you
+  // The length here is set such that alignment doesn't produce a 2 granule
+  // range.
+  EXPECT_THAT(tag_map.GetTags(8, 8), ::testing::ContainerEq(TagsVec{0}));
+
+  EXPECT_THAT(tag_map.GetTags(30, 2), ::testing::ContainerEq(TagsVec{1}));
+
+  // Here the length pushes the range into the next granule. When aligned
+  // this produces 2 granules.
+  EXPECT_THAT(tag_map.GetTags(30, 4),
+              ::testing::ContainerEq(TagsVec{1, llvm::None}));
+
+  // A range can also have gaps at the beginning or in the middle.
+  // Add more tags, 1 granule away from the first range.
+  tag_map.InsertTags(48, {3, 4});
+
+  // Untagged first granule
+  EXPECT_THAT(tag_map.GetTags(32, 32),
+              ::testing::ContainerEq(TagsVec{llvm::None, 3}));
+
+  // Untagged middle granule
+  EXPECT_THAT(tag_map.GetTags(16, 48),
+              ::testing::ContainerEq(TagsVec{1, llvm::None, 3}));
+}