1 files changed, 17 insertions, 3 deletions
diff --git a/compiler/dex/post_opt_passes.h b/compiler/dex/post_opt_passes.h
index 1ab8625..e9fa0eb 100644
--- a/compiler/dex/post_opt_passes.h
+++ b/compiler/dex/post_opt_passes.h
@@ -263,12 +263,19 @@ class PerformInitRegLocations : public PassMEMirSsaRep {
 };
 
 /**
- * @class TypeInference
+ * @class TypeInferencePass
  * @brief Type inference pass.
  */
-class TypeInference : public PassMEMirSsaRep {
+class TypeInferencePass : public PassMEMirSsaRep {
  public:
-  TypeInference() : PassMEMirSsaRep("TypeInference", kRepeatingPreOrderDFSTraversal) {
+  TypeInferencePass() : PassMEMirSsaRep("TypeInference", kRepeatingPreOrderDFSTraversal) {
+  }
+
+  void Start(PassDataHolder* data) const {
+    DCHECK(data != nullptr);
+    CompilationUnit* c_unit = down_cast<PassMEDataHolder*>(data)->c_unit;
+    DCHECK(c_unit != nullptr);
+    c_unit->mir_graph->InferTypesStart();
   }
 
   bool Worker(PassDataHolder* data) const {
@@ -280,6 +287,13 @@ class TypeInference : public PassMEMirSsaRep {
     DCHECK(bb != nullptr);
     return c_unit->mir_graph->InferTypes(bb);
   }
+
+  void End(PassDataHolder* data) const {
+    DCHECK(data != nullptr);
+    CompilationUnit* c_unit = down_cast<PassMEDataHolder*>(data)->c_unit;
+    DCHECK(c_unit != nullptr);
+    c_unit->mir_graph.get()->InferTypesEnd();
+  }
 };
 
 /**