はじめに

昔々、あるところにPrimeFactor社という会社がありました。 PrimeFactor社の主力商品は素数です。メルセンヌ素数、双子素数、安全素数など様々な商品を取り揃えており、 ソースコードにハードコードして出荷して莫大な利益を上げていました。

prime :: Integer
prime = 1000000007

PrimeFactor社は研究開発に熱心な会社でした。社の売り上げの多くを技術開発に当て、 ついに量子計算を用いた素因数分解計算機の開発に成功しました。

ある日のこと、お客様から一通の苦情が送られてきました。 「おたくの素数を使ってガロア体$F_{p}$を作っていたら逆元が存在しなくてプログラムが例外を投げたぞ。 素数のふりして合成数を出荷するなんて詐欺じゃないか」 これは一大事です。お客様に出荷した素数を確認したところ $p = 3713287801 = 571 \times 2281 \times 2851$であることがわかりました。

原因を調査するために製品の製造過程を調べたところ、品質検査を行うプログラムがFermatテストを行なっていることがわかりました。

checkPrime :: Integer -> IO Bool
checkPrime n = 
    and <$> (replicateM 100 $ do
      a <- randomRIO (2,n)
      pure $ modPow a (n-1) n ==  1)

$p = 3713287801$はカーマイケル数だったので運悪く検査を通過してしまったようです。 PrimeFactor社は素数検査アルゴリズムをMiller-Rabin法に切り替えました。 しかしながら、Miller-Rabin法も確率的アルゴリズムであるが故に、非常に低い確率ではあるものの、 合成数を出荷してしまう可能性は排除できません。

そこで、PrimeFactor社は再発防止策として出荷する素数に証明をつけることにしました。 すなわち、誤って合成数が出荷されても埋め込まれたソースコードをコンパイルする時にエラーとなるようにしたいです。 幸いなことに、PrimeFactor社内では最新鋭の素因数分解計算機によって合成数の素因数分解を高速に行うことができます。 一方で出荷先でコンパイル時に行える計算資源は非常に限られています。 どのようにすれば良いでしょうか？

本題

さて、しょうもない茶番にお付き合いいただき、ありがとうございます。 今回のお題は「型レベル自然数を使って、与えられた数が素数であった時のみにコンパイルが通るプログラムを作る」です。

素数であることの証拠

※しばらく数学的な議論が続きます。証明等に興味のない方はこのセクションの最後の素数検証アルゴリズムまで飛ばしても大丈夫です。

素数判定アルゴリズムを用いた手法

素数判定アルゴリズムは証拠として使えるでしょうか。

• 試し割り: これは$1 \leq i \leq \sqrt{n}$を満たす各$i$について$i$が$n$で割り切れないことを試すということです。 しかし、この計算は$O(\sqrt{n})$なのでをコンパイル時に行うのは現実的ではありません。（実行時でも厳しい）
• 確率的素数判定法: Miller-Rabin法などの確率的アルゴリズムは合成数を高速に高い確率で検出できますが、素数であることの証明には使えません。 また、コンパイル時に行うため、乱数生成が難しいという問題もあります。

補助データを用いた手法

入力$n$のみから素数であることを高速に証明するのは難しそうです。 補助データを用いる場合はどうでしょうか。

今回は以下の定理を利用します。

整数$n$が素数である$\iff$ある$a$が存在して$\min \{ i \mid 1 < i, a^{i} \mod n = 1\} = n-1$.

証明）

• $n$が素数である時は$(Z/Z_n)^* = \{1, \dots, n-1\}$が位数$n-1$の巡回群となることから従います。
• $n$が合成数かつ、そのような$a$が存在したと仮定します。$a^{n-1} \mod n = 1$より$a$は$n$と互いに素です。 この時、$\lambda(n)$をカーマイケル関数とすると $n - 1 \leq \lambda(n)$が成り立ちます。一方でオイラーのトーシェント関数を$\phi(n)$とすると $\lambda(n) \leq \phi(n)$です。今$n$が合成数なので$\phi(n) < n-1$です。これは矛盾です。

$\min \{ i \mid 1 < i, a^{i} \mod n = 1\}$は$a$のmultiplicative orderと呼ばれ、この記事では$\mathrm{order}_n(a)$ と書きます。

この定理から、$a$を補助データとして、$\mathrm{order}_n(a) = n-1$を確かめれば$n$が素数であることの証拠 になることがわかります。

$\mathrm{order}_n(a) = n - 1$を確かめる方法について考えます。 まず$a^{n-1} \mod n = 1$であることを確かめます。この時、$\mathrm{order}_n(a)$は$n-1$の 約数であることが証明できます。

証明） $k = \mathrm{order}_n(a)$とします。定義より、$a^k = 1$かつ$0 < i < k$となる$i$に対し$a^i \mod n \neq 1$です。 $k$が$n-1$の約数でないとすると、$n - 1 = qk + r$となる$0 < r < k$が存在します。 ここで$1 = a^{n-1} = a^{qk+r} = (a^k)^q a^r = 1 a^r = a^r \neq 1 (\mod n)$より矛盾します。

従って$\mathrm{order}_n(a) < n$の時、ある$n-1$の素因数$p$が存在して、$a^{\frac{n-1}{p}} \mod n = 1$です。

証明) $\mathrm{order}_n(a) q = n - 1$となる$q > 1$が存在します。$q$の素因数（$n-1$の素因数でもある)を$p$とすれば $a^{\frac{n-1}{p}} \mod n= 1$です。

このことから、$n -1$の全ての素因数$p_1 \dots p_m$について$a^{\frac{n-1}{p_i}} \mod n \neq 1$を確かめれば $\mathrm{order}_n(a) = n-1$となります。

素数検証アルゴリズム

まとめると、次のようなアルゴリズムで素数であることを検証できます。

入力：整数$n$, $a$, $p_1, \dots, p_m$

出力: 整数$n$が素数であるかつ、$a$, $p_1, \dots,p_m$がその証拠となる時True、そうでない時、False

アルゴリズム:

1. $p_1, \dots, p_m$が$n$の素因数全体でない場合、Falseを出力
2. $a^{n-1} \mod n \neq 1$の時、Falseを出力

3. 各 $1 \leq i \leq m$について

   1. $a^{\frac{n-1}{p_i}} \mod n = 1$ならばFalseを出力

4. Trueを出力

素数検証アルゴリズムの型レベル実装

それでは先ほどのアルゴリズムをHaskellの型レベル自然数を用いて実装していきます。

まずは今回使うモジュールと言語拡張一覧です。

{-# LANGUAGE GADTs, DataKinds, UndecidableInstances, KindSignatures, TypeOperators, TypeFamilies #-}
module Prime where
import Data.Type.Bool
import Data.Type.Equality
import GHC.TypeLits
import Data.Kind
import Data.Proxy

素因数チェック

まず最初に、$p_1,\ldots, p_n$が$n-1$の素因数全体であることを確かめます。

type family IsFactorization (n :: Nat) (l :: [Nat]) :: Bool where
    IsFactorization 0 l = False
    IsFactorization 1 '[] = True
    IsFactorization 1 l = False
    IsFactorization n (p ': l) = Mod n p == 0 && IsFactorizationSub n p l (Mod n p)
type family IsFactorization Sub n p l m where
    IsFactorizationSub 0 _ _ _ = TypeError (Text "Invalid argument 0")
    IsFactorizationSub n p l 0 = IsFactorizationSub (Div n p) p l (Mod (Div n p) p)
    IsFactorizationSub n p l _ = IsFactorization n l

IsFactorization n lは自然数nと素数のリストlを受け取り、lがnの素因数全体かを確かめる型レベル関数です。 IsFactorizaionSub n p l mは補助関数です。ここで注意することがあって、 型レベル計算は名前呼びの項書き換え系に型レベル関数ごとに引数の組み合わせでメモ化を行うものだと考えれば良いです。 そのため、一行目のIsFactorizationSub 0 _ _ _ = ...がない場合、2行目の IsFactorizationSub n p l 0 = IsFactorizationSub (Div n p) p l (Mod (Div n p) p)の第一引数は遅延評価されてDiv ... (Div (Div n p) p) ... pと積み重なってしまいます。 一行目がある場合、Div n pが0であるかを確かめる必要があるので、nを正格評価します。それに寄ってこのような積み重ねを回避できます。

ModPow

次に$a^d \mod n$を計算する型レベル関数ModPow a d nを定義します。

type family ModPow (a :: Nat) (d :: Nat) (n :: Nat) :: Nat where
    ModPow a 0 n = 1
    ModPow a d n = ModPowSub a d n (Mod d 2 == 0)

type family ModPowSub a d n b where
    ModPowSub a d n True = Square (ModPow a (Div d 2) n) n
    ModPowSub a d n False =
        Mod (Square (ModPow a (Div d 2) n) n GHC.TypeLits.* a) n

type family Square a n where
    Square 1 n = 1
    Square 0 n = 0
    Square a n = Mod (a GHC.TypeLits.* a) n

Square a nは$a \times a \mod n$を計算します。 ここでも先ほど説明したように、aを正格評価するためにa = 1, 0のケースを特別扱いしています。

位数(order)のチェック

最後に$\mathrm{order}_n(a) = n - 1$を確かめます。

type family IsPrimeCert (n :: Nat) (l :: [Nat]) (a :: Nat) :: Bool
    where
    IsPrimeCert n l a = ModPow a (n - 1) n == 1 && IsPrimeCertSub n l a

type family IsPrimeCertSub n l a
    where
    IsPrimeCertSub n '[] a = True
    IsPrimeCertSub n (p ':l) a = 
        Not (ModPow a (Div (n - 1) p) n == 1) && IsPrimeCertSub n l a

素数検証

a, lがnが素数であることの証拠になっていることを表すConstraintを CheckPrimeCert n l aで実装します。

type family CheckPrimeCert (n :: Nat) (l :: [Nat]) (a :: Nat) :: Constraint
    where
    CheckPrimeCert n l a =
        ((IsFactorization (n - 1) l
          || TypeError (ShowType l :<>: Text " is not the prime factors of ":<>: ShowType n))
        &&
        (IsPrimeCert n l a
          || TypeError (ShowType a :<>: Text " is not a generator of the multiplicative group modulo " :<>: ShowType n))) ~ True

TypeError Constraintを使うことでわかりやすい型エラーメッセージを出力することができます。

ユーザAPI

最後の、nが素数であることを表す型クラスKnownPrime nとその証拠を保持するGADT PrimeCert nを定義します。

class KnownPrime n where
  primeCert :: PrimeCert n

data PrimeCert (n :: Nat) where
  PrimeCert :: CheckPrimeCert n l a => Proxy l -> Proxy a -> PrimeCert n

使い方

nが素数であることを宣言するにはKnownPrime nのインスタンスを宣言してあげます。

{-# LANGUAGE DataKinds, TypeApplications #-}
module User where
import Prime
import Data.Proxy

instance KnownPrime 5 where
    primeCert = PrimeCert (Proxy @'[2]) (Proxy @2)
instance KnownPrime 11 where
    primeCert = PrimeCert (Proxy @'[2,5]) (Proxy @2)
instance KnownPrime 1000000007 where
    primeCert = PrimeCert (Proxy @'[2,500000003]) (Proxy @5)

-- TypeErrors
instance KnownPrime 13 where
    primeCert = PrimeCert (Proxy @'[2,5]) (Proxy @2)
instance KnownPrime 57 where
    primeCert = PrimeCert (Proxy @'[2, 7]) (Proxy @2)

コンパイルしてみると、正しい証拠を与えた時には型エラーが発生せず、誤った証拠を与えた時は型エラーになることがわかります。 また、$10^9 + 7$といった多少大きい数に対しても現実的な時間で型検査がすることが確認できます。

$ ghc User.hs
Loaded package environment from /Users/autotaker/playground/modulo-no-overhead/.ghc.environment.x86_64-darwin-8.6.3
[1 of 2] Compiling Prime            ( Prime.hs, Prime.o )
[2 of 2] Compiling User             ( User.hs, User.o )

User.hs:16:17: error:
    • '[2, 5] is not the prime factors of 13
    • In the expression: PrimeCert (Proxy @'[2, 5]) (Proxy @2)
      In an equation for ‘primeCert’:
          primeCert = PrimeCert (Proxy @'[2, 5]) (Proxy @2)
      In the instance declaration for ‘KnownPrime 13’
   |
16 |     primeCert = PrimeCert (Proxy @'[2,5]) (Proxy @2)
   |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

User.hs:19:17: error:
    • 2 is not a generator of the multiplicative group modulo 57
    • In the expression: PrimeCert (Proxy @'[2, 7]) (Proxy @2)
      In an equation for ‘primeCert’:
          primeCert = PrimeCert (Proxy @'[2, 7]) (Proxy @2)
      In the instance declaration for ‘KnownPrime 57’
   |
19 |     primeCert = PrimeCert (Proxy @'[2, 7]) (Proxy @2)
   |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Limitations

素因数リストの検証

察しの良い読者の方はお気づきかもしれませんが、 今回の話では証拠として与えた素因数リストが本当に素数のリストであることは 検証していません。それさえも検証したい場合は再帰的にKnownPrime pの制約を与えれば良いです。

type family KnownPrimes l :: Constraint where
    KnownPrimes '[] = ()
    KnownPrimes (p ': ps) = (KnownPrime p, KnownPrimes ps)
type family CheckPrimeCert n l a :: Constraint where
    CheckPrimeCert n l a =
      (KnownPrimes l
      , (IsFactorization (n - 1) l 
        && IsPrimeCert n l a) ~ True)

ただこの場合、インスタンス宣言の数が少し多くなるという問題があるので今回は採用しませんでした。

ボイラーポレートの自動生成

KnownPrime型クラスのインスタンスはTemplate Haskellを使って自動生成できそうです。ただ今回の記事では執筆時間の都合上割愛しました。

まとめ

型レベル計算によって素数である時のみコンパイルが通るようなプログラムを実現しました。型レベル計算は遅いので、補助データを用いることで型レベル計算の計算量を抑える必要があります。それらの補助データはユーザが与える、あるいはTemplate Haskellによって導出すれば現実的な時間で動作することが期待できます。 また、型レベル計算の実装時には型レベル自然数が正格に評価されるように注意する必要があります。

今回作成したプログラムはgistにおいてあります。