Merge pull request #459 from kmod/osr_analysis

Fix analysis issue that virtualenv was running into

src/analysis/fpc.h

@@ -45,20 +45,23 @@ public:
 };
 
 template <typename T>
-typename BBAnalyzer<T>::AllMap computeFixedPoint(CFG* cfg, const BBAnalyzer<T>& analyzer, bool reverse) {
+void computeFixedPoint(typename BBAnalyzer<T>::Map&& initial_map, CFGBlock* initial_block,
+                       const BBAnalyzer<T>& analyzer, bool reverse, typename BBAnalyzer<T>::AllMap& starting_states,
+                       typename BBAnalyzer<T>::AllMap& ending_states) {
     assert(!reverse);
 
     typedef typename BBAnalyzer<T>::Map Map;
     typedef typename BBAnalyzer<T>::AllMap AllMap;
-    AllMap starting_states;
-    AllMap ending_states;
+
+    assert(!starting_states.size());
+    assert(!ending_states.size());
 
     llvm::SmallPtrSet<CFGBlock*, 32> in_queue;
     std::priority_queue<CFGBlock*, llvm::SmallVector<CFGBlock*, 32>, CFGBlockMinIndex> q;
 
-    starting_states.insert(make_pair(cfg->getStartingBlock(), Map()));
-    q.push(cfg->getStartingBlock());
-    in_queue.insert(cfg->getStartingBlock());
+    starting_states.insert(make_pair(initial_block, std::move(initial_map)));
+    q.push(initial_block);
+    in_queue.insert(initial_block);
 
     int num_evaluations = 0;
     while (!q.empty()) {
@@ -124,12 +127,9 @@ typename BBAnalyzer<T>::AllMap computeFixedPoint(CFG* cfg, const BBAnalyzer<T>&
     }
 
     if (VERBOSITY("analysis")) {
-        printf("%ld BBs, %d evaluations = %.1f evaluations/block\n", cfg->blocks.size(), num_evaluations,
-               1.0 * num_evaluations / cfg->blocks.size());
+        printf("%d BBs, %d evaluations = %.1f evaluations/block\n", starting_states.size(), num_evaluations,
+               1.0 * num_evaluations / starting_states.size());
     }
-
-
-    return ending_states;
 }
 }
 

src/analysis/function_analysis.cpp

@@ -21,9 +21,11 @@
 
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/StringSet.h"
 
 #include "analysis/fpc.h"
 #include "analysis/scoping_analysis.h"
+#include "codegen/osrentry.h"
 #include "core/ast.h"
 #include "core/cfg.h"
 #include "core/common.h"
@@ -221,13 +223,10 @@ class DefinednessBBAnalyzer : public BBAnalyzer<DefinednessAnalysis::DefinitionL
 private:
     typedef DefinednessAnalysis::DefinitionLevel DefinitionLevel;
 
-    CFG* cfg;
     ScopeInfo* scope_info;
-    const ParamNames& arg_names;
 
 public:
-    DefinednessBBAnalyzer(CFG* cfg, ScopeInfo* scope_info, const ParamNames& arg_names)
-        : cfg(cfg), scope_info(scope_info), arg_names(arg_names) {}
+    DefinednessBBAnalyzer(ScopeInfo* scope_info) : scope_info(scope_info) {}
 
     virtual DefinitionLevel merge(DefinitionLevel from, DefinitionLevel into) const {
         assert(from != DefinednessAnalysis::Undefined);
@@ -348,15 +347,6 @@ public:
 void DefinednessBBAnalyzer::processBB(Map& starting, CFGBlock* block) const {
     DefinednessVisitor visitor(starting);
 
-    if (block == cfg->getStartingBlock()) {
-        for (auto e : arg_names.args)
-            visitor._doSet(scope_info->internString(e));
-        if (arg_names.vararg.size())
-            visitor._doSet(scope_info->internString(arg_names.vararg));
-        if (arg_names.kwarg.size())
-            visitor._doSet(scope_info->internString(arg_names.kwarg));
-    }
-
     for (int i = 0; i < block->body.size(); i++) {
         block->body[i]->accept(&visitor);
     }
@@ -369,20 +359,23 @@ void DefinednessBBAnalyzer::processBB(Map& starting, CFGBlock* block) const {
     }
 }
 
-DefinednessAnalysis::DefinednessAnalysis(const ParamNames& arg_names, CFG* cfg, ScopeInfo* scope_info)
-    : scope_info(scope_info) {
+void DefinednessAnalysis::run(llvm::DenseMap<InternedString, DefinednessAnalysis::DefinitionLevel>&& initial_map,
+                              CFGBlock* initial_block, ScopeInfo* scope_info) {
     Timer _t("DefinednessAnalysis()", 10);
 
-    results = computeFixedPoint(cfg, DefinednessBBAnalyzer(cfg, scope_info, arg_names), false);
+    // Don't run this twice:
+    assert(!defined_at_end.size());
 
-    for (const auto& p : results) {
-        RequiredSet& required = defined_at_end[p.first];
+    computeFixedPoint(std::move(initial_map), initial_block, DefinednessBBAnalyzer(scope_info), false,
+                      defined_at_beginning, defined_at_end);
+
+    for (const auto& p : defined_at_end) {
+        RequiredSet& required = defined_at_end_sets[p.first];
         for (const auto& p2 : p.second) {
             ScopeInfo::VarScopeType vst = scope_info->getScopeTypeOfName(p2.first);
             if (vst == ScopeInfo::VarScopeType::GLOBAL || vst == ScopeInfo::VarScopeType::NAME)
                 continue;
 
-            // printf("%d %s %d\n", p.first->idx, p2.first.c_str(), p2.second);
             required.insert(p2.first);
         }
     }
@@ -392,24 +385,45 @@ DefinednessAnalysis::DefinednessAnalysis(const ParamNames& arg_names, CFG* cfg,
 }
 
 DefinednessAnalysis::DefinitionLevel DefinednessAnalysis::isDefinedAtEnd(InternedString name, CFGBlock* block) {
-    auto& map = results[block];
+    assert(defined_at_end.count(block));
+    auto& map = defined_at_end[block];
     if (map.count(name) == 0)
         return Undefined;
     return map[name];
 }
 
 const DefinednessAnalysis::RequiredSet& DefinednessAnalysis::getDefinedNamesAtEnd(CFGBlock* block) {
-    return defined_at_end[block];
+    assert(defined_at_end_sets.count(block));
+    return defined_at_end_sets[block];
 }
 
-PhiAnalysis::PhiAnalysis(const ParamNames& arg_names, CFG* cfg, LivenessAnalysis* liveness, ScopeInfo* scope_info)
-    : definedness(arg_names, cfg, scope_info), liveness(liveness) {
+PhiAnalysis::PhiAnalysis(llvm::DenseMap<InternedString, DefinednessAnalysis::DefinitionLevel>&& initial_map,
+                         CFGBlock* initial_block, bool initials_need_phis, LivenessAnalysis* liveness,
+                         ScopeInfo* scope_info)
+    : definedness(), liveness(liveness) {
     Timer _t("PhiAnalysis()", 10);
 
-    for (CFGBlock* block : cfg->blocks) {
+    // I think this should always be the case -- if we're going to generate phis for the initial block,
+    // then we should include the initial arguments as an extra entry point.
+    assert(initials_need_phis == (initial_block->predecessors.size() > 0));
+
+    definedness.run(std::move(initial_map), initial_block, scope_info);
+
+    for (const auto& p : definedness.defined_at_end) {
+        CFGBlock* block = p.first;
         RequiredSet& required = required_phis[block];
-        if (block->predecessors.size() > 1) {
+
+        int npred = 0;
+        for (CFGBlock* pred : block->predecessors) {
+            if (definedness.defined_at_end.count(pred))
+                npred++;
+        }
+
+        if (npred > 1 || (initials_need_phis && block == initial_block)) {
             for (CFGBlock* pred : block->predecessors) {
+                if (!definedness.defined_at_end.count(pred))
+                    continue;
+
                 const RequiredSet& defined = definedness.getDefinedNamesAtEnd(pred);
                 for (const auto& s : defined) {
                     if (required.count(s) == 0 && liveness->isLiveAtEnd(s, pred)) {
@@ -430,15 +444,18 @@ const PhiAnalysis::RequiredSet& PhiAnalysis::getAllRequiredAfter(CFGBlock* block
     static RequiredSet empty;
     if (block->successors.size() == 0)
         return empty;
+    assert(required_phis.count(block->successors[0]));
     return required_phis[block->successors[0]];
 }
 
 const PhiAnalysis::RequiredSet& PhiAnalysis::getAllRequiredFor(CFGBlock* block) {
+    assert(required_phis.count(block));
     return required_phis[block];
 }
 
 bool PhiAnalysis::isRequired(InternedString name, CFGBlock* block) {
     assert(!startswith(name.str(), "!"));
+    assert(required_phis.count(block));
     return required_phis[block].count(name) != 0;
 }
 
@@ -456,21 +473,18 @@ bool PhiAnalysis::isRequiredAfter(InternedString name, CFGBlock* block) {
 bool PhiAnalysis::isPotentiallyUndefinedAfter(InternedString name, CFGBlock* block) {
     assert(!startswith(name.str(), "!"));
 
-    if (block->successors.size() != 1)
-        return false;
-
-    return isPotentiallyUndefinedAt(name, block->successors[0]);
+    for (auto b : block->successors) {
+        if (isPotentiallyUndefinedAt(name, b))
+            return true;
+    }
+    return false;
 }
 
 bool PhiAnalysis::isPotentiallyUndefinedAt(InternedString name, CFGBlock* block) {
     assert(!startswith(name.str(), "!"));
 
-    for (CFGBlock* pred : block->predecessors) {
-        DefinednessAnalysis::DefinitionLevel dlevel = definedness.isDefinedAtEnd(name, pred);
-        if (dlevel != DefinednessAnalysis::Defined)
-            return true;
-    }
-    return false;
+    assert(definedness.defined_at_beginning.count(block));
+    return definedness.defined_at_beginning[block][name] != DefinednessAnalysis::Defined;
 }
 
 LivenessAnalysis* computeLivenessInfo(CFG* cfg) {
@@ -478,6 +492,38 @@ LivenessAnalysis* computeLivenessInfo(CFG* cfg) {
 }
 
 PhiAnalysis* computeRequiredPhis(const ParamNames& args, CFG* cfg, LivenessAnalysis* liveness, ScopeInfo* scope_info) {
-    return new PhiAnalysis(args, cfg, liveness, scope_info);
+    llvm::DenseMap<InternedString, DefinednessAnalysis::DefinitionLevel> initial_map;
+
+    for (auto e : args.args)
+        initial_map[scope_info->internString(e)] = DefinednessAnalysis::Defined;
+    if (args.vararg.size())
+        initial_map[scope_info->internString(args.vararg)] = DefinednessAnalysis::Defined;
+    if (args.kwarg.size())
+        initial_map[scope_info->internString(args.kwarg)] = DefinednessAnalysis::Defined;
+
+    return new PhiAnalysis(std::move(initial_map), cfg->getStartingBlock(), false, liveness, scope_info);
+}
+
+PhiAnalysis* computeRequiredPhis(const OSREntryDescriptor* entry_descriptor, LivenessAnalysis* liveness,
+                                 ScopeInfo* scope_info) {
+    llvm::DenseMap<InternedString, DefinednessAnalysis::DefinitionLevel> initial_map;
+
+    llvm::StringSet<> potentially_undefined;
+    for (const auto& p : entry_descriptor->args) {
+        if (!startswith(p.first.str(), "!is_defined_"))
+            continue;
+        potentially_undefined.insert(p.first.str().substr(12));
+    }
+
+    for (const auto& p : entry_descriptor->args) {
+        if (p.first.str()[0] == '!')
+            continue;
+        if (potentially_undefined.count(p.first.str()))
+            initial_map[p.first] = DefinednessAnalysis::PotentiallyDefined;
+        else
+            initial_map[p.first] = DefinednessAnalysis::Defined;
+    }
+
+    return new PhiAnalysis(std::move(initial_map), entry_descriptor->backedge->target, true, liveness, scope_info);
 }
 }

src/analysis/function_analysis.h

@@ -52,6 +52,8 @@ public:
     bool isLiveAtEnd(InternedString name, CFGBlock* block);
 };
 
+class PhiAnalysis;
+
 class DefinednessAnalysis {
 public:
     enum DefinitionLevel {
@@ -62,16 +64,21 @@ public:
     typedef llvm::DenseSet<InternedString> RequiredSet;
 
 private:
-    llvm::DenseMap<CFGBlock*, llvm::DenseMap<InternedString, DefinitionLevel>> results;
-    llvm::DenseMap<CFGBlock*, RequiredSet> defined_at_end;
-    ScopeInfo* scope_info;
+    llvm::DenseMap<CFGBlock*, llvm::DenseMap<InternedString, DefinitionLevel>> defined_at_beginning, defined_at_end;
+    llvm::DenseMap<CFGBlock*, RequiredSet> defined_at_end_sets;
 
 public:
-    DefinednessAnalysis(const ParamNames& param_names, CFG* cfg, ScopeInfo* scope_info);
+    DefinednessAnalysis() {}
+
+    void run(llvm::DenseMap<InternedString, DefinitionLevel>&& initial_map, CFGBlock* initial_block,
+             ScopeInfo* scope_info);
 
     DefinitionLevel isDefinedAtEnd(InternedString name, CFGBlock* block);
     const RequiredSet& getDefinedNamesAtEnd(CFGBlock* block);
+
+    friend class PhiAnalysis;
 };
+
 class PhiAnalysis {
 public:
     typedef llvm::DenseSet<InternedString> RequiredSet;
@@ -83,18 +90,24 @@ private:
     llvm::DenseMap<CFGBlock*, RequiredSet> required_phis;
 
 public:
-    PhiAnalysis(const ParamNames&, CFG* cfg, LivenessAnalysis* liveness, ScopeInfo* scope_info);
+    // Initials_need_phis specifies that initial_map should count as an additional entry point
+    // that may require phis.
+    PhiAnalysis(llvm::DenseMap<InternedString, DefinednessAnalysis::DefinitionLevel>&& initial_map,
+                CFGBlock* initial_block, bool initials_need_phis, LivenessAnalysis* liveness, ScopeInfo* scope_info);
 
     bool isRequired(InternedString name, CFGBlock* block);
     bool isRequiredAfter(InternedString name, CFGBlock* block);
     const RequiredSet& getAllRequiredAfter(CFGBlock* block);
     const RequiredSet& getAllRequiredFor(CFGBlock* block);
+    // If "name" may be undefined at the beginning of any immediate successor block of "block":
     bool isPotentiallyUndefinedAfter(InternedString name, CFGBlock* block);
+    // If "name" may be undefined at the beginning of "block"
     bool isPotentiallyUndefinedAt(InternedString name, CFGBlock* block);
 };
 
 LivenessAnalysis* computeLivenessInfo(CFG*);
-PhiAnalysis* computeRequiredPhis(const ParamNames&, CFG*, LivenessAnalysis*, ScopeInfo* scope_Info);
+PhiAnalysis* computeRequiredPhis(const ParamNames&, CFG*, LivenessAnalysis*, ScopeInfo* scope_info);
+PhiAnalysis* computeRequiredPhis(const OSREntryDescriptor*, LivenessAnalysis*, ScopeInfo* scope_info);
 }
 
 #endif

src/analysis/type_analysis.cpp

@@ -721,7 +721,7 @@ public:
         return changed;
     }
 
-    static PropagatingTypeAnalysis* doAnalysis(CFG* cfg, SpeculationLevel speculation, ScopeInfo* scope_info,
+    static PropagatingTypeAnalysis* doAnalysis(SpeculationLevel speculation, ScopeInfo* scope_info,
                                                TypeMap&& initial_types, CFGBlock* initial_block) {
         Timer _t("PropagatingTypeAnalysis::doAnalysis()");
 
@@ -785,15 +785,16 @@ public:
         }
 
         if (VERBOSITY("types")) {
-            printf("Type analysis: %ld BBs, %d evaluations = %.1f evaluations/block\n", cfg->blocks.size(),
-                   num_evaluations, 1.0 * num_evaluations / cfg->blocks.size());
+            printf("Type analysis: %d BBs, %d evaluations = %.1f evaluations/block\n", starting_types.size(),
+                   num_evaluations, 1.0 * num_evaluations / starting_types.size());
         }
 
         if (VERBOSITY("types") >= 3) {
-            for (CFGBlock* b : cfg->blocks) {
+            for (const auto& p : starting_types) {
+                auto b = p.first;
                 printf("Types at beginning of block %d:\n", b->idx);
 
-                TypeMap& starting = starting_types[b];
+                const TypeMap& starting = p.second;
                 for (const auto& p : starting) {
                     ASSERT(p.second, "%s", p.first.c_str());
                     printf("%s: %s\n", p.first.c_str(), p.second->debugName().c_str());
@@ -836,17 +837,17 @@ TypeAnalysis* doTypeAnalysis(CFG* cfg, const ParamNames& arg_names, const std::v
 
     assert(i == arg_types.size());
 
-    return PropagatingTypeAnalysis::doAnalysis(cfg, speculation, scope_info, std::move(initial_types),
+    return PropagatingTypeAnalysis::doAnalysis(speculation, scope_info, std::move(initial_types),
                                                cfg->getStartingBlock());
 }
 
-TypeAnalysis* doTypeAnalysis(CFG* cfg, const OSREntryDescriptor* entry_descriptor, EffortLevel effort,
+TypeAnalysis* doTypeAnalysis(const OSREntryDescriptor* entry_descriptor, EffortLevel effort,
                              TypeAnalysis::SpeculationLevel speculation, ScopeInfo* scope_info) {
     // if (effort == EffortLevel::INTERPRETED) {
     // return new NullTypeAnalysis();
     //}
     TypeMap initial_types(entry_descriptor->args.begin(), entry_descriptor->args.end());
-    return PropagatingTypeAnalysis::doAnalysis(cfg, speculation, scope_info, std::move(initial_types),
+    return PropagatingTypeAnalysis::doAnalysis(speculation, scope_info, std::move(initial_types),
                                                entry_descriptor->backedge->target);
 }
 }

src/analysis/type_analysis.h

@@ -46,7 +46,7 @@ public:
 TypeAnalysis* doTypeAnalysis(CFG* cfg, const ParamNames& param_names,
                              const std::vector<ConcreteCompilerType*>& arg_types, EffortLevel effort,
                              TypeAnalysis::SpeculationLevel speculation, ScopeInfo* scope_info);
-TypeAnalysis* doTypeAnalysis(CFG* cfg, const OSREntryDescriptor* entry_descriptor, EffortLevel effort,
+TypeAnalysis* doTypeAnalysis(const OSREntryDescriptor* entry_descriptor, EffortLevel effort,
                              TypeAnalysis::SpeculationLevel speculation, ScopeInfo* scope_info);
 }
 

src/codegen/ast_interpreter.cpp

@@ -132,6 +132,7 @@ private:
     CompiledFunction* compiled_func;
     SourceInfo* source_info;
     ScopeInfo* scope_info;
+    PhiAnalysis* phis;
 
     SymMap sym_table;
     CFGBlock* next_block, *current_block;
@@ -223,9 +224,9 @@ void ASTInterpreter::gcVisit(GCVisitor* visitor) {
 }
 
 ASTInterpreter::ASTInterpreter(CompiledFunction* compiled_function)
-    : compiled_func(compiled_function), source_info(compiled_function->clfunc->source), scope_info(0), current_block(0),
-      current_inst(0), last_exception(NULL, NULL, NULL), passed_closure(0), created_closure(0), generator(0),
-      edgecount(0), frame_info(ExcInfo(NULL, NULL, NULL)) {
+    : compiled_func(compiled_function), source_info(compiled_function->clfunc->source), scope_info(0), phis(NULL),
+      current_block(0), current_inst(0), last_exception(NULL, NULL, NULL), passed_closure(0), created_closure(0),
+      generator(0), edgecount(0), frame_info(ExcInfo(NULL, NULL, NULL)) {
 
     CLFunction* f = compiled_function->clfunc;
     if (!source_info->cfg)
@@ -324,15 +325,19 @@ void ASTInterpreter::eraseDeadSymbols() {
     if (source_info->liveness == NULL)
         source_info->liveness = computeLivenessInfo(source_info->cfg);
 
-    if (source_info->phis == NULL)
-        source_info->phis = computeRequiredPhis(compiled_func->clfunc->param_names, source_info->cfg,
-                                                source_info->liveness, scope_info);
+    if (this->phis == NULL) {
+        PhiAnalysis*& phis = source_info->phis[/* entry_descriptor = */ NULL];
+        if (!phis)
+            phis = computeRequiredPhis(compiled_func->clfunc->param_names, source_info->cfg, source_info->liveness,
+                                       scope_info);
+        this->phis = phis;
+    }
 
     std::vector<InternedString> dead_symbols;
     for (auto& it : sym_table) {
         if (!source_info->liveness->isLiveAtEnd(it.first, current_block)) {
             dead_symbols.push_back(it.first);
-        } else if (source_info->phis->isRequiredAfter(it.first, current_block)) {
+        } else if (phis->isRequiredAfter(it.first, current_block)) {
             assert(scope_info->getScopeTypeOfName(it.first) != ScopeInfo::VarScopeType::GLOBAL);
         } else {
         }
@@ -472,10 +477,9 @@ Value ASTInterpreter::visit_jump(AST_Jump* node) {
 
             std::map<InternedString, Box*> sorted_symbol_table;
 
-            auto phis = compiled_func->clfunc->source->phis;
             for (auto& name : phis->definedness.getDefinedNamesAtEnd(current_block)) {
                 auto it = sym_table.find(name);
-                if (!compiled_func->clfunc->source->liveness->isLiveAtEnd(name, current_block))
+                if (!source_info->liveness->isLiveAtEnd(name, current_block))
                     continue;
 
                 if (phis->isPotentiallyUndefinedAfter(name, current_block)) {

src/codegen/codegen.cpp

@@ -35,7 +35,7 @@ namespace pyston {
 DS_DEFINE_RWLOCK(codegen_rwlock);
 
 SourceInfo::SourceInfo(BoxedModule* m, ScopingAnalysis* scoping, AST* ast, const std::vector<AST_stmt*>& body)
-    : parent_module(m), scoping(scoping), ast(ast), cfg(NULL), liveness(NULL), phis(NULL), body(body) {
+    : parent_module(m), scoping(scoping), ast(ast), cfg(NULL), liveness(NULL), body(body) {
     switch (ast->type) {
         case AST_TYPE::ClassDef:
         case AST_TYPE::Lambda:

src/codegen/irgen.cpp

@@ -342,6 +342,8 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
     CompiledFunction* cf = irstate->getCurFunction();
     ConcreteCompilerType* rtn_type = irstate->getReturnType();
     // llvm::MDNode* func_info = irstate->getFuncDbgInfo();
+    PhiAnalysis* phi_analysis = source->phis[entry_descriptor];
+    assert(phi_analysis);
 
     if (entry_descriptor != NULL)
         assert(blocks.count(source->cfg->getStartingBlock()) == 0);
@@ -508,18 +510,11 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
         CFGBlock* block = traversal_order[_i].first;
         CFGBlock* pred = traversal_order[_i].second;
 
-        if (VERBOSITY("irgen") >= 3)
-            printf("processing block %d\n", block->idx);
-
-        if (!blocks.count(block)) {
-            if (VERBOSITY("irgen") >= 3)
-                printf("Skipping this block\n");
-            // created_phis[block] = NULL;
-            // ending_symbol_tables[block] = NULL;
-            // phi_ending_symbol_tables[block] = NULL;
-            // llvm_exit_blocks[block] = NULL;
+        if (!blocks.count(block))
             continue;
-        }
+
+        if (VERBOSITY("irgen") >= 2)
+            printf("processing block %d\n", block->idx);
 
         std::unique_ptr<IRGenerator> generator(createIRGenerator(irstate, llvm_entry_blocks, block, types));
         llvm::BasicBlock* entry_block_end = llvm_entry_blocks[block];
@@ -634,9 +629,9 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
 
 
             std::set<InternedString> names;
-            for (const auto& s : source->phis->getAllRequiredFor(block)) {
+            for (const auto& s : phi_analysis->getAllRequiredFor(block)) {
                 names.insert(s);
-                if (source->phis->isPotentiallyUndefinedAfter(s, block->predecessors[0])) {
+                if (phi_analysis->isPotentiallyUndefinedAfter(s, block->predecessors[0])) {
                     names.insert(getIsDefinedName(s, source->getInternedStrings()));
                 }
             }
@@ -807,7 +802,7 @@ static void emitBBs(IRGenState* irstate, TypeAnalysis* types, const OSREntryDesc
 
             // printf("(%d %ld) -> (%d %ld)\n", bpred->idx, phi_ending_symbol_tables[bpred]->size(), b->idx,
             // phis->size());
-            assert(sameKeyset(phi_ending_symbol_tables[bpred], phis));
+            ASSERT(sameKeyset(phi_ending_symbol_tables[bpred], phis), "%d->%d", bpred->idx, b->idx);
             assert(phi_ending_symbol_tables[bpred]->size() == phis->size());
         }
 
@@ -1042,7 +1037,7 @@ CompiledFunction* doCompile(SourceInfo* source, ParamNames* param_names, const O
         speculation_level = TypeAnalysis::SOME;
     TypeAnalysis* types;
     if (entry_descriptor)
-        types = doTypeAnalysis(source->cfg, entry_descriptor, effort, speculation_level, source->getScopeInfo());
+        types = doTypeAnalysis(entry_descriptor, effort, speculation_level, source->getScopeInfo());
     else
         types = doTypeAnalysis(source->cfg, *param_names, spec->arg_types, effort, speculation_level,
                                source->getScopeInfo());
@@ -1060,7 +1055,7 @@ CompiledFunction* doCompile(SourceInfo* source, ParamNames* param_names, const O
         computeBlockSetClosure(blocks);
     }
 
-    IRGenState irstate(cf, source, param_names, getGCBuilder(), dbg_funcinfo);
+    IRGenState irstate(cf, source, source->phis[entry_descriptor], param_names, getGCBuilder(), dbg_funcinfo);
 
     emitBBs(&irstate, types, entry_descriptor, blocks);
 

src/codegen/irgen/hooks.cpp

@@ -232,8 +232,13 @@ CompiledFunction* compileFunction(CLFunction* f, FunctionSpecialization* spec, E
         if (source->liveness == NULL)
             source->liveness = computeLivenessInfo(source->cfg);
 
-        if (source->phis == NULL)
-            source->phis = computeRequiredPhis(f->param_names, source->cfg, source->liveness, source->getScopeInfo());
+        PhiAnalysis*& phis = source->phis[entry_descriptor];
+        if (!phis) {
+            if (entry_descriptor)
+                phis = computeRequiredPhis(entry_descriptor, source->liveness, source->getScopeInfo());
+            else
+                phis = computeRequiredPhis(f->param_names, source->cfg, source->liveness, source->getScopeInfo());
+        }
     }
 
 

src/codegen/irgen/irgenerator.cpp

@@ -2278,7 +2278,7 @@ private:
                 p.second->decvref(emitter);
                 symbol_table.erase(getIsDefinedName(p.first));
                 symbol_table.erase(p.first);
-            } else if (source->phis->isRequiredAfter(p.first, myblock)) {
+            } else if (irstate->getPhis()->isRequiredAfter(p.first, myblock)) {
                 assert(scope_info->getScopeTypeOfName(p.first) != ScopeInfo::VarScopeType::GLOBAL);
                 ConcreteCompilerType* phi_type = types->getTypeAtBlockEnd(p.first, myblock);
                 // printf("Converting %s from %s to %s\n", p.first.c_str(),
@@ -2302,9 +2302,10 @@ private:
             }
         }
 
-        const PhiAnalysis::RequiredSet& all_phis = source->phis->getAllRequiredAfter(myblock);
+        const PhiAnalysis::RequiredSet& all_phis = irstate->getPhis()->getAllRequiredAfter(myblock);
         for (PhiAnalysis::RequiredSet::const_iterator it = all_phis.begin(), end = all_phis.end(); it != end; ++it) {
-            // printf("phi will be required for %s\n", it->c_str());
+            if (VERBOSITY() >= 3)
+                printf("phi will be required for %s\n", it->c_str());
             assert(scope_info->getScopeTypeOfName(*it) != ScopeInfo::VarScopeType::GLOBAL);
             CompilerVariable*& cur = symbol_table[*it];
 
@@ -2316,7 +2317,7 @@ private:
                 ConcreteCompilerVariable* is_defined
                     = static_cast<ConcreteCompilerVariable*>(_popFake(defined_name, true));
 
-                if (source->phis->isPotentiallyUndefinedAfter(*it, myblock)) {
+                if (irstate->getPhis()->isPotentiallyUndefinedAfter(*it, myblock)) {
                     // printf("is potentially undefined later, so marking it defined\n");
                     if (is_defined) {
                         _setFake(defined_name, is_defined);
@@ -2415,7 +2416,7 @@ public:
         // We have one successor, but they have more than one predecessor.
         // We're going to sort out which symbols need to go in phi_st and which belong inst.
         for (SymbolTable::iterator it = st->begin(); it != st->end();) {
-            if (allowableFakeEndingSymbol(it->first) || source->phis->isRequiredAfter(it->first, myblock)) {
+            if (allowableFakeEndingSymbol(it->first) || irstate->getPhis()->isRequiredAfter(it->first, myblock)) {
                 ASSERT(it->second->isGrabbed(), "%s", it->first.c_str());
                 assert(it->second->getVrefs() == 1);
                 // this conversion should have already happened... should refactor this.

src/codegen/irgen/irgenerator.h

@@ -56,6 +56,7 @@ class IRGenState {
 private:
     CompiledFunction* cf;
     SourceInfo* source_info;
+    PhiAnalysis* phis;
     ParamNames* param_names;
     GCBuilder* gc;
     llvm::MDNode* func_dbg_info;
@@ -68,9 +69,9 @@ private:
 
 
 public:
-    IRGenState(CompiledFunction* cf, SourceInfo* source_info, ParamNames* param_names, GCBuilder* gc,
+    IRGenState(CompiledFunction* cf, SourceInfo* source_info, PhiAnalysis* phis, ParamNames* param_names, GCBuilder* gc,
                llvm::MDNode* func_dbg_info)
-        : cf(cf), source_info(source_info), param_names(param_names), gc(gc), func_dbg_info(func_dbg_info),
+        : cf(cf), source_info(source_info), phis(phis), param_names(param_names), gc(gc), func_dbg_info(func_dbg_info),
           scratch_space(NULL), frame_info(NULL), frame_info_arg(NULL), scratch_size(0) {
         assert(cf->func);
         assert(!cf->clfunc); // in this case don't need to pass in sourceinfo
@@ -92,6 +93,8 @@ public:
 
     SourceInfo* getSourceInfo() { return source_info; }
 
+    PhiAnalysis* getPhis() { return phis; }
+
     ScopeInfo* getScopeInfo();
     ScopeInfo* getScopeInfoForNode(AST* node);
 

src/core/types.h

@@ -244,7 +244,7 @@ public:
     AST* ast;
     CFG* cfg;
     LivenessAnalysis* liveness;
-    PhiAnalysis* phis;
+    std::unordered_map<const OSREntryDescriptor*, PhiAnalysis*> phis;
     bool is_generator;
 
     InternedStringPool& getInternedStrings();

test/tests/osr_maybe_undefined3.py

-# fail-if: '-O' not in EXTRA_JIT_ARGS
-# - wip
-
 # Regression test: make sure we can handle variables that are only defined
 # on excluded parts of an osr compilation
 def f():

test/unittests/CMakeLists.txt

@@ -16,3 +16,4 @@ endmacro()
 add_unittest(gc)
 add_unittest(analysis)
 add_custom_command(TARGET analysis_unittest POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_SOURCE_DIR}/test/unittests/analysis_listcomp.py ${CMAKE_BINARY_DIR}/test/unittests/analysis_listcomp.py)
+add_custom_command(TARGET analysis_unittest POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_SOURCE_DIR}/test/unittests/analysis_osr.py ${CMAKE_BINARY_DIR}/test/unittests/analysis_osr.py)

test/unittests/analysis.cpp

@@ -7,6 +7,7 @@
 
 #include "analysis/function_analysis.h"
 #include "analysis/scoping_analysis.h"
+#include "codegen/osrentry.h"
 #include "codegen/parser.h"
 #include "core/ast.h"
 #include "core/cfg.h"
@@ -16,7 +17,7 @@ using namespace pyston;
 
 class AnalysisTest : public ::testing::Test {
 protected:
-    virtual void SetUp() {
+    static void SetUpTestCase() {
         initCodegen();
     }
 };
@@ -35,7 +36,7 @@ TEST_F(AnalysisTest, augassign) {
     ASSERT_FALSE(scope_info->getScopeTypeOfName(module->interned_strings->get("a")) == ScopeInfo::VarScopeType::GLOBAL);
     ASSERT_FALSE(scope_info->getScopeTypeOfName(module->interned_strings->get("b")) == ScopeInfo::VarScopeType::GLOBAL);
 
-    SourceInfo* si = new SourceInfo(createModule("__main__", fn), scoping, func, func->body);
+    SourceInfo* si = new SourceInfo(createModule("augassign", fn), scoping, func, func->body);
 
     CFG* cfg = computeCFG(si, func->body);
     LivenessAnalysis* liveness = computeLivenessInfo(cfg);
@@ -51,3 +52,83 @@ TEST_F(AnalysisTest, augassign) {
     PhiAnalysis* phis = computeRequiredPhis(ParamNames(func), cfg, liveness, scope_info);
 }
 
+void doOsrTest(bool is_osr, bool i_maybe_undefined) {
+    const std::string fn("test/unittests/analysis_osr.py");
+    AST_Module* module = caching_parse_file(fn.c_str());
+    assert(module);
+
+    ScopingAnalysis *scoping = new ScopingAnalysis(module);
+
+    assert(module->body[0]->type == AST_TYPE::FunctionDef);
+    AST_FunctionDef* func = static_cast<AST_FunctionDef*>(module->body[0]);
+
+    ScopeInfo* scope_info = scoping->getScopeInfoForNode(func);
+    SourceInfo* si = new SourceInfo(createModule("osr" + std::to_string((is_osr << 1) + i_maybe_undefined), fn),
+            scoping, func, func->body);
+
+    CFG* cfg = computeCFG(si, func->body);
+    LivenessAnalysis* liveness = computeLivenessInfo(cfg);
+
+    // cfg->print();
+
+    InternedString i_str = module->interned_strings->get("i");
+    InternedString idi_str = module->interned_strings->get("!is_defined_i");
+    InternedString iter_str = module->interned_strings->get("#iter_3");
+
+    CFGBlock* loop_backedge = cfg->blocks[5];
+    ASSERT_EQ(6, loop_backedge->idx);
+    ASSERT_EQ(1, loop_backedge->body.size());
+
+    ASSERT_EQ(AST_TYPE::Jump, loop_backedge->body[0]->type);
+    AST_Jump* backedge = ast_cast<AST_Jump>(loop_backedge->body[0]);
+    ASSERT_LE(backedge->target->idx, loop_backedge->idx);
+
+    PhiAnalysis* phis;
+
+    if (is_osr) {
+        OSREntryDescriptor* entry_descriptor = OSREntryDescriptor::create(NULL, backedge);
+        entry_descriptor->args[i_str] = NULL;
+        if (i_maybe_undefined)
+            entry_descriptor->args[idi_str] = NULL;
+        entry_descriptor->args[iter_str] = NULL;
+        phis = computeRequiredPhis(entry_descriptor, liveness, scope_info);
+    } else {
+        phis = computeRequiredPhis(ParamNames(func), cfg, liveness, scope_info);
+    }
+
+    // First, verify that we require phi nodes for the block we enter into.
+    // This is somewhat tricky since the osr entry represents an extra entry
+    // into the BB which the analysis might not otherwise track.
+
+    auto required_phis = phis->getAllRequiredFor(backedge->target);
+    EXPECT_EQ(1, required_phis.count(i_str));
+    EXPECT_EQ(0, required_phis.count(idi_str));
+    EXPECT_EQ(1, required_phis.count(iter_str));
+    EXPECT_EQ(2, required_phis.size());
+
+    EXPECT_EQ(!is_osr || i_maybe_undefined, phis->isPotentiallyUndefinedAt(i_str, backedge->target));
+    EXPECT_FALSE(phis->isPotentiallyUndefinedAt(iter_str, backedge->target));
+    EXPECT_EQ(!is_osr || i_maybe_undefined, phis->isPotentiallyUndefinedAfter(i_str, loop_backedge));
+    EXPECT_FALSE(phis->isPotentiallyUndefinedAfter(iter_str, loop_backedge));
+
+    // Now, let's verify that we don't need a phi after the loop
+
+    CFGBlock* if_join = cfg->blocks[7];
+    ASSERT_EQ(8, if_join->idx);
+    ASSERT_EQ(2, if_join->predecessors.size());
+
+    if (is_osr)
+        EXPECT_EQ(0, phis->getAllRequiredFor(if_join).size());
+    else
+        EXPECT_EQ(1, phis->getAllRequiredFor(if_join).size());
+}
+
+TEST_F(AnalysisTest, osr_initial) {
+    doOsrTest(false, false);
+}
+TEST_F(AnalysisTest, osr1) {
+    doOsrTest(true, false);
+}
+TEST_F(AnalysisTest, osr2) {
+    doOsrTest(true, true);
+}

test/unittests/analysis_osr.py

+def f():
+    if True:
+        for i in xrange(20000):
+            pass
+    else:
+        a = 1
+f()
+