diff --git a/src/trace_processor/importers/proto/heap_graph_module.cc b/src/trace_processor/importers/proto/heap_graph_module.cc
index ad2b806..cc1d7ad 100644
--- a/src/trace_processor/importers/proto/heap_graph_module.cc
+++ b/src/trace_processor/importers/proto/heap_graph_module.cc
@@ -36,72 +36,6 @@
 using ObjectTable = tables::HeapGraphObjectTable;
 using ReferenceTable = tables::HeapGraphReferenceTable;
 
-const char* HeapGraphRootTypeToString(int32_t type) {
-  switch (type) {
-    case protos::pbzero::HeapGraphRoot::ROOT_UNKNOWN:
-      return "ROOT_UNKNOWN";
-    case protos::pbzero::HeapGraphRoot::ROOT_JNI_GLOBAL:
-      return "ROOT_JNI_GLOBAL";
-    case protos::pbzero::HeapGraphRoot::ROOT_JNI_LOCAL:
-      return "ROOT_JNI_LOCAL";
-    case protos::pbzero::HeapGraphRoot::ROOT_JAVA_FRAME:
-      return "ROOT_JAVA_FRAME";
-    case protos::pbzero::HeapGraphRoot::ROOT_NATIVE_STACK:
-      return "ROOT_NATIVE_STACK";
-    case protos::pbzero::HeapGraphRoot::ROOT_STICKY_CLASS:
-      return "ROOT_STICKY_CLASS";
-    case protos::pbzero::HeapGraphRoot::ROOT_THREAD_BLOCK:
-      return "ROOT_THREAD_BLOCK";
-    case protos::pbzero::HeapGraphRoot::ROOT_MONITOR_USED:
-      return "ROOT_MONITOR_USED";
-    case protos::pbzero::HeapGraphRoot::ROOT_THREAD_OBJECT:
-      return "ROOT_THREAD_OBJECT";
-    case protos::pbzero::HeapGraphRoot::ROOT_INTERNED_STRING:
-      return "ROOT_INTERNED_STRING";
-    case protos::pbzero::HeapGraphRoot::ROOT_FINALIZING:
-      return "ROOT_FINALIZING";
-    case protos::pbzero::HeapGraphRoot::ROOT_DEBUGGER:
-      return "ROOT_DEBUGGER";
-    case protos::pbzero::HeapGraphRoot::ROOT_REFERENCE_CLEANUP:
-      return "ROOT_REFERENCE_CLEANUP";
-    case protos::pbzero::HeapGraphRoot::ROOT_VM_INTERNAL:
-      return "ROOT_VM_INTERNAL";
-    case protos::pbzero::HeapGraphRoot::ROOT_JNI_MONITOR:
-      return "ROOT_JNI_MONITOR";
-    default:
-      return "ROOT_UNKNOWN";
-  }
-}
-
-const char* HeapGraphTypeKindToString(int32_t type) {
-  switch (type) {
-    case protos::pbzero::HeapGraphType::KIND_NORMAL:
-      return "KIND_NORMAL";
-    case protos::pbzero::HeapGraphType::KIND_NOREFERENCES:
-      return "KIND_NOREFERENCES";
-    case protos::pbzero::HeapGraphType::KIND_STRING:
-      return "KIND_STRING";
-    case protos::pbzero::HeapGraphType::KIND_ARRAY:
-      return "KIND_ARRAY";
-    case protos::pbzero::HeapGraphType::KIND_CLASS:
-      return "KIND_CLASS";
-    case protos::pbzero::HeapGraphType::KIND_CLASSLOADER:
-      return "KIND_CLASSLOADER";
-    case protos::pbzero::HeapGraphType::KIND_DEXCACHE:
-      return "KIND_DEXCACHE";
-    case protos::pbzero::HeapGraphType::KIND_SOFT_REFERENCE:
-      return "KIND_SOFT_REFERENCE";
-    case protos::pbzero::HeapGraphType::KIND_WEAK_REFERENCE:
-      return "KIND_WEAK_REFERENCE";
-    case protos::pbzero::HeapGraphType::KIND_FINALIZER_REFERENCE:
-      return "KIND_FINALIZER_REFERENCE";
-    case protos::pbzero::HeapGraphType::KIND_PHANTOM_REFERENCE:
-      return "KIND_PHANTOM_REFERENCE";
-    default:
-      return "KIND_UNKNOWN";
-  }
-}
-
 // Iterate over a repeated field of varints, independent of whether it is
 // packed or not.
 template <int32_t field_no, typename T, typename F>
@@ -237,8 +171,13 @@
         entry.kind() == protos::pbzero::HeapGraphType::KIND_ARRAY ||
         entry.kind() == protos::pbzero::HeapGraphType::KIND_STRING;
 
-    StringId kind = context_->storage->InternString(
-        HeapGraphTypeKindToString(entry.kind()));
+    protos::pbzero::HeapGraphType::Kind kind =
+        protos::pbzero::HeapGraphType::KIND_UNKNOWN;
+    if (protos::pbzero::HeapGraphType_Kind_MIN <= entry.kind() &&
+        entry.kind() <= protos::pbzero::HeapGraphType_Kind_MAX) {
+      kind = protos::pbzero::HeapGraphType::Kind(entry.kind());
+    }
+
     std::optional<uint64_t> location_id;
     if (entry.has_location_id())
       location_id = entry.location_id();
@@ -265,11 +204,15 @@
   }
   for (auto it = heap_graph.roots(); it; ++it) {
     protos::pbzero::HeapGraphRoot::Decoder entry(*it);
-    const char* str = HeapGraphRootTypeToString(entry.root_type());
-    auto str_view = base::StringView(str);
 
     HeapGraphTracker::SourceRoot src_root;
-    src_root.root_type = context_->storage->InternString(str_view);
+    if (protos::pbzero::HeapGraphRoot_Type_MIN <= entry.root_type() &&
+        entry.root_type() <= protos::pbzero::HeapGraphRoot_Type_MAX) {
+      src_root.root_type =
+          protos::pbzero::HeapGraphRoot::Type(entry.root_type());
+    } else {
+      src_root.root_type = protos::pbzero::HeapGraphRoot::ROOT_UNKNOWN;
+    }
     // grep-friendly: object_ids
     bool parse_error =
         ForEachVarInt<protos::pbzero::HeapGraphRoot::kObjectIdsFieldNumber>(
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker.cc b/src/trace_processor/importers/proto/heap_graph_tracker.cc
index e337113..02c3ae2 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker.cc
+++ b/src/trace_processor/importers/proto/heap_graph_tracker.cc
@@ -21,6 +21,7 @@
 #include "perfetto/base/flat_set.h"
 #include "perfetto/ext/base/string_splitter.h"
 #include "perfetto/ext/base/string_utils.h"
+#include "protos/perfetto/trace/profiling/heap_graph.pbzero.h"
 #include "src/trace_processor/importers/proto/profiler_util.h"
 #include "src/trace_processor/tables/profiler_tables_py.h"
 
@@ -57,42 +58,6 @@
   }
 }
 
-base::FlatSet<ObjectTable::Id> GetChildren(TraceStorage* storage,
-                                           ObjectTable::RowReference object) {
-  auto cls_row_ref =
-      *storage->heap_graph_class_table().FindById(object.type_id());
-
-  StringId kind = cls_row_ref.kind();
-  std::optional<StringId> weakref_kind =
-      storage->string_pool().GetId("KIND_WEAK_REFERENCE");
-  std::optional<StringId> softref_kind =
-      storage->string_pool().GetId("KIND_SOFT_REFERENCE");
-  std::optional<StringId> finalizerref_kind =
-      storage->string_pool().GetId("KIND_FINALIZER_REFERENCE");
-  std::optional<StringId> phantomref_kind =
-      storage->string_pool().GetId("KIND_PHANTOM_REFERENCE");
-
-  if ((weakref_kind && kind == *weakref_kind) ||
-      (softref_kind && kind == *softref_kind) ||
-      (finalizerref_kind && kind == *finalizerref_kind) ||
-      (phantomref_kind && kind == *phantomref_kind)) {
-    // Do not follow weak / soft / finalizer / phantom references.
-    return {};
-  }
-
-  base::FlatSet<ObjectTable::Id> children;
-  ForReferenceSet(storage, object,
-                  [object, &children](ReferenceTable::RowReference ref) {
-                    PERFETTO_CHECK(ref.owner_id() == object.id());
-                    auto opt_owned = ref.owned_id();
-                    if (opt_owned) {
-                      children.insert(*opt_owned);
-                    }
-                    return true;
-                  });
-  return children;
-}
-
 struct ClassDescriptor {
   StringId name;
   std::optional<StringId> location;
@@ -176,56 +141,6 @@
 
 }  // namespace
 
-void MarkRoot(TraceStorage* storage,
-              ObjectTable::RowReference row_ref,
-              StringId type) {
-  row_ref.set_root_type(type);
-
-  // DFS to mark reachability for all children
-  std::vector<ObjectTable::RowReference> stack({row_ref});
-  while (!stack.empty()) {
-    ObjectTable::RowReference cur_node = stack.back();
-    stack.pop_back();
-
-    if (cur_node.reachable())
-      continue;
-    cur_node.set_reachable(true);
-
-    for (ObjectTable::Id child_node : GetChildren(storage, cur_node)) {
-      auto child_ref =
-          *storage->mutable_heap_graph_object_table()->FindById(child_node);
-      stack.push_back(child_ref);
-    }
-  }
-}
-
-void UpdateShortestPaths(TraceStorage* storage,
-                         ObjectTable::RowReference row_ref) {
-  // Calculate shortest distance to a GC root.
-  std::deque<std::pair<int32_t, ObjectTable::RowReference>> reachable_nodes{
-      {0, row_ref}};
-  while (!reachable_nodes.empty()) {
-    auto pair = reachable_nodes.front();
-
-    int32_t distance = pair.first;
-    ObjectTable::RowReference cur_row_ref = pair.second;
-
-    reachable_nodes.pop_front();
-    int32_t cur_distance = cur_row_ref.root_distance();
-    if (cur_distance == -1 || cur_distance > distance) {
-      cur_row_ref.set_root_distance(distance);
-
-      for (ObjectTable::Id child_node : GetChildren(storage, cur_row_ref)) {
-        auto child_row_ref =
-            *storage->mutable_heap_graph_object_table()->FindById(child_node);
-        int32_t child_distance = child_row_ref.root_distance();
-        if (child_distance == -1 || child_distance > distance + 1)
-          reachable_nodes.emplace_back(distance + 1, child_row_ref);
-      }
-    }
-  }
-}
-
 std::optional<base::StringView> GetStaticClassTypeName(base::StringView type) {
   static const base::StringView kJavaClassTemplate("java.lang.Class<");
   if (!type.empty() && type.at(type.size() - 1) == '>' &&
@@ -283,7 +198,21 @@
           "libcore.util.NativeAllocationRegistry$CleanerThunk.this$0")),
       native_size_str_id_(
           storage_->InternString("libcore.util.NativeAllocationRegistry.size")),
-      cleaner_next_str_id_(storage_->InternString("sun.misc.Cleaner.next")) {}
+      cleaner_next_str_id_(storage_->InternString("sun.misc.Cleaner.next")) {
+  for (size_t i = 0; i < root_type_string_ids_.size(); i++) {
+    auto val = static_cast<protos::pbzero::HeapGraphRoot::Type>(i);
+    auto str_view =
+        base::StringView(protos::pbzero::HeapGraphRoot_Type_Name(val));
+    root_type_string_ids_[i] = storage_->InternString(str_view);
+  }
+
+  for (size_t i = 0; i < type_kind_string_ids_.size(); i++) {
+    auto val = static_cast<protos::pbzero::HeapGraphType::Kind>(i);
+    auto str_view =
+        base::StringView(protos::pbzero::HeapGraphType_Kind_Name(val));
+    type_kind_string_ids_[i] = storage_->InternString(str_view);
+  }
+}
 
 HeapGraphTracker::SequenceState& HeapGraphTracker::GetOrCreateSequence(
     uint32_t seq_id) {
@@ -429,26 +358,27 @@
   sequence_state.interned_location_names.emplace(intern_id, strid);
 }
 
-void HeapGraphTracker::AddInternedType(uint32_t seq_id,
-                                       uint64_t intern_id,
-                                       StringId strid,
-                                       std::optional<uint64_t> location_id,
-                                       uint64_t object_size,
-                                       std::vector<uint64_t> field_name_ids,
-                                       uint64_t superclass_id,
-                                       uint64_t classloader_id,
-                                       bool no_fields,
-                                       StringId kind) {
+void HeapGraphTracker::AddInternedType(
+    uint32_t seq_id,
+    uint64_t intern_id,
+    StringId strid,
+    std::optional<uint64_t> location_id,
+    uint64_t object_size,
+    std::vector<uint64_t> field_name_ids,
+    uint64_t superclass_id,
+    uint64_t classloader_id,
+    bool no_fields,
+    protos::pbzero::HeapGraphType::Kind kind) {
   SequenceState& sequence_state = GetOrCreateSequence(seq_id);
-  sequence_state.interned_types[intern_id].name = strid;
-  sequence_state.interned_types[intern_id].location_id = location_id;
-  sequence_state.interned_types[intern_id].object_size = object_size;
-  sequence_state.interned_types[intern_id].field_name_ids =
-      std::move(field_name_ids);
-  sequence_state.interned_types[intern_id].superclass_id = superclass_id;
-  sequence_state.interned_types[intern_id].classloader_id = classloader_id;
-  sequence_state.interned_types[intern_id].no_fields = no_fields;
-  sequence_state.interned_types[intern_id].kind = kind;
+  InternedType& type = sequence_state.interned_types[intern_id];
+  type.name = strid;
+  type.location_id = location_id;
+  type.object_size = object_size;
+  type.field_name_ids = std::move(field_name_ids);
+  type.superclass_id = superclass_id;
+  type.classloader_id = classloader_id;
+  type.no_fields = no_fields;
+  type.kind = kind;
 }
 
 void HeapGraphTracker::AddInternedFieldName(uint32_t seq_id,
@@ -618,7 +548,7 @@
     }
     if (location_name)
       type_row_ref.set_location(*location_name);
-    type_row_ref.set_kind(interned_type.kind);
+    type_row_ref.set_kind(InternTypeKindString(interned_type.kind));
 
     base::StringView normalized_type =
         NormalizeTypeName(storage_->GetString(interned_type.name));
@@ -668,8 +598,9 @@
       auto it_and_success = roots_[std::make_pair(sequence_state.current_upid,
                                                   sequence_state.current_ts)]
                                 .emplace(*ptr);
-      if (it_and_success.second)
-        MarkRoot(storage_, row_ref, root.root_type);
+      if (it_and_success.second) {
+        MarkRoot(row_ref, InternRootTypeString(root.root_type));
+      }
     }
   }
 
@@ -817,9 +748,88 @@
   }
 }
 
-void FindPathFromRoot(TraceStorage* storage,
-                      ObjectTable::RowReference row_ref,
-                      PathFromRoot* path) {
+base::FlatSet<ObjectTable::Id> HeapGraphTracker::GetChildren(
+    ObjectTable::RowReference object) {
+  auto cls_row_ref =
+      *storage_->heap_graph_class_table().FindById(object.type_id());
+
+  StringId kind = cls_row_ref.kind();
+
+  if (kind == InternTypeKindString(
+                  protos::pbzero::HeapGraphType::KIND_WEAK_REFERENCE) ||
+      kind == InternTypeKindString(
+                  protos::pbzero::HeapGraphType::KIND_SOFT_REFERENCE) ||
+      kind == InternTypeKindString(
+                  protos::pbzero::HeapGraphType::KIND_FINALIZER_REFERENCE) ||
+      kind == InternTypeKindString(
+                  protos::pbzero::HeapGraphType::KIND_PHANTOM_REFERENCE)) {
+    // Do not follow weak / soft / finalizer / phantom references.
+    return {};
+  }
+
+  base::FlatSet<ObjectTable::Id> children;
+  ForReferenceSet(storage_, object,
+                  [object, &children](ReferenceTable::RowReference ref) {
+                    PERFETTO_CHECK(ref.owner_id() == object.id());
+                    auto opt_owned = ref.owned_id();
+                    if (opt_owned) {
+                      children.insert(*opt_owned);
+                    }
+                    return true;
+                  });
+  return children;
+}
+
+void HeapGraphTracker::MarkRoot(ObjectTable::RowReference row_ref,
+                                StringId type) {
+  row_ref.set_root_type(type);
+
+  // DFS to mark reachability for all children
+  std::vector<ObjectTable::RowReference> stack({row_ref});
+  while (!stack.empty()) {
+    ObjectTable::RowReference cur_node = stack.back();
+    stack.pop_back();
+
+    if (cur_node.reachable())
+      continue;
+    cur_node.set_reachable(true);
+
+    for (ObjectTable::Id child_node : GetChildren(cur_node)) {
+      auto child_ref =
+          *storage_->mutable_heap_graph_object_table()->FindById(child_node);
+      stack.push_back(child_ref);
+    }
+  }
+}
+
+void HeapGraphTracker::UpdateShortestPaths(ObjectTable::RowReference row_ref) {
+  // Calculate shortest distance to a GC root.
+  std::deque<std::pair<int32_t, ObjectTable::RowReference>> reachable_nodes{
+      {0, row_ref}};
+  while (!reachable_nodes.empty()) {
+    auto pair = reachable_nodes.front();
+
+    int32_t distance = pair.first;
+    ObjectTable::RowReference cur_row_ref = pair.second;
+
+    reachable_nodes.pop_front();
+    int32_t cur_distance = cur_row_ref.root_distance();
+    if (cur_distance == -1 || cur_distance > distance) {
+      cur_row_ref.set_root_distance(distance);
+
+      for (ObjectTable::Id child_node : GetChildren(cur_row_ref)) {
+        auto child_row_ref =
+            *storage_->mutable_heap_graph_object_table()->FindById(child_node);
+        int32_t child_distance = child_row_ref.root_distance();
+        if (child_distance == -1 || child_distance > distance + 1)
+          reachable_nodes.emplace_back(distance + 1, child_row_ref);
+      }
+    }
+  }
+}
+
+void HeapGraphTracker::FindPathFromRoot(ObjectTable::RowReference row_ref,
+                                        PathFromRoot* path) {
   // We have long retention chains (e.g. from LinkedList). If we use the stack
   // here, we risk running out of stack space. This is why we use a vector to
   // simulate the stack.
@@ -844,7 +854,7 @@
 
     ClassTable::Id type_id = object_row_ref.type_id();
 
-    auto type_row_ref = *storage->heap_graph_class_table().FindById(type_id);
+    auto type_row_ref = *storage_->heap_graph_class_table().FindById(type_id);
     std::optional<StringId> opt_class_name_id =
         type_row_ref.deobfuscated_name();
     if (!opt_class_name_id) {
@@ -854,9 +864,9 @@
     StringId class_name_id = *opt_class_name_id;
     std::optional<StringId> root_type = object_row_ref.root_type();
     if (root_type) {
-      class_name_id = storage->InternString(base::StringView(
-          storage->GetString(class_name_id).ToStdString() + " [" +
-          storage->GetString(*root_type).ToStdString() + "]"));
+      class_name_id = storage_->InternString(base::StringView(
+          storage_->GetString(class_name_id).ToStdString() + " [" +
+          storage_->GetString(*root_type).ToStdString() + "]"));
     }
     auto it = path->nodes[parent_id].children.find(class_name_id);
     if (it == path->nodes[parent_id].children.end()) {
@@ -876,15 +886,14 @@
       // size to the relevant node in the resulting tree.
       output_tree_node->size += object_row_ref.self_size();
       output_tree_node->count++;
-      base::FlatSet<ObjectTable::Id> children_set =
-          GetChildren(storage, object_row_ref);
+      base::FlatSet<ObjectTable::Id> children_set = GetChildren(object_row_ref);
       children.assign(children_set.begin(), children_set.end());
       PERFETTO_CHECK(children.size() == children_set.size());
 
       if (object_row_ref.native_size()) {
-        StringId native_class_name_id = storage->InternString(
+        StringId native_class_name_id = storage_->InternString(
             base::StringView(std::string("[native] ") +
-                             storage->GetString(class_name_id).ToStdString()));
+                             storage_->GetString(class_name_id).ToStdString()));
         std::map<StringId, size_t>::iterator native_it;
         bool inserted_new_node;
         std::tie(native_it, inserted_new_node) =
@@ -910,7 +919,7 @@
       PERFETTO_CHECK(i < children.size());
       ObjectTable::Id child = children[i];
       auto child_row_ref =
-          *storage->mutable_heap_graph_object_table()->FindById(child);
+          *storage_->mutable_heap_graph_object_table()->FindById(child);
       if (++i == children.size())
         stack.pop_back();
 
@@ -971,11 +980,11 @@
 
   // First pass to calculate shortest paths
   for (ObjectTable::RowNumber root : roots) {
-    UpdateShortestPaths(storage_, root.ToRowReference(object_table));
+    UpdateShortestPaths(root.ToRowReference(object_table));
   }
   PathFromRoot init_path;
   for (ObjectTable::RowNumber root : roots) {
-    FindPathFromRoot(storage_, root.ToRowReference(object_table), &init_path);
+    FindPathFromRoot(root.ToRowReference(object_table), &init_path);
   }
 
   std::vector<int64_t> node_to_cumulative_size(init_path.nodes.size());
@@ -1047,6 +1056,26 @@
   return false;
 }
 
+StringId HeapGraphTracker::InternRootTypeString(
+    protos::pbzero::HeapGraphRoot::Type root_type) {
+  size_t idx = static_cast<size_t>(root_type);
+  if (idx >= root_type_string_ids_.size()) {
+    idx = static_cast<size_t>(protos::pbzero::HeapGraphRoot::ROOT_UNKNOWN);
+  }
+
+  return root_type_string_ids_[idx];
+}
+
+StringId HeapGraphTracker::InternTypeKindString(
+    protos::pbzero::HeapGraphType::Kind kind) {
+  size_t idx = static_cast<size_t>(kind);
+  if (idx >= type_kind_string_ids_.size()) {
+    idx = static_cast<size_t>(protos::pbzero::HeapGraphType::KIND_UNKNOWN);
+  }
+
+  return type_kind_string_ids_[idx];
+}
+
 HeapGraphTracker::~HeapGraphTracker() = default;
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker.h b/src/trace_processor/importers/proto/heap_graph_tracker.h
index c2f184f..f5d3a1b 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker.h
+++ b/src/trace_processor/importers/proto/heap_graph_tracker.h
@@ -23,6 +23,7 @@
 #include <utility>
 #include <vector>
 
+#include "perfetto/base/flat_set.h"
 #include "perfetto/ext/base/string_view.h"
 
 #include "protos/perfetto/trace/profiling/heap_graph.pbzero.h"
@@ -55,15 +56,6 @@
   std::set<tables::HeapGraphObjectTable::Id> visited;
 };
 
-void MarkRoot(TraceStorage*,
-              tables::HeapGraphObjectTable::RowReference,
-              StringId type);
-void UpdateShortestPaths(TraceStorage* s,
-                         tables::HeapGraphObjectTable::RowReference row_ref);
-void FindPathFromRoot(TraceStorage* storage,
-                      tables::HeapGraphObjectTable::RowReference,
-                      PathFromRoot* path);
-
 std::optional<base::StringView> GetStaticClassTypeName(base::StringView type);
 size_t NumberOfArrays(base::StringView type);
 NormalizedType GetNormalizedType(base::StringView type);
@@ -89,7 +81,7 @@
   };
 
   struct SourceRoot {
-    StringId root_type;
+    protos::pbzero::HeapGraphRoot::Type root_type;
     std::vector<uint64_t> object_ids;
   };
 
@@ -114,7 +106,7 @@
                        uint64_t superclass_id,
                        uint64_t classloader_id,
                        bool no_fields,
-                       StringId kind);
+                       protos::pbzero::HeapGraphType::Kind kind);
   void AddInternedFieldName(uint32_t seq_id,
                             uint64_t intern_id,
                             base::StringView str);
@@ -162,7 +154,7 @@
     uint64_t superclass_id;
     bool no_fields;
     uint64_t classloader_id;
-    StringId kind;
+    protos::pbzero::HeapGraphType::Kind kind;
   };
   struct SequenceState {
     UniquePid current_upid = 0;
@@ -218,6 +210,8 @@
   InternedType* GetSuperClass(SequenceState* sequence_state,
                               const InternedType* current_type);
   bool IsTruncated(UniquePid upid, int64_t ts);
+  StringId InternRootTypeString(protos::pbzero::HeapGraphRoot::Type);
+  StringId InternTypeKindString(protos::pbzero::HeapGraphType::Kind);
 
   // Returns the object pointed to by `field` in `obj`.
   std::optional<tables::HeapGraphObjectTable::Id> GetReferenceByFieldName(
@@ -231,6 +225,13 @@
   // all the other tables have been fully populated.
   void PopulateNativeSize(const SequenceState& seq);
 
+  base::FlatSet<tables::HeapGraphObjectTable::Id> GetChildren(
+      tables::HeapGraphObjectTable::RowReference);
+  void MarkRoot(tables::HeapGraphObjectTable::RowReference, StringId type);
+  void UpdateShortestPaths(tables::HeapGraphObjectTable::RowReference row_ref);
+  void FindPathFromRoot(tables::HeapGraphObjectTable::RowReference,
+                        PathFromRoot* path);
+
   TraceStorage* const storage_;
   std::map<uint32_t, SequenceState> sequence_state_;
 
@@ -253,6 +254,16 @@
   StringId cleaner_thunk_this0_str_id_;
   StringId native_size_str_id_;
   StringId cleaner_next_str_id_;
+
+  std::array<StringId, 15> root_type_string_ids_ = {};
+  static_assert(protos::pbzero::HeapGraphRoot_Type_MIN == 0);
+  static_assert(protos::pbzero::HeapGraphRoot_Type_MAX + 1 ==
+                std::tuple_size<decltype(root_type_string_ids_)>{});
+
+  std::array<StringId, 12> type_kind_string_ids_ = {};
+  static_assert(protos::pbzero::HeapGraphType_Kind_MIN == 0);
+  static_assert(protos::pbzero::HeapGraphType_Kind_MAX + 1 ==
+                std::tuple_size<decltype(type_kind_string_ids_)>{});
 };
 
 }  // namespace trace_processor
diff --git a/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc b/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
index 15be4ac..6ae4f09 100644
--- a/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
+++ b/src/trace_processor/importers/proto/heap_graph_tracker_unittest.cc
@@ -67,8 +67,6 @@
 
   HeapGraphTracker tracker(context.storage.get());
 
-  StringPool::Id normal_kind = context.storage->InternString("KIND_NORMAL");
-
   constexpr uint64_t kLocation = 0;
   tracker.AddInternedLocationName(kSeqId, kLocation,
                                   context.storage->InternString("location"));
@@ -94,34 +92,34 @@
       kSeqId, kTypeBitmap,
       context.storage->InternString("android.graphics.Bitmap"), kLocation,
       /*object_size=*/0,
-      /*reference_field_name_ids=*/{}, /*superclass_id=*/0,
-      /*classloader_id=*/0, /*no_reference_fields=*/false,
-      /*kind=*/normal_kind);
+      /*field_name_ids=*/{}, /*superclass_id=*/0,
+      /*classloader_id=*/0, /*no_fields=*/false,
+      protos::pbzero::HeapGraphType::KIND_NORMAL);
 
-  tracker.AddInternedType(
-      kSeqId, kTypeCleaner, context.storage->InternString("sun.misc.Cleaner"),
-      kLocation, /*object_size=*/0,
-      /*reference_field_name_ids=*/{kReferent, kThunk, kNext},
-      /*superclass_id=*/0,
-      /*classloader_id=*/0, /*no_reference_fields=*/false,
-      /*kind=*/normal_kind);
+  tracker.AddInternedType(kSeqId, kTypeCleaner,
+                          context.storage->InternString("sun.misc.Cleaner"),
+                          kLocation, /*object_size=*/0,
+                          /*field_name_ids=*/{kReferent, kThunk, kNext},
+                          /*superclass_id=*/0,
+                          /*classloader_id=*/0, /*no_fields=*/false,
+                          protos::pbzero::HeapGraphType::KIND_NORMAL);
 
   tracker.AddInternedType(
       kSeqId, kTypeCleanerThunk,
       context.storage->InternString(
           "libcore.util.NativeAllocationRegistry$CleanerThunk"),
       kLocation, /*object_size=*/0,
-      /*reference_field_name_ids=*/{kThis0}, /*superclass_id=*/0,
-      /*classloader_id=*/0, /*no_reference_fields=*/false,
-      /*kind=*/normal_kind);
+      /*field_name_ids=*/{kThis0}, /*superclass_id=*/0,
+      /*classloader_id=*/0, /*no_fields=*/false,
+      protos::pbzero::HeapGraphType::KIND_NORMAL);
 
   tracker.AddInternedType(
       kSeqId, kTypeNativeAllocationRegistry,
       context.storage->InternString("libcore.util.NativeAllocationRegistry"),
       kLocation, /*object_size=*/0,
-      /*reference_field_name_ids=*/{}, /*superclass_id=*/0,
-      /*classloader_id=*/0, /*no_reference_fields=*/false,
-      /*kind=*/normal_kind);
+      /*field_name_ids=*/{}, /*superclass_id=*/0,
+      /*classloader_id=*/0, /*no_fields=*/false,
+      protos::pbzero::HeapGraphType::KIND_NORMAL);
 
   enum Objects : uint64_t {
     kObjBitmap = 1,
@@ -226,9 +224,6 @@
   StringPool::Id b = context.storage->InternString("B");
   StringPool::Id weak_ref = context.storage->InternString("WeakReference");
 
-  StringPool::Id normal_kind = context.storage->InternString("KIND_NORMAL");
-  StringPool::Id weak_ref_kind =
-      context.storage->InternString("KIND_WEAK_REFERENCE");
   tracker.AddInternedFieldName(kSeqId, kField, field);
 
   tracker.AddInternedLocationName(kSeqId, kLocation,
@@ -236,24 +231,24 @@
   tracker.AddInternedType(kSeqId, kX, x, kLocation, /*object_size=*/0,
                           /*field_name_ids=*/{}, /*superclass_id=*/0,
                           /*classloader_id=*/0, /*no_fields=*/false,
-                          /*kind=*/normal_kind);
+                          protos::pbzero::HeapGraphType::KIND_NORMAL);
   tracker.AddInternedType(kSeqId, kY, y, kLocation, /*object_size=*/0,
                           /*field_name_ids=*/{}, /*superclass_id=*/0,
                           /*classloader_id=*/0, /*no_fields=*/false,
-                          /*kind=*/normal_kind);
+                          protos::pbzero::HeapGraphType::KIND_NORMAL);
   tracker.AddInternedType(kSeqId, kA, a, kLocation, /*object_size=*/0,
                           /*field_name_ids=*/{}, /*superclass_id=*/0,
                           /*classloader_id=*/0, /*no_fields=*/false,
-                          /*kind=*/normal_kind);
+                          protos::pbzero::HeapGraphType::KIND_NORMAL);
   tracker.AddInternedType(kSeqId, kB, b, kLocation, /*object_size=*/0,
                           /*field_name_ids=*/{}, /*superclass_id=*/0,
                           /*classloader_id=*/0, /*no_fields=*/false,
-                          /*kind=*/normal_kind);
+                          protos::pbzero::HeapGraphType::KIND_NORMAL);
   tracker.AddInternedType(kSeqId, kWeakRef, weak_ref, kLocation,
                           /*object_size=*/0,
                           /*field_name_ids=*/{}, /*superclass_id=*/0,
                           /*classloader_id=*/0, /*no_fields=*/false,
-                          /*kind=*/weak_ref_kind);
+                          protos::pbzero::HeapGraphType::KIND_WEAK_REFERENCE);
   {
     HeapGraphTracker::SourceObject obj;
     obj.object_id = 999;
@@ -312,7 +307,7 @@
   }
 
   HeapGraphTracker::SourceRoot root;
-  root.root_type = context.storage->InternString("ROOT");
+  root.root_type = protos::pbzero::HeapGraphRoot::ROOT_UNKNOWN;
   root.object_ids.emplace_back(1);
   root.object_ids.emplace_back(999);
   tracker.AddRoot(kSeqId, kPid, kTimestamp, root);
diff --git a/src/trace_processor/metrics/sql/android/android_sysui_notifications_blocking_calls_metric.sql b/src/trace_processor/metrics/sql/android/android_sysui_notifications_blocking_calls_metric.sql
index e03dca9..ec17f42 100644
--- a/src/trace_processor/metrics/sql/android/android_sysui_notifications_blocking_calls_metric.sql
+++ b/src/trace_processor/metrics/sql/android/android_sysui_notifications_blocking_calls_metric.sql
@@ -32,6 +32,8 @@
        s.name GLOB 'NotificationStackScrollLayout#onMeasure'
     OR s.name GLOB 'NotificationToplineView#onMeasure'
     OR s.name GLOB 'ExpNotRow#*'
+    OR s.name GLOB 'NotificationShadeWindowView#onMeasure'
+    OR s.name GLOB 'ImageFloatingTextView#onMeasure'
 )
 GROUP BY s.name;
 
diff --git a/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.out b/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.out
index 5b69129..a30c810 100644
--- a/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.out
+++ b/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.out
@@ -14,6 +14,20 @@
         min_dur_ns: 10000000
     }
     blocking_calls {
+        name: "ImageFloatingTextView#onMeasure"
+        cnt: 1
+        total_dur_ns: 10000000
+        max_dur_ns: 10000000
+        min_dur_ns: 10000000
+    }
+    blocking_calls {
+        name: "NotificationShadeWindowView#onMeasure"
+        cnt: 1
+        total_dur_ns: 10000000
+        max_dur_ns: 10000000
+        min_dur_ns: 10000000
+    }
+    blocking_calls {
         name: "NotificationStackScrollLayout#onMeasure"
         cnt: 1
         total_dur_ns: 10000000
diff --git a/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.py b/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.py
index c5df8b4..cab8c8f 100644
--- a/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.py
+++ b/test/trace_processor/diff_tests/android/android_sysui_notifications_blocking_calls_metric.py
@@ -24,7 +24,8 @@
 # List of blocking calls
 blocking_call_names = [
     'NotificationStackScrollLayout#onMeasure', 'ExpNotRow#onMeasure(MessagingStyle)',
-    'ExpNotRow#onMeasure(BigTextStyle)',
+    'ExpNotRow#onMeasure(BigTextStyle)', 'NotificationShadeWindowView#onMeasure',
+    'ImageFloatingTextView#onMeasure',
     'Should not be in the metric'
 ]