Rustコトハジメ

Rustで競技プログラミングをやる人の覚書のようなものです。

tokio::fsはただspawn_blockingしてるだけ

tokio::fsはファイルシステムの操作をasyncにして提供するライブラリなのだが、ファイルシステムの操作はシステムコールを呼んでいるためブロッキングである。これをどうやって非同期に偽造してるかというと、

// 例としてread
pub async fn read(path: impl AsRef<Path>) -> io::Result<Vec<u8>> {
    let path = path.as_ref().to_owned();
    asyncify(move || std::fs::read(path)).await
}

pub(crate) async fn asyncify<F, T>(f: F) -> io::Result<T>
where
    F: FnOnce() -> io::Result<T> + Send + 'static,
    T: Send + 'static,
{
    match sys::run(f).await {
        Ok(res) => res,
        Err(_) => Err(io::Error::new(
            io::ErrorKind::Other,
            "background task failed",
        )),
    }
}

asyncifyという関数でラップしてるだけ。では、asyncifyは何なのかというと、中で呼ばれてるsys::runは、spawn_blockingである。

以下完全にはわからない

では、spawn_blockingは一旦何なのかというと、

In general, issuing a blocking call or performing a lot of compute in a future without yielding is not okay, as it may prevent the executor from driving other futures forward. This function runs the provided closure on a thread dedicated to blocking operations. See the CPU-bound tasks and blocking code section for more information.

RustのFutureが採用する協調的マルチタスキングでは、自分が何もしない時には自分からCPUを放棄しないといけない。ふつうにブロッキングコールをすると、これが出来ないから、こうやって包んで放棄を偽造しないといけない。

ではこうすると何が起こるかというと、このspawn_blockingをawaitしてるFutureが再度ランタイムによってスケジュールされてpollされる時までは他のタスクがCPUを使うことが出来る。なぜそう思うかというと、以下のBlockingTaskの中でpollの最後でブロッキング関数を実行しているから。

impl<T, R> Future for BlockingTask<T>
where
    T: FnOnce() -> R,
{
    type Output = R;

    fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<R> {
        let me = &mut *self;
        let func = me
            .func
            .take()
            .expect("[internal exception] blocking task ran twice.");

        // This is a little subtle:
        // For convenience, we'd like _every_ call tokio ever makes to Task::poll() to be budgeted
        // using coop. However, the way things are currently modeled, even running a blocking task
        // currently goes through Task::poll(), and so is subject to budgeting. That isn't really
        // what we want; a blocking task may itself want to run tasks (it might be a Worker!), so
        // we want it to start without any budgeting.
        crate::coop::stop();

        Poll::Ready(func())
    }
}

しかし本当ならば、ブロッキングタスクが終わってない場合にはNotReadyを返すというコードになってるべきではないかと思うのだが、なぜそうなっていないのだろうか?(他のところで担保されてるんだろうか?)

Futureのselect!マクロ

async fn task_one() { /* ... */ }
async fn task_two() { /* ... */ }

async fn race_tasks() {
    let t1 = task_one().fuse();
    let t2 = task_two().fuse();

    pin_mut!(t1, t2);

    select! {
        () = t1 => println!("task one completed first"),
        () = t2 => println!("task two completed first"),
    }
}

async fn count() {
    let mut a_fut = future::ready(4);
    let mut b_fut = future::ready(6);
    let mut total = 0;

    loop {
        select! {
            a = a_fut => total += a,
            b = b_fut => total += b,
            complete => break,
            default => unreachable!(), // never runs (futures are ready, then complete)
        };
    }
    assert_eq!(total, 10);
}

というサンプルコードを要点だけ説明する。

  • パターンマッチ文である。Readyが返ってきたものについて矢印の右辺を実行
  • 2番目のサンプルからわかるように、何度も呼ぶと残ってるFutureも実行する。1回だけなら、最初に返ってきたFutureだけ実行する
  • completeとdefaultは特殊なパターンマッチ(loopを抜けるために存在しているようなもの)
  • FuseFutureを要求する。なぜかというと、loopの中で使えるようにするため。一度ReadyとなったものはNotReadyを返してほしい。こうしないとパニックする可能性がある
  • ピンしているのは、Unpinを要求するため。Unpinを要求しているのは、select!自身がFutureを消費せず、select!後にもFutureにアクセスすることを許すため

2番目のa_fut, b_futはFusedFutureを実装していないように見えるがreadyが実装。

impl<T> FusedFuture for Ready<T> {
    fn is_terminated(&self) -> bool {
        self.0.is_none()
    }
}

なぜFuture::pollはPinをとるのか?

Futureトレイトはこのような形をしている。

trait Future {
    type Output;
    fn poll(
        // Note the change from `&mut self` to `Pin<&mut Self>`:
        self: Pin<&mut Self>,
        // and the change from `wake: fn()` to `cx: &mut Context<'_>`:
        cx: &mut Context<'_>,
    ) -> Poll<Self::Output>;
}

executerはpollを呼び、値がReadyとなることを期待する。もしNotReadyである場合は、下からwakeされるまで待つ。Futureというのは、ステートマシンだといえる。

本質的には、

trait SimpleFuture {
    type Output;
    fn poll(&mut self, wake: fn()) -> Poll<Self::Output>;
}

のようなものと変わりないが、wakeにコンテキスト(データ)が必要な場合に対応するために、Contextという下に隠す設計になっている。

なぜこのPinが必要かというと、async fn構文のためである。async fn構文を使うことによって、awaitするところを状態としたステートマシンを自動生成することが出来る。async fnは中で巨大なFutureを生成する。

さて、もしここで、async fnの中で参照を使っていて、そのFutureをmoveしようとするとどうなるかというと、この時にポインタの値が変わってしまう。moveが出来ないというのはRustプログラミングの中では致命的なので、この制約を取り去りたい。例えば実際に、mapする場合などはmoveする。

シンプルな例でいうとこのような場合である。

async {
    let mut x = [0; 128];
    let read_into_buf_fut = read_into_buf(&mut x);
    read_into_buf_fut.await;
    println!("{:?}", x);
}
struct ReadIntoBuf<'a> {
    buf: &'a mut [u8], // points to `x` below
}

struct AsyncFuture {
    x: [u8; 128],
    read_into_buf_fut: ReadIntoBuf<'what_lifetime?>,
}

それでは困るため、Pinを使う。Pinは、&mut T, &T, Box<T>などに対してTのアドレスを固定(ピン)する。やり方の一つとしては、Box::pinがある。

pub fn pin(x: T) -> Pin<Box<T>>

Pinの逆Unpinは、中のTを取り出したり、書き換えたりする能力を持つ。moveしても問題が起こらない型については、Unpinが実装されている。

ECR50: F「Relatively Prime Powers」

かなり難しい数学問題。ラフな解法のみ。

問題

xを因数分解した時、指数部分のGCDが1である場合、この数をエレガントと呼ぶ。

2からnまでの数字の中でエレガントな数は何個あるか?

制約: n=1018

解法

まず、指数部分のGCDが2であるというのはどういう状況かというと、例えば, 22 34のような状況である。これは、指数部分をGCDで割ることによって(21 32)2と出来る。

n以下にこの括弧部分が何個あるかというと、n1/2個である。

従って求める解は、包除原理より、

sum[i=1..] { n1/i mu(i) }

のように書ける。n1/iの部分はiが60くらいで0になるから、これはO(60)くらいで求まることになる。

オイラー関数関連

x以下でxと素なものの個数

これはオイラー関数で求まる。xの素因数をp1,p2,..pkとして、x(1-1/p1)(1-1/p2)...(1-1/pk)

既約分数との関係(言い換え)

これは、

1/x,2/x,...,x/xの中で「既約分数」の個数を求めてるのと同じこと。

x以下でxと素なものの和

例えば、x=15として、素なものを列挙すると1,2,4,7,8,11,13,14となる。個数は15(1-1/3)(1-1/5)=8より一致。

この和は、(1+14)+(2+13)+(4+11)+(7+8)で求まる。つまり、オイラー関数をphi(x)とすると、x*phi(x)/2となる。

証明: xと素な整数aをとる。gcd(a,x)=1より、gcd(a,x-a)=1である。これはつまり、x以下の整数のうちxと素なもののうち、aを見つけるとx-aも自動的に見つかることを意味している。今、a*a!=xであるから、上のようにペアをとることが可能。

ABC160: F「Distributing Integer」TDPC木の全方位木版

問題

n頂点の木が与えられる。

頂点kを根として、そこから隣接する頂点に番号を書いていく方法は何通りあるか?

全kについて答えよ。

制約: n=105

解法

考え方は、

rustforbeginners.hatenablog.com

とほぼ同じ。

CumRLのために

  1. dp値
  2. 部分木サイズのfact
  3. 部分木サイズ

を辺の値として持つ必要がある。これをCumRLでは、

  1. 累積積
  2. 累積積
  3. 累積和

でfoldする。

実装

838ms

競技プログラミングでたまに出る「全方位木DP」を解説 - テストステ論

で説明した抽象フレームワークで実装した。

fn solve() {
    let out = stdout();
    let mut out = BufWriter::new(out.lock());
    let mc = ModComb::new(300000, 1000000007);
    #[derive(Clone)]
    struct Z {
        mc: ModComb,
    }
    #[derive(Clone, Debug)]
    struct DP {
        dp: Mod,
        szfact: Mod,
        szsum: usize,
    }
    impl Foldable for Z {
        type T = DP;
        fn id() -> DP {
            DP {
                dp: Mod::new(1),
                szfact: Mod::new(1),
                szsum: 0,
            }
        }
        fn fold(acc: DP, x: DP) -> DP {
            DP {
                dp: acc.dp * x.dp,
                szfact: acc.szfact * x.szfact,
                szsum: acc.szsum + x.szsum,
            }
        }
    }
    impl ZenHoable for Z {
        type NVal = i64;
        type EVal = i64;
        fn f(&self, _: i64, _: i64, dp: &CumRL<Self>, l: usize, r: usize) -> DP {
            let LCum = dp.lcum(l);
            let RCum = dp.rcum(r);
            let dpfact = LCum.dp * RCum.dp;
            let szfact = LCum.szfact * RCum.szfact;
            // 部分木のサイズ和
            let szsum = LCum.szsum + RCum.szsum;

            DP {
                dp: Mod::new(self.mc.fact(szsum)) * dpfact / szfact,
                szfact: Mod::new(self.mc.fact(szsum+1)),
                szsum: szsum+1,
            }
        }
    }
    input!{
        n:usize,
        e:[(usize,usize);n-1],
    }
    let z = Z {
        mc: mc.clone(),
    };
    let mut g = ZenHo::new(z, vec![0; n]);
    for (u,v) in e {
        let u=u-1;
        let v=v-1;
        g.add_edge(u,v,0);
        g.add_edge(v,u,0);
    }
    g.build(0);
    for u in 0..n {
        let m = g.g[u].len();
        let mut dpfact = Mod::new(1);
        let mut szfact = Mod::new(1);
        let mut szsum = 0;
        for j in 0..m {
            let v = g.g[u][j];
            let dp = g.calc(v, u);
            dpfact *= dp.dp;
            szfact *= dp.szfact;
            szsum += dp.szsum;
        }
        let dp = Mod::new(mc.fact(szsum)) * dpfact / szfact;
        writeln!(out,"{}",dp);
    }
}

CFR629 (Div. 3): E「Tree Queries」

解法はわかっていたのに、最後のところでミスってAC出来なかった・・・。最近、見えてたんだけど解けなかったが多い。

https://codeforces.com/contest/1328/problem/E

問題

n頂点の木が与えられる。

m個のクエリk v1 v2 ... vkが与えられる。

各クエリについて、頂点uであって、v1,...vkが0-uパスの上にあるか、あるいは距離が1であるものが存在するか答えよ。

制約: n=105, m=105, クエリ中のvの合計は105

解法

viの中の関係だけで解きたいことは察せる。

0-uパスを考えて、もしviが条件を満たすとすると、viからの親の集合piは、木上で一直線になるはずである。

これをLCAを使って確かめる。

一直線であるというのは、piを深さ順にソートした場合に、p[i] = lca(p[i],p[i+1])がすべてのiについて成り立っているということである。O(vの数 logn)

コンテスト中は、この判定処理を、p[0] = lca(p[0], p[i])としてしまった。なぜか55まで通った。p[0]が深さ1のところで、p[1..]がすべてその部分木にある場合が反例。このケースが55まで入っていないというのがビビった。

実装

250msくらい。

fn solve() {
    let out = stdout();
    let mut out = BufWriter::new(out.lock());
    input!{
        new_stdin_parser = parser,
        n: usize, m: usize,
        e: [(usize,usize);n-1],
    }

    let mut g = HLDecomposition::new(n);
    for (u,v) in e {
        let u = u-1;
        let v = v-1;
        g.connect(u,v);
    }
    g.build(0);

    for _ in 0..m {
        input!{
            parser = parser,
            k: usize,
            v: [usize;k],
        }
        let mut xs = vec![];
        for i in 0..k {
            let x = v[i]-1;
            if x != 0 && g.par[x].unwrap() != 0 {
                let up = g.par[x].unwrap();
                let depth = g.depth[up] as i64;
                xs.push((depth, up));
            } else {
                // let depth = g.depth[x] as i64;
                // xs.push((depth, x));
            }
        }
        xs.sort();
        xs.dedup();

        let mut y = vec![];
        for x in xs {
            y.push(x.1);
        }
        // dbg!(&y);
        let m = y.len();
        let mut ok = true;
        if m > 0 {
            for i in 0..m-1 {
                let a = y[i];
                let b = y[i+1];
                let lca = g.lca(a, b);
                if lca != a {
                    ok = false;
                }
            }
        }
        writeln!(out,"{}",if ok {"YES"} else {"NO"});
    }
}