この記事は Haskell Advent Calendar 2023 の 18 日目の記事です．
はじめに
2023 年 10 月にリリースされた GHC 9.8 で，GHC の書換え規則における高階型のパターン照合が強化されました．
強化された書換え規則を実際に試してみたので，書換え規則の基本を振り返りつつ，簡単に紹介したいと思います．
Rewrite rules now support a limited form of higher order matching when a pattern variable is applied to distinct locally bound variables, as proposed in GHC Proposal #555. For example:
forall f. foo (\x -> f x)Now matches:
foo (\x -> x*2 + x)
GHC の書換え規則とは
GHC における書換え規則 (rewrite rule) は，コンパイル時に，決められたパターンに合致する式を別の式に置き換える機能です．
プログラムのコンパイル時の最適化手法として利用されています．
わかりやすい例として，map を 2 回適用する次のような式を考えてみます．
map f (map g xs) -- (1)
この式は，map を 1 回だけ適用する次の式と計算結果が同じになります．
map (f . g) xs -- (2)
書換え規則により式 (1) を式 (2) に書き換えてから計算することで，map の適用回数が減り，中間のデータ構造を減らすことができます．
このような，式の書換えにより中間のデータ構造を減らして計算を効率化する手法は，融合 (fusion) や森林伐採 (deforestation) として知られています．
なお，GHC における書換え規則は，その正しさをコンパイラが保証してくれるわけではありません．
式 (1) と (2) の等価性は簡単に証明できそうですが，そのような証明は自分で行わなければならないという点に注意が必要です．
GHC の書換え規則では，プログラム中に次のような RULES プラグマを書くことで，式の書換えを定義します．
{-# RULES
"map/map" forall f g xs. map f (map g xs) = map (f . g) xs
#-}
この RULES プラグマの例では，左辺 map f (map g xs) に合致する式を右辺 map (f . g) xs の形の式に書き換えるための書換え規則が定義されています．
map/map は書換え規則の名前です．
この書換え規則が動く様子を，簡単なプログラムで確かめてみます．
import Prelude hiding (map) -- Prelude の map を隠す
map :: (a -> b) -> [a] -> [b]
{-# NOINLINE map #-} -- 書換え規則が確実に適用されるようにインライン化を抑制
map _ [] = []
map f (x:xs) = f x : map f xs
-- 書換え規則を定義
{-# RULES
"map/map" forall f g xs. map f (map g xs) = map (f . g) xs
#-}
main = print $ map abs (map negate [1])
このプログラムを GHC でコンパイルする際，-ddump-rule-firings オプションや -ddump-rule-rewrites オプションを指定することで，コンパイラが書換え規則を適用する様子を確かめることができます．
また，書換え規則が適用されるようにするために，最適化を有効にするオプション -O も指定しておきます．
% ghc -O -ddump-rule-rewrites test.hs
[1 of 2] Compiling Main ( test.hs, test.o )
...
Rule fired
Rule: map/map
Module: (Main)
Before: Main.map
TyArg GHC.Num.Integer.Integer
TyArg GHC.Num.Integer.Integer
ValArg GHC.Num.Integer.integerAbs
ValArg Main.map
@GHC.Num.Integer.Integer
@GHC.Num.Integer.Integer
GHC.Num.Integer.integerNegate
(GHC.Base.build
@GHC.Num.Integer.Integer
(\ (@a_dMK)
(c_dML [OS=OneShot] :: GHC.Num.Integer.Integer -> a_dMK -> a_dMK)
(n_dMM [OS=OneShot] :: a_dMK) ->
c_dML (GHC.Num.Integer.IS 1#) n_dMM))
After: (\ (@b_aLg)
(@b_aLo)
(@a_aLj)
(f_aCk :: b_aLo -> b_aLg)
(g_aCl :: a_aLj -> b_aLo)
(xs_aCm :: [a_aLj]) ->
Main.map
@a_aLj @b_aLg (GHC.Base.. @b_aLo @b_aLg @a_aLj f_aCk g_aCl) xs_aCm)
@GHC.Num.Integer.Integer
@GHC.Num.Integer.Integer
@GHC.Num.Integer.Integer
GHC.Num.Integer.integerAbs
GHC.Num.Integer.integerNegate
(GHC.Base.build
@GHC.Num.Integer.Integer
(\ (@a_dMK)
(c_dML [OS=OneShot] :: GHC.Num.Integer.Integer -> a_dMK -> a_dMK)
(n_dMM [OS=OneShot] :: a_dMK) ->
c_dML (GHC.Num.Integer.IS 1#) n_dMM))
Cont: Stop[DiscArgCtxt] [GHC.Num.Integer.Integer]
[2 of 2] Linking test
% ./test
[1]
ちょっと見づらいですが，書換え規則 map/map が適用されている様子がわかります．
書換え前は map abs (negate [1]) で，書換え後は map (abs . negate) [1] (出力を見る限りは (\f g xs -> map (f . g) xs) abs negate [1] ですが) という形になっています．
出力の (GHC.Base.build ...) 部分は，GHC.Base.build による [1] の表現である build (\c n -> c 1 n) に相当します．
書換え規則と高階型
fold/build などの融合変換の研究は多くなされており，GHC のライブラリでも書換え規則が多用されています．
その中で，高度な書換えを実現しようとしたとき，高階型に対する書換え規則の表現力が必要となることが以前から指摘されていました．
GHC Proposal #555 には，提案の動機のひとつとして，次のような例が挙げられています．
{-# RULES
"map" [~1] forall f xs. map f xs = build (\c n -> foldr (\x ys -> c (f x) ys) n xs)
"mapList" [1] forall f. foldr (\x ys -> f x : ys) [] = map f
#-}
1 つ目の規則 map は，関数 map の適用を foldr と build を使った形に書き換えます．
2 つ目の規則 mapList は，foldr/build 融合が発生しなかった部分を単純な map の関数適用に戻すためのものです．
各規則の [~1] と [1] は，コンパイル時のどの段階で規則を適用するかを表すものですが，ここでは説明を省略します．
ここで，式 map f' (map g' xs) に対してこれらの書換え規則を適用することを考えてみます．
map f' (map g' xs)
|
| map 規則の適用
v
build (\c n ->
foldr (\x ys -> c (f' x) ys) n
(build (\c' n' -> foldr (\x' ys' -> c' (g' x') ys') n' xs)))
|
| foldr/build 融合変換
v
foldr (\x ys -> f' (g' x) : ys) [] xs
|
| mapList 規則の適用
v
map (\x -> f' (g' x)) xs
最後に規則 mapList の適用がありますが，ここで問題となる高階型のパターン照合が出てきます．
mapList 規則のパターン \x ys -> f x : ys が，適用対象の式の \x ys -> f' (g' x) : ys に照合する (つまり，f が \x -> f' (g' x) に照合する) ことになります．
GHC 9.8 で，このようなローカル変数束縛をもつ高階型のパターン照合を書換え規則で扱えるようになりました．
利用できる高階型のパターンの詳細については，GHC Proposal #555 を参照してください．
高階型のパターンをもつ書換え規則を試してみる
実際に，GHC 9.8 で高階型のパターンをもつ書換え規則を試してみたいと思います．
今回は，次のプログラムで動作を確認してみます．
関数 foo の適用を bar の適用に書き換える，高階型のパターンをもつ書換え規則 foo/bar を定義しています．
foo :: (Int -> Int) -> Int -> Int
{-# NOINLINE foo #-}
foo f x = x
bar :: (Int -> Int) -> Int -> Int
bar f x = f x
{-# RULES
"foo/bar" forall f x. foo (\y -> f y) x = bar f x
#-}
main = print $ foo (\x -> x + x) 1
比較のために，機能強化が入る前後の GHC 9.6 環境と GHC 9.8 環境での実行結果を確認してみます．
GHC 9.6 環境での実行結果:
% ghc --version
The Glorious Glasgow Haskell Compilation System, version 9.6.3
% ghc -O -ddump-rule-firings test.hs
[1 of 2] Compiling Main ( test.hs, test.o )
Rule fired: Class op show (BUILTIN)
Rule fired: Class op + (BUILTIN)
[2 of 2] Linking test
% ./test
1
GHC 9.8 環境での実行結果:
% ghc --version
The Glorious Glasgow Haskell Compilation System, version 9.8.1
% ghc -O -ddump-rule-firings test.hs
[1 of 2] Compiling Main ( test.hs, test.o )
Rule fired: Class op show (BUILTIN)
Rule fired: Class op + (BUILTIN)
Rule fired: foo/bar (Main)
Rule fired: +# (BUILTIN)
Rule fired: <# (BUILTIN)
[2 of 2] Linking test
% ./test
2
GHC 9.8 環境では foo/bar 規則が適用されていることがわかります．
期待通り，高階型のパターンをもつ書換え規則を動かすことができました！
まとめ
GHC 9.8 における書換え規則のパターン照合の強化内容について見てきました．
書換え規則はプログラム最適化の強力な手段であり，高階型のパターン照合が強化されたことでより一層強力になりました．
大いなる力には大いなる責任が伴うということで，使いどころを見極めながら使っていきたいですね．