moves the TypeRecorder ownership to the ICInfo

src/asm_writing/icinfo.cpp

@@ -23,6 +23,7 @@
 #include "asm_writing/assembler.h"
 #include "asm_writing/mc_writer.h"
 #include "codegen/patchpoints.h"
+#include "codegen/type_recording.h"
 #include "codegen/unwinding.h"
 #include "core/common.h"
 #include "core/options.h"
@@ -339,14 +340,12 @@ static llvm::DenseMap<AST*, ICInfo*> ics_by_ast_node;
 
 ICInfo::ICInfo(void* start_addr, void* slowpath_rtn_addr, void* continue_addr, StackInfo stack_info, int size,
                llvm::CallingConv::ID calling_conv, LiveOutSet _live_outs, assembler::GenericRegister return_register,
-               TypeRecorder* type_recorder, std::vector<Location> ic_global_decref_locations,
-               assembler::RegisterSet allocatable_registers)
+               std::vector<Location> ic_global_decref_locations, assembler::RegisterSet allocatable_registers)
     : next_slot_to_try(0),
       stack_info(stack_info),
       calling_conv(calling_conv),
       live_outs(std::move(_live_outs)),
       return_register(return_register),
-      type_recorder(type_recorder),
       retry_in(0),
       retry_backoff(1),
       times_rewritten(0),
@@ -416,7 +415,7 @@ std::unique_ptr<ICInfo> registerCompiledPatchpoint(uint8_t* start_addr, uint8_t*
 
     ICInfo* icinfo
         = new ICInfo(start_addr, slowpath_rtn_addr, continue_addr, stack_info, ic->size, ic->getCallingConvention(),
-                     std::move(live_outs), return_register, ic->type_recorder, decref_info, ic->allocatable_regs);
+                     std::move(live_outs), return_register, decref_info, ic->allocatable_regs);
 
     assert(!ics_by_return_addr.count(slowpath_rtn_addr));
     ics_by_return_addr[slowpath_rtn_addr] = icinfo;
@@ -487,9 +486,10 @@ ICInfo* ICInfo::getICInfoForNode(AST* node) {
         return it->second;
     return NULL;
 }
-void ICInfo::associateNodeWithICInfo(AST* node) {
+void ICInfo::associateNodeWithICInfo(AST* node, std::unique_ptr<TypeRecorder> type_recorder) {
     assert(!this->node);
     this->node = node;
+    this->type_recorder = std::move(type_recorder);
     ics_by_ast_node[node] = this;
 }
 void ICInfo::appendDecrefInfosTo(std::vector<DecrefInfo>& dest_decref_infos) {

src/asm_writing/icinfo.h

@@ -93,7 +93,7 @@ private:
     const llvm::CallingConv::ID calling_conv;
     LiveOutSet live_outs;
     const assembler::GenericRegister return_register;
-    TypeRecorder* const type_recorder;
+    std::unique_ptr<TypeRecorder> type_recorder;
     int retry_in, retry_backoff;
     int times_rewritten;
     assembler::RegisterSet allocatable_registers;
@@ -113,7 +113,7 @@ private:
 public:
     ICInfo(void* start_addr, void* slowpath_rtn_addr, void* continue_addr, StackInfo stack_info, int size,
            llvm::CallingConv::ID calling_conv, LiveOutSet live_outs, assembler::GenericRegister return_register,
-           TypeRecorder* type_recorder, std::vector<Location> ic_global_decref_locations,
+           std::vector<Location> ic_global_decref_locations,
            assembler::RegisterSet allocatable_registers = assembler::RegisterSet::stdAllocatable());
     ~ICInfo();
     void* const start_addr, *const slowpath_rtn_addr, *const continue_addr;
@@ -123,6 +123,7 @@ public:
 
     llvm::CallingConv::ID getCallingConvention() { return calling_conv; }
     const LiveOutSet& getLiveOuts() { return live_outs; }
+    TypeRecorder* getTypeRecorder() { return type_recorder.get(); }
 
     std::unique_ptr<ICSlotRewrite> startRewrite(const char* debug_name);
     void invalidate(ICSlotInfo* entry);
@@ -145,7 +146,7 @@ public:
     friend class ICSlotRewrite;
 
     static ICInfo* getICInfoForNode(AST* node);
-    void associateNodeWithICInfo(AST* node);
+    void associateNodeWithICInfo(AST* node, std::unique_ptr<TypeRecorder> type_recorder);
 
     void appendDecrefInfosTo(std::vector<DecrefInfo>& dest_decref_infos);
 };
@@ -179,7 +180,7 @@ public:
     ICInfo* getICInfo() { return ic_entry->ic; }
     int getSlotSize() { return ic_entry->size; }
     uint8_t* getSlotStart() { return ic_entry->start_addr; }
-    TypeRecorder* getTypeRecorder() { return getICInfo()->type_recorder; }
+    TypeRecorder* getTypeRecorder() { return getICInfo()->getTypeRecorder(); }
     assembler::GenericRegister returnRegister() { return getICInfo()->return_register; }
     int getScratchSize() { return getICInfo()->stack_info.scratch_size; }
     int getScratchRspOffset() {

src/asm_writing/rewriter.h

@@ -179,6 +179,7 @@ public:
     RewriterVar* cmp(AST_TYPE::AST_TYPE cmp_type, RewriterVar* other, Location loc = Location::any());
     RewriterVar* toBool(Location loc = Location::any());
 
+    RefType getType() const { return reftype; }
     RewriterVar* setType(RefType type);
     RewriterVar* setNullable(bool nullable) {
         this->nullable = nullable;

src/codegen/baseline_jit.h

@@ -251,6 +251,7 @@ private:
         StackInfo stack_info;
         AST* node;
         std::vector<Location> decref_infos;
+        std::unique_ptr<TypeRecorder> type_recorder;
     };
 
     llvm::SmallVector<PPInfo, 8> pp_infos;
@@ -349,13 +350,13 @@ private:
     RewriterVar* emitCallWithAllocatedArgs(void* func_addr, const llvm::ArrayRef<RewriterVar*> args,
                                            const llvm::ArrayRef<RewriterVar*> additional_uses);
     std::pair<RewriterVar*, RewriterAction*> emitPPCall(void* func_addr, llvm::ArrayRef<RewriterVar*> args,
-                                                        unsigned short pp_size, AST* ast_node = NULL,
-                                                        TypeRecorder* type_recorder = NULL,
+                                                        unsigned short pp_size, bool should_record_type = false,
+                                                        AST* ast_node = NULL,
                                                         llvm::ArrayRef<RewriterVar*> additional_uses = {});
 
     static void assertNameDefinedHelper(const char* id);
-    static Box* callattrHelper(Box* obj, BoxedString* attr, CallattrFlags flags, TypeRecorder* type_recorder,
-                               Box** args, std::vector<BoxedString*>* keyword_names);
+    static Box* callattrHelper(Box* obj, BoxedString* attr, CallattrFlags flags, Box** args,
+                               std::vector<BoxedString*>* keyword_names);
     static Box* createDictHelper(uint64_t num, Box** keys, Box** values);
     static Box* createListHelper(uint64_t num, Box** data);
     static Box* createSetHelper(uint64_t num, Box** data);
@@ -364,15 +365,14 @@ private:
     static BORROWED(Box*) hasnextHelper(Box* b);
     static BORROWED(Box*) nonzeroHelper(Box* b);
     static BORROWED(Box*) notHelper(Box* b);
-    static Box* runtimeCallHelper(Box* obj, ArgPassSpec argspec, TypeRecorder* type_recorder, Box** args,
-                                  std::vector<BoxedString*>* keyword_names);
+    static Box* runtimeCallHelper(Box* obj, ArgPassSpec argspec, Box** args, std::vector<BoxedString*>* keyword_names);
 
     void _emitGetLocal(RewriterVar* val_var, const char* name);
     void _emitJump(CFGBlock* b, RewriterVar* block_next, ExitInfo& exit_info);
     void _emitOSRPoint();
     void _emitPPCall(RewriterVar* result, void* func_addr, llvm::ArrayRef<RewriterVar*> args, unsigned short pp_size,
                      AST* ast_node, llvm::ArrayRef<RewriterVar*> vars_to_bump);
-    void _emitRecordType(RewriterVar* type_recorder_var, RewriterVar* obj_cls_var);
+    void _emitRecordType(RewriterVar* obj_cls_var);
     void _emitReturn(RewriterVar* v);
     void _emitSideExit(STOLEN(RewriterVar*) var, RewriterVar* val_constant, CFGBlock* next_block,
                        RewriterVar* false_path);

src/codegen/patchpoints.cpp

@@ -45,10 +45,8 @@ int ICSetupInfo::totalSize() const {
 static std::vector<std::pair<std::unique_ptr<PatchpointInfo>, void* /* addr of func to call */>> new_patchpoints;
 
 std::unique_ptr<ICSetupInfo> ICSetupInfo::initialize(bool has_return_value, int size, ICType type,
-                                                     TypeRecorder* type_recorder,
                                                      assembler::RegisterSet allocatable_regs) {
-    auto rtn
-        = std::unique_ptr<ICSetupInfo>(new ICSetupInfo(type, size, has_return_value, type_recorder, allocatable_regs));
+    auto rtn = std::unique_ptr<ICSetupInfo>(new ICSetupInfo(type, size, has_return_value, allocatable_regs));
 
 
     // We use size == CALL_ONLY_SIZE to imply that the call isn't patchable
@@ -431,7 +429,7 @@ std::unique_ptr<ICSetupInfo> createHasnextIC() {
 }
 
 std::unique_ptr<ICSetupInfo> createDeoptIC() {
-    return ICSetupInfo::initialize(true, 0, ICSetupInfo::Deopt, NULL);
+    return ICSetupInfo::initialize(true, 0, ICSetupInfo::Deopt);
 }
 
 } // namespace pyston

src/codegen/patchpoints.h

@@ -65,19 +65,13 @@ public:
     };
 
 private:
-    ICSetupInfo(ICType type, int size, bool has_return_value, TypeRecorder* type_recorder,
-                assembler::RegisterSet allocatable_regs)
-        : type(type),
-          size(size),
-          has_return_value(has_return_value),
-          type_recorder(type_recorder),
-          allocatable_regs(allocatable_regs) {}
+    ICSetupInfo(ICType type, int size, bool has_return_value, assembler::RegisterSet allocatable_regs)
+        : type(type), size(size), has_return_value(has_return_value), allocatable_regs(allocatable_regs) {}
 
 public:
     const ICType type;
     const int size;
     const bool has_return_value;
-    TypeRecorder* const type_recorder;
     assembler::RegisterSet allocatable_regs;
 
     int totalSize() const;
@@ -97,9 +91,9 @@ public:
         return llvm::CallingConv::C;
     }
 
-    static std::unique_ptr<ICSetupInfo>
-    initialize(bool has_return_value, int size, ICType type, TypeRecorder* type_recorder = NULL,
-               assembler::RegisterSet allocatable_regs = assembler::RegisterSet::stdAllocatable());
+    static std::unique_ptr<ICSetupInfo> initialize(bool has_return_value, int size, ICType type,
+                                                   assembler::RegisterSet allocatable_regs
+                                                   = assembler::RegisterSet::stdAllocatable());
 };
 
 struct PatchpointInfo {

src/codegen/type_recording.cpp

@@ -14,21 +14,12 @@
 
 #include "codegen/type_recording.h"
 
-#include <unordered_map>
-
+#include "asm_writing/icinfo.h"
 #include "core/options.h"
 #include "core/types.h"
 
 namespace pyston {
 
-static std::unordered_map<AST*, std::unique_ptr<TypeRecorder>> type_recorders;
-TypeRecorder* getTypeRecorderForNode(AST* node) {
-    std::unique_ptr<TypeRecorder>& r = type_recorders[node];
-    if (r == NULL)
-        r = llvm::make_unique<TypeRecorder>();
-    return r.get();
-}
-
 Box* recordType(TypeRecorder* self, Box* obj) {
     // The baseline JIT directly generates machine code for this function inside JitFragmentWriter::_emitRecordType.
     // When changing this function one has to also change the bjit code.
@@ -52,11 +43,11 @@ Box* recordType(TypeRecorder* self, Box* obj) {
 }
 
 BoxedClass* predictClassFor(AST* node) {
-    auto it = type_recorders.find(node);
-    if (it == type_recorders.end())
+    ICInfo* ic = ICInfo::getICInfoForNode(node);
+    if (!ic || !ic->getTypeRecorder())
         return NULL;
 
-    return it->second->predict();
+    return ic->getTypeRecorder()->predict();
 }
 
 BoxedClass* TypeRecorder::predict() {

src/codegen/type_recording.h

@@ -44,8 +44,6 @@ public:
     friend Box* recordType(TypeRecorder*, Box*);
 };
 
-TypeRecorder* getTypeRecorderForNode(AST* node);
-
 BoxedClass* predictClassFor(AST* node);
 }
 

src/runtime/ics.cpp

@@ -249,8 +249,7 @@ RuntimeIC::RuntimeIC(void* func_addr, int total_size) {
 
         // printf("Allocated runtime IC at %p\n", addr);
 
-        std::unique_ptr<ICSetupInfo> setup_info(
-            ICSetupInfo::initialize(true, patchable_size, ICSetupInfo::Generic, NULL));
+        std::unique_ptr<ICSetupInfo> setup_info(ICSetupInfo::initialize(true, patchable_size, ICSetupInfo::Generic));
         uint8_t* pp_start = (uint8_t*)addr + PROLOGUE_SIZE;
         uint8_t* pp_end = pp_start + patchable_size + CALL_SIZE;