Glasgow Haskell Compiler の書換え規則における高階型のパターン照合

この記事は [Haskell Advent Calendar 2023](https://qiita.com/advent-calendar/2023/haskell) の 18 日目の記事です．

## はじめに

2023 年 10 月にリリースされた GHC 9.8 で，GHC の書換え規則における高階型のパターン照合が強化されました．
強化された書換え規則を実際に試してみたので，書換え規則の基本を振り返りつつ，簡単に紹介したいと思います．

[GHC 9.8.1 リリースノート](https://downloads.haskell.org/~ghc/9.8.1/docs/users_guide/9.8.1-notes.html) より:

> 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](https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0555-template-patterns.rst). For example:
> 
> ```haskell
> forall f. foo (\x -> f x)
> ```
> 
> Now matches:
> 
> ```haskell
> foo (\x -> x*2 + x)
> ```

## GHC の書換え規則とは

GHC における書換え規則 (rewrite rule) は，コンパイル時に，決められたパターンに合致する式を別の式に置き換える機能です．
プログラムのコンパイル時の最適化手法として利用されています．

* ["Rewrite rules" - GHC User’s Guide](https://downloads.haskell.org/ghc/latest/docs/users_guide/exts/rewrite_rules.html)

わかりやすい例として，`map` を 2 回適用する次のような式を考えてみます．

```haskell
map f (map g xs)      -- (1)
```

この式は，`map` を 1 回だけ適用する次の式と計算結果が同じになります．

```haskell
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` は書換え規則の名前です．

この書換え規則が動く様子を，簡単なプログラムで確かめてみます．

```haskell
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](https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0555-template-patterns.rst) には，提案の動機のひとつとして，次のような例が挙げられています．

```
{-# 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](https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0555-template-patterns.rst) を参照してください．

## 高階型のパターンをもつ書換え規則を試してみる

実際に，GHC 9.8 で高階型のパターンをもつ書換え規則を試してみたいと思います．

今回は，次のプログラムで動作を確認してみます．
関数 `foo` の適用を `bar` の適用に書き換える，高階型のパターンをもつ書換え規則 `foo/bar` を定義しています．

```haskell
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 における書換え規則のパターン照合の強化内容について見てきました．
書換え規則はプログラム最適化の強力な手段であり，高階型のパターン照合が強化されたことでより一層強力になりました．
大いなる力には大いなる責任が伴うということで，使いどころを見極めながら使っていきたいですね．

## 参考

* [Version 9.8.1 Release Notes](https://downloads.haskell.org/ghc/9.8.1/docs/users_guide/9.8.1-notes.html)
* [GHC Proposal #555: Higher Order Patterns in Rewrite Rules](https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0555-template-patterns.rst)
* ["Rewrite rules" - GHC User’s Guide](https://downloads.haskell.org/ghc/latest/docs/users_guide/exts/rewrite_rules.html) (日本語訳: [書き換え規則 - 栄光のグラスゴーHaskellコンパイルシステム利用の手引き](http://www.kotha.net/ghcguide_ja/latest/rewrite-rules.html))
* [GHC 9.8の新機能 - Zenn](https://zenn.dev/mod_poppo/articles/whats-new-in-ghc-9-8)
* [Rewrite Rulesについて軽く - Qiita](https://qiita.com/ruicc/items/4c33fe2501fa956cd9a5)