1 files changed, 11 insertions, 2 deletions
diff --git a/compiler/optimizing/code_generator.h b/compiler/optimizing/code_generator.h
index 4b11e7c699..2e2c3c00af 100644
--- a/compiler/optimizing/code_generator.h
+++ b/compiler/optimizing/code_generator.h
@@ -34,6 +34,7 @@
 #include "stack_map_stream.h"
 #include "string_reference.h"
 #include "utils/label.h"
+#include "utils/type_reference.h"
 
 namespace art {
 
@@ -343,7 +344,7 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
   void BuildStackMaps(MemoryRegion region, const DexFile::CodeItem& code_item);
   size_t ComputeStackMapsSize();
   size_t GetNumberOfJitRoots() const {
-    return jit_string_roots_.size();
+    return jit_string_roots_.size() + jit_class_roots_.size();
   }
 
   // Fills the `literals` array with literals collected during code generation.
@@ -611,6 +612,8 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
         block_order_(nullptr),
         jit_string_roots_(StringReferenceValueComparator(),
                           graph->GetArena()->Adapter(kArenaAllocCodeGenerator)),
+        jit_class_roots_(TypeReferenceValueComparator(),
+                         graph->GetArena()->Adapter(kArenaAllocCodeGenerator)),
         disasm_info_(nullptr),
         stats_(stats),
         graph_(graph),
@@ -681,6 +684,7 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
   virtual void EmitJitRootPatches(uint8_t* code ATTRIBUTE_UNUSED,
                                   const uint8_t* roots_data ATTRIBUTE_UNUSED) {
     DCHECK_EQ(jit_string_roots_.size(), 0u);
+    DCHECK_EQ(jit_class_roots_.size(), 0u);
   }
 
   // Frame size required for this method.
@@ -711,7 +715,12 @@ class CodeGenerator : public DeletableArenaObject<kArenaAllocCodeGenerator> {
   // Maps a StringReference (dex_file, string_index) to the index in the literal table.
   // Entries are intially added with a 0 index, and `EmitJitRoots` will compute all the
   // indices.
-  ArenaSafeMap<StringReference, size_t, StringReferenceValueComparator> jit_string_roots_;
+  ArenaSafeMap<StringReference, uint32_t, StringReferenceValueComparator> jit_string_roots_;
+
+  // Maps a ClassReference (dex_file, type_index) to the index in the literal table.
+  // Entries are intially added with a pointer in the handle zone, and `EmitJitRoots`
+  // will compute all the indices.
+  ArenaSafeMap<TypeReference, uint64_t, TypeReferenceValueComparator> jit_class_roots_;
 
   DisassemblyInformation* disasm_info_;