Back out "Support static awaits for co-routine objects"

Summary:
Original commit changeset: 4ce3ad9

Original Phabricator Diff: D35194272 (8a01c37)

Reviewed By: jbower-fb

Differential Revision: D35995916

fbshipit-source-id: 437cf00

Include/genobject.h

@@ -140,15 +140,6 @@ static inline void _PyAwaitable_SetAwaiter(PyObject *receiver, PyObject *awaiter
     }
 }
 
-void
-_PyCoro_SetAwaiter(PyCoroObject *coro, PyCoroObject *awaiter);
-
-PySendResult
-_PyGen_DoSend(PyThreadState *tstate,
-              PyGenObject *gen,
-              PyObject *arg,
-              PyObject **pResult);
-
 /* Asynchronous Generators */
 
 typedef struct {

Jit/codegen/autogen.cpp

@@ -477,7 +477,6 @@ void translateYieldValue(Environ* env, const Instruction* instr) {
 void translateYieldFrom(Environ* env, const Instruction* instr) {
   asmjit::x86::Builder* as = env->as;
   bool skip_initial_send = instr->isYieldFromSkipInitialSend();
-  bool is_coroutine = instr->isYieldFromCoroutine();
 
   // Make sure tstate is in RDI for use in epilogue and here.
   int tstate_loc = instr->getInput(0)->getPhyRegOrStackSlot();
@@ -529,17 +528,11 @@ void translateYieldFrom(Environ* env, const Instruction* instr) {
       PhyLocation(iter_loc).toString().c_str());
   as->mov(x86::rdi, x86::ptr(x86::rbp, iter_loc));
 
-  if (is_coroutine) {
-    emitCall(*env, reinterpret_cast<uint64_t>(JITRT_YieldFromCoroutine), instr);
-  } else if (instr->isYieldFromHandleStopAsyncIteration()) {
-    emitCall(
-        *env,
-        reinterpret_cast<uint64_t>(JITRT_YieldFromHandleStopAsyncIteration),
-        instr);
-  } else {
-    emitCall(*env, reinterpret_cast<uint64_t>(JITRT_YieldFrom), instr);
-  }
-
+  uint64_t func = reinterpret_cast<uint64_t>(
+      instr->isYieldFromHandleStopAsyncIteration()
+          ? JITRT_YieldFromHandleStopAsyncIteration
+          : JITRT_YieldFrom);
+  emitCall(*env, func, instr);
   // Yielded or final result value now in RAX. If the result was NULL then
   // done will be set so we'll correctly jump to the following CheckExc.
   const auto yf_result_phys_reg = PhyLocation::RAX;
@@ -1110,10 +1103,6 @@ BEGIN_RULES(Instruction::kYieldFrom)
   GEN(ANY, CALL(translateYieldFrom))
 END_RULES
 
-BEGIN_RULES(Instruction::kYieldFromCoroutine)
-  GEN(ANY, CALL(translateYieldFrom))
-END_RULES
-
 BEGIN_RULES(Instruction::kYieldFromSkipInitialSend)
   GEN(ANY, CALL(translateYieldFrom))
 END_RULES

Jit/compiler.cpp

@@ -98,7 +98,6 @@ void Compiler::runPasses(
   runPass<jit::hir::SSAify>(irfunc, callback);
   runPass<jit::hir::Simplify>(irfunc, callback);
   runPass<jit::hir::DynamicComparisonElimination>(irfunc, callback);
-  runPass<jit::hir::AwaitOptimization>(irfunc, callback);
   runPass<jit::hir::GuardTypeRemoval>(irfunc, callback);
   runPass<jit::hir::PhiElimination>(irfunc, callback);
   if (_PyJIT_IsHIRInlinerEnabled()) {

Jit/hir/builder.cpp

@@ -3901,7 +3901,7 @@ void HIRBuilder::emitGetAwaitable(
   tc.emit<CallCFunc>(
       1,
       iter,
-      CallCFunc::Func::kJITRT_GetAwaitableIter,
+      CallCFunc::Func::k_PyCoro_GetAwaitableIter,
       std::vector<Register*>{iterable});
   if (prev_op == BEFORE_ASYNC_WITH || prev_op == WITH_CLEANUP_START) {
     BasicBlock* error_block = cfg.AllocateBlock();
@@ -3918,7 +3918,28 @@ void HIRBuilder::emitGetAwaitable(
     tc.emit<CheckExc>(iter, iter, tc.frame);
   }
 
+  // For coroutines only, runtime assert it isn't already awaiting by checking
+  // if it has a sub-iterator using _PyGen_yf().
+  TranslationContext block_assert_not_awaited_coro{
+      cfg.AllocateBlock(), tc.frame};
+  TranslationContext block_done{cfg.AllocateBlock(), tc.frame};
+  tc.emit<CondBranchCheckType>(
+      iter,
+      Type::fromTypeExact(&PyCoro_Type),
+      block_assert_not_awaited_coro.block,
+      block_done.block);
+  Register* yf = temps_.AllocateStack();
+  block_assert_not_awaited_coro.emit<CallCFunc>(
+      1, yf, CallCFunc::Func::k_PyGen_yf, std::vector<Register*>{iter});
+  TranslationContext block_coro_already_awaited{cfg.AllocateBlock(), tc.frame};
+  block_assert_not_awaited_coro.emit<CondBranch>(
+      yf, block_coro_already_awaited.block, block_done.block);
+  block_coro_already_awaited.emit<RaiseStatic>(
+      0, PyExc_RuntimeError, "coroutine is being awaited already", tc.frame);
+
   stack.push(iter);
+
+  tc.block = block_done.block;
 }
 
 void HIRBuilder::emitBuildString(

Jit/hir/hir.h

@@ -1527,7 +1527,7 @@ class INSTR_CLASS(CallCFunc, (TOptObject | TCUInt64), HasOutput, Operands<>) {
 // List of allowed functions
 #define CallCFunc_FUNCS(X)      \
   X(_PyAsyncGenValueWrapperNew) \
-  X(JITRT_GetAwaitableIter)     \
+  X(_PyCoro_GetAwaitableIter)   \
   X(_PyGen_yf)                  \
   X(_PyEval_GetAIter)           \
   X(_PyEval_GetANext)           \
@@ -1555,10 +1555,6 @@ class INSTR_CLASS(CallCFunc, (TOptObject | TCUInt64), HasOutput, Operands<>) {
     return kFuncNames[static_cast<size_t>(func_)];
   }
 
-  Func func() const {
-    return func_;
-  }
-
  private:
   const Func func_;
 

Jit/hir/optimization.cpp

@@ -35,7 +35,6 @@ PassRegistry::PassRegistry() {
   addPass(CopyPropagation::Factory);
   addPass(CleanCFG::Factory);
   addPass(DynamicComparisonElimination::Factory);
-  addPass(AwaitOptimization::Factory);
   addPass(PhiElimination::Factory);
   addPass(InlineFunctionCalls::Factory);
   addPass(Simplify::Factory);
@@ -291,70 +290,6 @@ void DynamicComparisonElimination::Run(Function& irfunc) {
   reflowTypes(irfunc);
 }
 
-void AwaitOptimization::Run(Function& irfunc) {
-  if (irfunc.code && !(irfunc.code->co_flags & CO_COROUTINE)) {
-    // Skip in non-async functions
-    return;
-  }
-
-  LivenessAnalysis liveness{irfunc};
-  liveness.Run();
-  auto last_uses = liveness.GetLastUses();
-
-  std::vector<Instr*> removals;
-  for (auto& block : irfunc.cfg.blocks) {
-    for (auto it = block.begin(); it != block.end();) {
-      auto& instr = *it;
-      ++it;
-      if (!instr.IsCallCFunc()) {
-        continue;
-      }
-
-      auto& callc = static_cast<CallCFunc&>(instr);
-      if (callc.func() != CallCFunc::Func::kJITRT_GetAwaitableIter) {
-        continue;
-      }
-
-      auto input_instr = modelReg(callc.GetOperand(0))->instr();
-      if (!input_instr->IsInvokeStaticFunction()) {
-        continue;
-      }
-
-      auto invoke = static_cast<InvokeStaticFunction*>(input_instr);
-      if (((reinterpret_cast<PyCodeObject*>(invoke->func()->func_code))
-               ->co_flags &
-           CO_COROUTINE) == 0) {
-        continue;
-      }
-
-      auto& dying_regs = map_get(last_uses, &callc, kEmptyRegSet);
-      if (dying_regs.count(callc.GetOperand(0)) == 0) {
-        // Invoke output lives after GetAwaitableIterable, can't re-write
-        // because this isn't just a simple "await invoke()"
-        continue;
-      }
-
-      removals.push_back(&callc);
-      auto assign = Assign::create(callc.GetOutput(), callc.GetOperand(0));
-      assign->InsertBefore(callc);
-
-      ++it;
-      auto& next = *it;
-      if (next.IsCheckExc()) {
-        removals.push_back(&next);
-        auto assign = Assign::create(next.GetOutput(), next.GetOperand(0));
-        assign->InsertBefore(next);
-      }
-    }
-  }
-
-  for (auto instr : removals) {
-    instr->unlink();
-    delete instr;
-  }
-  reflowTypes(irfunc);
-}
-
 void CopyPropagation::Run(Function& irfunc) {
   std::vector<Instr*> assigns;
   for (auto block : irfunc.cfg.GetRPOTraversal()) {

Jit/hir/optimization.h

@@ -99,20 +99,6 @@ class DynamicComparisonElimination : public Pass {
   PyCFunction isinstance_func_{nullptr};
 };
 
-class AwaitOptimization : public Pass {
- public:
-  AwaitOptimization() : Pass("AwaitOptimization") {}
-
-  void Run(Function& irfunc) override;
-
-  static std::unique_ptr<AwaitOptimization> Factory() {
-    return std::make_unique<AwaitOptimization>();
-  }
-
- private:
-  DISALLOW_COPY_AND_ASSIGN(AwaitOptimization);
-};
-
 // Eliminate Assign instructions by propagating copies.
 class CopyPropagation : public Pass {
  public:

Jit/hir/preload.cpp

@@ -139,7 +139,8 @@ static std::unique_ptr<InvokeTarget> resolve_target_descr(
   target->container_is_immutable = _PyClassLoader_IsImmutable(container);
   if (return_pytype != NULL) {
     if (coroutine) {
-      target->return_type = TCoroutine;
+      // TODO properly handle coroutine returns awaitable type
+      target->return_type = TObject;
     } else {
       target->return_type =
           to_jit_type({std::move(return_pytype), optional, exact});

Jit/hir/type.cpp

@@ -31,7 +31,6 @@ const std::unordered_map<Type, PyTypeObject*>& typeToPyType() {
         {TBytes, &PyBytes_Type},
         {TCell, &PyCell_Type},
         {TCode, &PyCode_Type},
-        {TCoroutine, &PyCoro_Type},
         {TDict, &PyDict_Type},
         {TBaseException, reinterpret_cast<PyTypeObject*>(PyExc_BaseException)},
         {TFloat, &PyFloat_Type},