突如として、型推論はどんな風にしたら実装できるのか知りたくなった。型理論とか全く知らないが、実装して動かしてみて原理を知りたいと思った。そこであれこれググる:

Haskellで実装してみる

上記のリンクたちを参考にして、Haskellで構文木を渡すとその型を推論する関数を作ってみる。言語として使えるようになるまでにするには敷居が高いので、ひとまず型の推論だけで、型検査などは行わない。

評価したい式の構文木は構築されていて与えられるものとして、パーサなどの実装も行わない(パーサ編はこちら

(内容はほとんど上記の参照したサイトで解説されている内容そのままで、主にHaskellでどうやって型変数を扱うのか、どうやってSTモナドを扱うかと格闘していた気がする。)

構文木のノードの定義

ひとまず扱える基本的な型は整数のみで、構文は1引数の関数とその呼び出し、また変数の参照のみ、とする:

-- Node.hs
module Node (Expr(..)) where

data Expr = Natural Integer  -- 整数:値
          | Var String       -- 変数:名前
          | Fun String Expr  -- 関数:パラメータ名、結果の式
          | App Expr Expr    -- 適用:関数、引数
  deriving (Show)
  • 関数は全て引数は1つのみで、カリー化されてるものとする
  • 関数適用も同様に、1度に1つしか適用しないものとする

型の定義

式が持つ型は、整数型TIntと関数型TFun(引数の型と結果の型を持つ)と、型変数TVar(識別番号を持つ)。

-- Type.hs
module Type (Type(..)) where

data Type = TInt            -- 整数
          | TFun Type Type  -- 関数:引数の型、結果の型
          | TVar Int        -- 型変数:識別番号
  deriving (Eq)

instance Show Type where
  show TInt        = "Int"
  show (TFun p e)  = pp p ++ " -> " ++ show e
    where pp fun@(TFun _ _) = "(" ++ show fun ++ ")"
          pp t              = show t
  show (TVar i)    = "t" ++ show i

TVarは未確定の型を保持する型変数で、後述の単一化(unify)で決定されることになる(決定せずに残る場合もある)。

型変数の型が確定した場合にその型を保持するのは別テーブルで管理し、識別番号で参照する。

型推論

型推論を行う関数inferは、組み込み関数など型の確定しているグローバルな環境[(String, Type)]と、推論したい式の構文木Exprを受け取り、推論した型Typeを返す:

-- TypeInferencer.hs
module TypeInferencer (infer) where

import Prelude hiding (lookup)
import Control.Monad.ST (ST, runST)
import Data.Map (empty, fromList, insert, lookup, Map)
import Data.STRef (newSTRef, readSTRef, writeSTRef, STRef)

import Node (Expr(..))
import Type (Type(..))

-- Environment
type Env = Map String Type

-- (next var index, {index=>type})
type VarInfo = (Int, Map Int Type)

infer :: [(String, Type)] -> Expr -> Type
infer env expr = runST $ do
  varInfoRef <- newSTRef (0, empty)
  t <- doInfer (fromList env) varInfoRef expr
  (_, varDict) <- readSTRef varInfoRef
  return $ refer t varDict

doInfer :: Env -> STRef s VarInfo -> Expr -> ST s Type
doInfer env varInfoRef expr =
  case expr of
    Natural i  -> return TInt
    Var x -> do
      case lookup x env of
        Just t  -> return t
        Nothing -> error ("not found: " ++ x)
    Fun parm e -> do
      tparm <- createVar varInfoRef
      te <- doInfer (insert parm tparm env) varInfoRef e
      return $ TFun tparm te
    App f arg -> do
      funType <- doInfer env varInfoRef f
      argType <- doInfer env varInfoRef arg
      retType <- createVar varInfoRef
      unify funType (TFun argType retType) varInfoRef
      return retType
  • 型変数の値が決定した時に状態を書き換えるために、STモナドを使用することにして、実際の型推論 doInfer ではSTモナド内で行う
    • 結果は runST で取り出す
  • 型の推論は、関数の適用時に関数の引数の型から単一化して、どちらかが型変数の場合には反映させる

単一化

-- TypeInferencer.hs
unify :: Type -> Type -> STRef s VarInfo -> ST s ()
unify (TFun p1 e1) (TFun p2 e2) varInfoRef = do
  unify p1 p2 varInfoRef
  unify e1 e2 varInfoRef
unify t1@(TVar i1) t2@(TVar i2) varInfoRef
  | i1 == i2  = return ()
unify (TVar i1) t2 varInfoRef = unifyVar i1 t2 varInfoRef
unify t1 (TVar i2) varInfoRef = unifyVar i2 t1 varInfoRef
unify t1 t2 varInfoRef
  | t1 == t2  = return ()
  | otherwise = cannotUnify t1 t2 varInfoRef

unifyVar :: Int -> Type -> STRef s VarInfo -> ST s ()
unifyVar index type2 varInfoRef = do
  isOccur <- occur type2 index varInfoRef
  if isOccur
    then  error "occurs error"
    else do
      (nextIdx, varMap) <- readSTRef varInfoRef
      case lookup index varMap of
        Just vt  -> unify vt type2 varInfoRef
        Nothing  -> writeSTRef varInfoRef (nextIdx, insert index type2 varMap)

出現確認

-- TypeInferencer.hs
occur :: Type -> Int -> STRef s VarInfo -> ST s Bool
occur (TFun p e) n varInfoRef = (||) <$> occur p n varInfoRef <*> occur e n varInfoRef
occur (TVar i) n varInfoRef
  | i == n    = return True
  | otherwise = do
      (_, varMap) <- readSTRef varInfoRef
      case lookup i varMap of
        Just vt  -> occur vt n varInfoRef
        Nothing  -> return False
occur _ _ _   = return False

細々

-- TypeInferencer.hs
createVar :: STRef s VarInfo -> ST s Type
createVar varInfoRef = do
  (nextIdx, varMap) <- readSTRef varInfoRef
  writeSTRef varInfoRef (nextIdx + 1, varMap)
  return $ TVar nextIdx

refer :: Type -> Map Int Type -> Type
refer (TFun p e) varMap = TFun (refer p varMap) (refer e varMap)
refer t@(TVar v) varMap = case lookup v varMap of
  Just vt  -> refer vt varMap
  Nothing  -> t
refer t varMap          = t

cannotUnify t1 t2 varInfoRef = do
  (_, varMap) <- readSTRef varInfoRef
  error ("cannot unify: " ++ show (refer t1 varMap) ++ " <=> " ++ show (refer t2 varMap))

動かしてみる

Prelude> :l TypeInferencer.hs
[1 of 3] Compiling Node             ( Node.hs, interpreted )
[2 of 3] Compiling Type             ( Type.hs, interpreted )
[3 of 3] Compiling TypeInferencer   ( TypeInferencer.hs, interpreted )
Ok, modules loaded: Type, TypeInferencer, Node.
*TypeInferencer> let globalEnv = [("+", TFun TInt (TFun TInt TInt))]  -- (+) :: Int -> (Int -> Int)
*TypeInferencer> infer globalEnv $ Natural 123
Int  -- 123 :: Int
*TypeInferencer> infer globalEnv $ Fun "x" (App (App (Var "+") (Var "x")) (Natural 1))
Int -> Int  -- \x -> x + 1 :: Int -> Int

関数合成の型を推論してみる:

*TypeInferencer> infer globalEnv $ Fun "f" (Fun "g" (Fun "x" (App (Var "f") (App (Var "g") (Var "x")))))
(t3 -> t4) -> (t2 -> t3) -> t2 -> t4  -- \f -> \g -> \x -> f (g x) :: (b -> c) -> (a -> b) -> a -> c
*TypeInferencer> :t (.)
(.) :: (b -> c) -> (a -> b) -> a -> c

ソースはこちら

雑感

  • 確かに、ここまでだけだと実装は非常に簡単にできる
    • 関数が常に引数が1つという制約でかなり簡略化されてると思う
  • 実際に型検査しようとすると、型クラスを導入する必要があるんじゃないかと思い、そうなると結構大変になるんじゃないかと想像