Repository: Kotlin/kotlinx.coroutines
File: kotlinx-coroutines-core/jvm/src/scheduling/WorkQueue.kt (Kotlin)

Code before:
        }

        return null
    }

    // Polls for blocking task, invoked only by the owner
    fun pollBlocking(): Task? {
        while (true) { // Poll the slot
            val lastScheduled = lastScheduledTask.value ?: break
            if (!lastScheduled.isBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        return pollWithMode(onlyBlocking = true /* only blocking */)
    }

    fun pollCpu(): Task? {
        while (true) { // Poll the slot
            val lastScheduled = lastScheduledTask.value ?: break
            if (lastScheduled.isBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        return pollWithMode(onlyBlocking = false /* only cpu */)
    }

    private fun pollWithMode(/* Only blocking OR only CPU */ onlyBlocking: Boolean): Task? {
        val start = consumerIndex.value
        var end = producerIndex.value
        // CPU or (BLOCKING & hasBlocking)
        val shouldProceed = !onlyBlocking || blockingTasksInBuffer.value > 0
        while (start != end && shouldProceed) {
            val task = tryExtractFromTheMiddle(--end, onlyBlocking)
            if (task != null) {
                return task
            }
        }
        return null
    }

    private fun tryExtractFromTheMiddle(index: Int, onlyBlocking: Boolean): Task? {
        if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
        val arrayIndex = index and MASK
        val value = buffer[arrayIndex]
        if (value != null && value.isBlocking == onlyBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
            if (onlyBlocking) blockingTasksInBuffer.decrementAndGet()
            return value

Review comment (refactor):
Why not move this code to `pollWithMode` as well? Only one line would need to change to
```kotlin
            if (lastScheduled.isBlocking == onlyBlocking) break
```

Code after:
        while (start != end) {
            if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
            return tryExtractFromTheMiddle(start++, onlyBlocking) ?: continue
        }

        return null
    }

    // Polls for blocking task, invoked only by the owner
    // NB: ONLY for runSingleTask method
    fun pollBlocking(): Task? = pollWithExclusiveMode(onlyBlocking = true /* only blocking */)

    // Polls for CPU task, invoked only by the owner
    // NB: ONLY for runSingleTask method
    fun pollCpu(): Task? = pollWithExclusiveMode(onlyBlocking = false /* only cpu */)

    private fun pollWithExclusiveMode(/* Only blocking OR only CPU */ onlyBlocking: Boolean): Task? {
        while (true) { // Poll the slot
            val lastScheduled = lastScheduledTask.value ?: break
            if (lastScheduled.isBlocking != onlyBlocking) break
            if (lastScheduledTask.compareAndSet(lastScheduled, null)) {
                return lastScheduled
            } // Failed -> someone else stole it
        }

        // Failed to poll the slot, scan the queue
        val start = consumerIndex.value
        var end = producerIndex.value
        // Bail out if there is no blocking work for us
        while (start != end) {
            if (onlyBlocking && blockingTasksInBuffer.value == 0) return null
            val task = tryExtractFromTheMiddle(--end, onlyBlocking)
            if (task != null) {
                return task
            }
        }
        return null
    }

    private fun tryExtractFromTheMiddle(index: Int, onlyBlocking: Boolean): Task? {
        val arrayIndex = index and MASK
        val value = buffer[arrayIndex]
        if (value != null && value.isBlocking == onlyBlocking && buffer.compareAndSet(arrayIndex, value, null)) {
            if (onlyBlocking) blockingTasksInBuffer.decrementAndGet()
            return value
        }
        return null
    }

    fun offloadAllWorkTo(globalQueue: GlobalQueue) {
        lastScheduledTask.getAndSet(null)?.let { globalQueue.addLast(it) }