Repository: dotnet/runtime
File: src/coreclr/jit/async.cpp (C++)

Code before:
    for (BasicBlock* block = m_compiler->fgFirstBB; block != nullptr; block = nextBlock)
    {
        nextBlock = block->Next();
        for (GenTree* tree : LIR::AsRange(block))
        {
            if (!tree->IsCall() || !tree->AsCall()->IsAsync() || tree->AsCall()->IsTailCall())
            {
                continue;
            }

            if (tree->AsCall()->GetAsyncInfo().IsTailAwait)
            {
                TransformTailAwait(block, tree->AsCall(), &nextBlock);
                result = PhaseStatus::MODIFIED_EVERYTHING;
                break;
            }

            JITDUMP(FMT_BB " contains await(s)\n", block->bbNum);
            worklist.Push(block);
            break;
        }
    }

    JITDUMP("Found %d blocks with awaits\n", worklist.Height());

    if (worklist.Height() <= 0)
    {
        return result;
    }

    m_compiler->compSuspensionPoints = new (m_compiler, CMK_Async)
        jitstd::vector<ICorDebugInfo::AsyncSuspensionPoint>(m_compiler->getAllocator(CMK_Async));
    m_compiler->compAsyncVars = new (m_compiler, CMK_Async)
        jitstd::vector<ICorDebugInfo::AsyncContinuationVarInfo>(m_compiler->getAllocator(CMK_Async));

    m_asyncInfo = m_compiler->eeGetAsyncInfo();

    // Create the shared return BB now to put it in the right place in the block order.
    GetSharedReturnBB();

    // Compute liveness to be used for determining what must be captured on
    // suspension.
    if (m_compiler->m_dfsTree == nullptr)
    {
        m_compiler->m_dfsTree = m_compiler->fgComputeDfs<false>();
    }

    m_compiler->lvaComputePreciseRefCounts(/* isRecompute */ true, /* setSlotNumbers */ false);
    m_compiler->fgAsyncLiveness();
    INDEBUG(m_compiler->mostRecentlyActivePhase = PHASE_ASYNC);
    VarSetOps::AssignNoCopy(m_compiler, m_compiler->compCurLife, VarSetOps::MakeEmpty(m_compiler));

Review comment (suggestion):
The initial scan stops at the first async call in a block (break at the end of the loop) and only handles tail awaits (IsTailAwait) in this scan. If a block contains a non-tail await followed by a tail await, the tail await won’t be transformed here, and later the worklist-based Transform path doesn’t special-case IsTailAwait, so it would be treated like a normal await (creating a suspension point) and lose tail-await semantics. Consider either scanning the whole block for IsTailAwait, or teaching the worklist Transform loop to dispatch tail awaits to TransformTailAwait when encountered.

Code after:
    for (BasicBlock* block = m_compiler->fgFirstBB; block != nullptr; block = nextBlock)
    {
        bool hasAwait = false;
        nextBlock     = block->Next();
        for (GenTree* tree : LIR::AsRange(block))
        {
            if (!tree->IsCall() || !tree->AsCall()->IsAsync() || tree->AsCall()->IsTailCall())
            {
                continue;
            }

            if (tree->AsCall()->GetAsyncInfo().IsTailAwait)
            {
                TransformTailAwait(block, tree->AsCall(), &nextBlock);
                result = PhaseStatus::MODIFIED_EVERYTHING;
                break;
            }

            JITDUMP(FMT_BB " contains await(s)\n", block->bbNum);
            hasAwait = true;
        }

        if (hasAwait)
        {
            worklist.Push(block);
        }
    }

    JITDUMP("Found %d blocks with awaits\n", worklist.Height());

    if (worklist.Height() <= 0)
    {
        return result;
    }

    m_compiler->compSuspensionPoints = new (m_compiler, CMK_Async)
        jitstd::vector<ICorDebugInfo::AsyncSuspensionPoint>(m_compiler->getAllocator(CMK_Async));
    m_compiler->compAsyncVars = new (m_compiler, CMK_Async)
        jitstd::vector<ICorDebugInfo::AsyncContinuationVarInfo>(m_compiler->getAllocator(CMK_Async));

    m_asyncInfo = m_compiler->eeGetAsyncInfo();

    // Create the shared return BB now to put it in the right place in the block order.
    GetSharedReturnBB();

    // Compute liveness to be used for determining what must be captured on
    // suspension.
    if (m_compiler->m_dfsTree == nullptr)
    {
        m_compiler->m_dfsTree = m_compiler->fgComputeDfs<false>();
    }