select! 与 join!

($(biased;)? else => $else:expr $(,)? ) => {{
    $else
}};

(biased; $p:pat = $($t:tt)* ) => {
    $crate::select!(@{ start=0; () } $p = $($t)*)
};

( $p:pat = $($t:tt)* ) => {
    // Randomly generate a starting point. This makes `select!` a bit more
    // fair and avoids always polling the first future.
    $crate::select!(@{ start={ $crate::macros::support::thread_rng_n(BRANCHES) }; () } $p = $($t)*)
};

// Use a bitfield to track which futures completed
let mask = Ident::new(
    if branches <= 8 {
        "u8"
    } else if branches <= 16 {
        "u16"
    } else if branches <= 32 {
        "u32"
    } else if branches <= 64 {
        "u64"
    } else {
        panic!("up to 64 branches supported");
    },
    Span::call_site(),
);

TokenStream::from(quote! {
    pub(super) enum Out<#( #variants ),*> {
        #( #variants(#variants), )*
        // Include a `Disabled` variant signifying that all select branches
        // failed to resolve.
        Disabled,
    }

    pub(super) type Mask = #mask;
})

const BRANCHES: u32 = $crate::count!( $($count)* );

let mut disabled: __tokio_select_util::Mask = Default::default();

// First, invoke all the pre-conditions. For any that return true,
// set the appropriate bit in `disabled`.
$(
    if !$c {
        let mask: __tokio_select_util::Mask = 1 << $crate::count!( $($skip)* );
        disabled |= mask;
    }
)*

$crate::macros::support::poll_fn(|cx| {
    // Return `Pending` when the task budget is depleted since budget-aware futures
    // are going to yield anyway and other futures will not cooperate.
    $crate::macros::support::ready!($crate::macros::support::poll_budget_available(cx));

    // ...
    let start = $start;

    for i in 0..BRANCHES {
        let branch;
        #[allow(clippy::modulo_one)]
        {
            branch = (start + i) % BRANCHES;
        }
        match branch {
            $(
                #[allow(unreachable_code)]
                $crate::count!( $($skip)* ) => {
                    // ...
                    let mask = 1 << branch;

                    if disabled & mask == mask {
                        // The future has been disabled.
                        continue;
                    }

                    // ...
                    let ( $($skip,)* fut, .. ) = &mut *futures;

                    // ...
                    let mut fut = unsafe { $crate::macros::support::Pin::new_unchecked(fut) };

                    // Try polling it
                    let out = match $crate::macros::support::Future::poll(fut, cx) {
                        $crate::macros::support::Poll::Ready(out) => out,
                        $crate::macros::support::Poll::Pending => {
                            // ...
                            is_pending = true;
                            continue;
                        }
                    };

                    // Disable the future from future polling.
                    disabled |= mask;

                    // ...
                    match &out {
                        $crate::select_priv_clean_pattern!($bind) => {}
                        _ => continue,
                    }

                    // The select is complete, return the value
                    return $crate::macros::support::Poll::Ready($crate::select_variant!(__tokio_select_util::Out, ($($skip)*))(out));
                }
            )*
            _ => unreachable!("reaching this means there probably is an off by one bug"),
        }
    }

    if is_pending {
        $crate::macros::support::Poll::Pending
    } else {
        // All branches have been disabled.
        $crate::macros::support::Poll::Ready(__tokio_select_util::Out::Disabled)
    }
}).await

// Create a scope to separate polling from handling the output. This
// adds borrow checker flexibility when using the macro.
let mut output = {
    // ...
    let mut futures = &mut futures;

    $crate::macros::support::poll_fn(|cx| { /* ... */ }).await
};

match output {
    $(
        $crate::select_variant!(__tokio_select_util::Out, ($($skip)*) ($bind)) => $handle,
    )*
    __tokio_select_util::Out::Disabled => $else,
    _ => unreachable!("failed to match bind"),
}

use std::future::ready;

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let mut wins = [0u32; 3];
  for _ in 0..30_000 {
    let i = tokio::select! {
      v = ready(0) => v,
      v = ready(1) => v,
      v = ready(2) => v,
    };
    wins[i] += 1;
  }
  println!("random start: {wins:?}");

  let mut wins = [0u32; 3];
  for _ in 0..30_000 {
    let i = tokio::select! {
      biased;
      v = ready(0) => v,
      v = ready(1) => v,
      v = ready(2) => v,
    };
    wins[i] += 1;
  }
  println!("biased:       {wins:?}");

  // a branch whose output does not match its pattern is disabled, not chosen
  let r = tokio::select! {
    Some(v) = ready(None::<u32>) => format!("got {v}"),
    else => "all branches disabled -> else".to_string(),
  };
  println!("{r}");
}

random start: [9916, 10062, 10022]
biased:       [30000, 0, 0]
all branches disabled -> else

pub enum MaybeDone<Fut: Future> {
    /// A not-yet-completed future.
    Future { #[pin] future: Fut },
    /// The output of the completed future.
    Done { output: Fut::Output },
    /// The empty variant after the result of a [`MaybeDone`] has been
    /// taken using the [`take_output`](MaybeDone::take_output) method.
    Gone,
}

pub fn num_skip(&mut self) -> u32 {
    let num_skip = self.next;
    self.next += 1;
    if self.next == COUNT {
        self.next = 0;
    }
    num_skip
}

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

// Returns Pending `n` times (waking itself each time), logging each poll.
struct Noisy { name: char, left: u32, log: std::rc::Rc<std::cell::RefCell<String>> }

impl Future for Noisy {
  type Output = ();
  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
    self.log.borrow_mut().push(self.name);
    if self.left == 0 {
      return Poll::Ready(());
    }
    self.left -= 1;
    cx.waker().wake_by_ref();
    Poll::Pending
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  for biased in [false, true] {
    let log = std::rc::Rc::new(std::cell::RefCell::new(String::new()));
    let mk = |name| Noisy { name, left: 3, log: log.clone() };
    if biased {
      tokio::join!(biased; mk('A'), mk('B'), mk('C'));
    } else {
      tokio::join!(mk('A'), mk('B'), mk('C'));
    }
    let s = log.borrow();
    let rounds: Vec<&str> = (0..s.len()).step_by(3).map(|i| &s[i..i + 3]).collect();
    println!("{}: {}", if biased { "biased " } else { "default" }, rounds.join(" | "));
  }
}

default: ABC | BCA | CAB | ABC
biased : ABC | ABC | ABC | ABC

async fn recv_and_log(rx: &mut mpsc::Receiver<u32>) -> Option<u32> {
  let v = rx.recv().await?;
  sleep(Duration::from_millis(3)).await; // e.g. write an audit log
  Some(v)
}

use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::sleep;

// Not cancel safe: the message is already taken out of the channel
// before the second .await. Dropping the future there loses it.
async fn recv_and_log(rx: &mut mpsc::Receiver<u32>) -> Option<u32> {
  let v = rx.recv().await?;
  sleep(Duration::from_millis(3)).await; // e.g. write an audit log
  Some(v)
}

async fn run(safe: bool) {
  let (tx, mut rx) = mpsc::channel(16);
  tokio::spawn(async move {
    for i in 0..10 {
      tx.send(i).await.unwrap();
      sleep(Duration::from_millis(4)).await;
    }
  });
  let mut got = Vec::new();
  let mut ticks = 0;
  let mut tick = tokio::time::interval(Duration::from_millis(5));
  loop {
    tokio::select! {
      biased;
      v = async { if safe { rx.recv().await } else { recv_and_log(&mut rx).await } } => {
        match v { Some(v) => got.push(v), None => break }
      }
      _ = tick.tick() => ticks += 1,
    }
  }
  println!("{:<14} received {:?} ({} items, {} ticks)",
    if safe { "rx.recv():" } else { "recv_and_log:" }, got, got.len(), ticks);
}

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
  run(true).await;
  run(false).await;
}

rx.recv():     received [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] (10 items, 9 ticks)
recv_and_log:  received [0, 3, 4, 5, 8, 9] (6 items, 9 ticks)

A1 B1 C1 D1 C2 D2 A2 D3 A3 A4