AtCoder に登録したら解くべき精選過去問 10 問を C 言語(のコード中に埋め込んだ，アセンブラで生成した機械語)で解いてみた ② (第5問から第6問まで)
この記事は前回の続きです．前回は「C 言語で解いてみた」というほぼほぼ詐欺のようなタイトルでしたが，今回はちゃんと書きました．
そもそも何をしているの？
さて，ここら辺で C 言語の中に機械語を埋め込む方法について説明しておきましょう．提出コードの基本的な形は前回も書いたとおり次のようになっています．
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <string.h>
#include <stdio.h>
const char main_[] = { /* ここに解を求めるプログラムの中身を書く */ };
int main(void){
int fd = open("/dev/zero", O_RDONLY);
/* 返り値 */ (*solve)(/* 引数 */) = mmap(NULL, sizeof main_, PROT_WRITE | PROT_EXEC, MAP_PRIVATE, fd, 0);
memcpy(solve, main_, sizeof main_);
/* 入力を scanf などで受け取って solve 関数に渡し，返り値を出力する */
close(fd);
munmap(solve, sizeof main_);
return 0;
}これについて説明します．
まず，関数というのは配列です．その中には機械語がつらつらと書かれています．ならば，動的に確保した配列の中に機械語を書き込めば，それは関数になるわけです．
C で配列を動的に確保する方法は複数ありますが，その一つが mmap 関数を使うやり方です． mmap 関数とは次のような関数です．
#include <sys/mman.h> void *mmap(void *addr, size_t len, int prot, int flags, int fildes, off_t off);
第 5 引数 fildes にファイルを指定すると，その中身をメモリ上にマップして読み書きすることができる，そんな関数です．ここで，普通のファイルの代わりに /dev/zero を指定してやれば，メモリ上に 0 で埋め尽くされた領域が作られて使えるようになる，つまりメモリを確保することができるわけです．
またこの関数は，第 3 引数に prot という int 値をとります．ここでは PROT_READ ， PROT_WRITE ， PROT_EXEC の組み合わせまたは PROT_NONE が指定され，これによって確保される領域に対してどのようなアクセスができるか決まります． PROT_READ だけを指定すると読み込み専用になり， PROT_READ | PROT_WRITE を指定すると読み書きができます．そして， PROT_EXEC を指定すると実行が可能になるのです．
よって，このプログラムは， open 関数で /dev/zero を開き，それを用いて mmap 関数で実行可能なメモリ領域を確保し， memcpy 関数を使ってそこに機械語の書かれた配列の内容をコピーし，関数として呼び出しているわけです．
第5問 B - Some Sums
カラクリも分かったところで早速 5 問目に取り掛かりましょう．
入力は整数 n, a, b です． solve 関数は int solve(int n, int a, int b) とします．
1 以上 n 以下の整数を小さい順に見ていくのですが，この際にその各桁の和を計算していきます．
の各桁の和を
と表すと，
が 10 の倍数でなければ
，
が 10 で 1 回だけ割り切れるならば
であり，一般に
が 10 で
回割り切れるならば
が成り立ちます．
何回割り切れるか調べる際に div 命令を使いますが，これが実はちょっと楽にできるようになっています． div 命令は rdx レジスタと rax レジスタをあわせて 128bit としたものを割り算し，商を rax レジスタに，余りを rdx レジスタに格納します．ゆえに，そのままの状態で再度 div 命令を行うと，商をもう一度割り算することができるのです．実際にやってみると，ある数が 10 で何回割れるか知りたければ，
mov (調べたい数), %rax xor %rdx, %rdx div: div (整数10が格納されたレジスタ) inc (答えを格納するレジスタ) cmp $0, %rdx je div dec (答えを格納するレジスタ)
となります(一度余計にインクリメントしてからデクリメントしているのは，ジャンプ命令を減らして全体を短くするためです)．
1 以上 n 以下の数を走査するためのカウンタには rcx レジスタを使い，各桁の和は r8 レジスタに入れておきます．答えを格納する rax レジスタは割り算に利用するので r9 レジスタに一時的に格納し，同様に第 3 引数の入っている rdx レジスタの中身も r10 レジスタに避難させます．割る数 10 は r11 レジスタに入れておきます．いやー，レジスタがたくさんあるって素晴らしいですね．
最終的に solve 関数は以下のようになります．
mov $1, %rcx # <- 1 ... n xor %r8, %r8 # <- 各桁の和 xor %r9, %r9 # <- 答え mov %rdx, %r10 # <- 第 3 引数 (B) mov $10, %r11 # <- 割る数 10 loop: mov %rcx, %rax xor %rdx, %rdx div: div %r11 sub $9, %r8 cmp $0, %rdx je div add $10, %r8 cmp %r8, %rsi # 第 2 引数 (A) と比較 ja out cmp %r8, %r10 # 第 3 引数 (B) と比較 jb out add %rcx, %r9 # 範囲内なら答えに足す out: inc %rcx cmp %rdi, %rcx # カウンタが第 1 引数 (N) 以下の間繰り返す jbe loop mov %r9, %rax ret
第 6 問 B - Card Game for Two
ついにソートを書くときが来ました．配列を降順ソートして交互に見ていけばよいです．ソートアルゴリズムにも色々ありますが，今回はクイックソートを実装しようと思います．
入力は整数 n と配列 a です．配列 a の要素は 32bit 整数型とし， solve 関数は int solve(int n, int *a) とします．
クイックソートを書くうえで，今までと違う部分があります．それは，「コンパイル時に定められた個数だけの変数を使って書くことができない」という点です．クイックソートはよく再帰関数を使って実装されますが，再帰関数は呼び出される度に新たな自動変数を確保します．どこに確保するかというと，(再帰が深すぎたときに起こるスタックオーバーフローというエラーからも分かる通り)スタック領域です．今回の問題ではそのスタック領域を使ってみたいと思います．
といっても，スタックを使うのはそんなに難しくはありません．使うのは次の 2 つの命令です．
push (レジスタ) pop (レジスタ)
push 命令は，レジスタの中身をスタックにプッシュします． pop 命令は，スタックのトップをレジスタに格納してスタックから削除します(C++ の std::stack でいうところの top と pop を同時にやる)．これを使えば，関数の再帰呼び出しと実質同じことができるわけです(DFSとかも再帰関数でやるやり方とライブラリのスタックでやるやり方がありますね)．
また，このスタックはプログラムが様々な目的に使用しているので，初期状態が「空」ではない上に，関数呼び出し前と呼び出し後で同じ状態を保っていなければいけません．そこで，最初にスタックに何らかのマジックナンバーをプッシュしておき，マジックナンバーが取り出された時点で入れた分は全て取り出されたと分かるようにします．
クイックソートでは，配列内の範囲 [l, r) をソートするのが部分問題となります．よってスタックにプッシュするのは左端 l と右端 r です．今回は右端→左端の順にプッシュし，左端→右端の順にポップされるようにします．
push (マジックナンバー) push (全体の右端) push (全体の左端) loop: pop (レジスタ1) # <-今見たい部分の左端 cmp (マジックナンバー), (レジスタ1) je sortend # <- 取り出されたのがマジックナンバーならソート終了 pop (レジスタ2) # <-今見たい部分の右端 # ここで部分問題を解く jmp loop sortend:
配列のインデックスを保存しておくと，後で sizeof(int) = 4 倍して配列の先頭アドレスに足す操作をしなくてはなりません．そこでインデックスの代わりに全部アドレスで管理してしまうことにします．マジックナンバーは 0 にします．
スタックからポップした左端は r8 レジスタ，右端は r9 レジスタに入れることにします．ここまでで一通りコードを書くと次のようになります．
push $0 mov %rdi, %rax # <- 第 1 引数には配列の個数 n が入っている sal $2, %rax # <- 4 倍する add %rsi, %rax # <- 第 2 引数 a を足す: 配列の終端アドレス push %rax push %rsi loop: pop %r8 cmp $0, %r8 je sortend pop %r9 # ここで部分問題を解く jmp loop sortend:
さて，簡単のためピボットは先頭の要素とします．見たい範囲の左端 (r8 レジスタ)と右端 (r9 レジスタ) をそれぞれ r10 レジスタと r11 レジスタにコピーします．
mov %r8, %r10 mov %r9, %r11
r10 レジスタを右に動かしていき，ピボットより小さい値を探します．
mov (%r8), %eax # <- ピボット left: add $4, %r10 movsx (%r10), %ecx cmp %eax, %ecx jae left
同様に r11 レジスタを左に動かしていき，ピボットより大きい値を探します．
right: sub $4, %r11 movsx (%r11), %edx cmp %eax, %edx jb right
この時点で r10 が r11 よりも左にあれば， 2 つの値をスワップし，ふたたび r10 と r11 を動かしていきます．
cmp %r10, %r11 jbe end mov %ecx, (%r11) mov %edx, (%r10) jmp left # 上に戻る end:
r10 が r11 よりも右にあれば， r10 より左にはピボットより大きい値しかなく， r11 より右にはピボットより小さい値しかないので， r11 の場所の値とピボットを入れ替えて，次の部分問題を解きます．次の部分問題は，左端から r11 までと， r11 の右隣から右端までです．
mov %edx, (%r8) mov %eax, (%r11) push %r11 push %r8 add $4, %r11 push %r9 push %r11 jmp loop # 上に戻る
これを繰り返すことでソートができます．再帰の終了条件は，見ている部分の要素が 1 以下になったときで，これは部分問題を解く前に見ます．
mov %r9, %rcx # <- 右端のアドレスから sub %r8, %rcx # <- 左端のアドレスを引いて cmp $4, %rcx # <- 4 以下だったら jbe loop # <- continue
ソート部分を書き終えました．あとはこの配列を先頭から見て，足して引いて足して引いて……を繰り返したものが答えです．
xor %rax, %rax # <- 答えを入れるレジスタを 0 初期化 xor %rcx, %rcx # <- カウンタも 0 初期化 xor %rdx, %rdx # <- 0 なら足し， 1 なら引く loop2: cmp %rcx, %rdi je end cmp $0, %rdx # 足すか引くか je add sub (%rsi, %rcx, 4), %rax # 引く dec %rdx # 次は足す jmp sub add: add (%rsi, %rcx, 4), %rax # 足す inc %rdx #次は引く sub: inc %rcx # カウンタをインクリメント ja loop2 # 上に戻る
solve 関数全体と， AC コードのリンクです．https://atcoder.jp/contests/abc088/submissions/12041414
push $0 mov %rdi, %rax sal $2, %rax add %rsi, %rax push %rax push %rsi loop: pop %r8 cmp $0, %r8 je sortend pop %r9 mov %r9, %rcx sub %r8, %rcx cmp $4, %rcx jbe loop mov %r8, %r10 mov %r9, %r11 mov (%r8), %eax left: add $4, %r10 movsx (%r10), %ecx cmp %eax, %ecx jae left right: sub $4, %r11 movsx (%r11), %edx cmp %eax, %edx jb right cmp %r10, %r11 jbe end mov %ecx, (%r11) mov %edx, (%r10) jmp left end: mov %edx, (%r8) mov %eax, (%r11) push %r11 push %r8 add $4, %r11 push %r9 push %r11 jmp loop sortend: xor %rax, %rax xor %rcx, %rcx xor %rdx, %rdx loop2: cmp %rcx, %rdi je end cmp $0, %rdx je add sub (%rsi, %rcx, 4), %rax dec %rdx jmp sub add: add (%rsi, %rcx, 4), %rax inc %rdx sub: inc %rcx ja loop2 ret
お疲れ様でした．この調子で第 10 問までいければいいなと思っています．
