Attempt to use unification-fd.

Librarian.hs

{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Librarian where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Unification
import Control.Unification.IntVar

import Data.Functor.Fixedpoint
import Data.Functor.Identity
import qualified Data.Map as M

import Safe (atMay)

type Name = String

data Expr
  = Ref Int
  | Global Name
  | Lam (Maybe Type) Expr
  | App Expr Expr
  | Num Integer
  | Bool Bool

data TypeF a
  = LamT a a
  | NumT
  | BoolT
  deriving (Foldable, Functor, Show, Traversable)

instance Unifiable TypeF where
  zipMatch (LamT x x') (LamT y y') =
    Just (LamT (Right (x, y)) (Right (x', y')))
  zipMatch NumT NumT   = Just NumT
  zipMatch BoolT BoolT = Just BoolT
  zipMatch _ _         = Nothing

type Type = Fix TypeF

type UType = UTerm TypeF IntVar

type Environment = M.Map Name Type

type Infer = ReaderT ([UType], Environment) (ExceptT String (IntBindingT TypeF Identity))

instance Fallible TypeF IntVar a where
  occursFailure   = error "infinite type"
  mismatchFailure = error "type mismatch"

localEnv :: Infer [UType]
localEnv = fst <$> ask

globalEnv :: Infer Environment
globalEnv = snd <$> ask

localType :: UType -> Infer UType -> Infer UType
localType t = local (\(ts, gs) -> (t:ts, gs))

infer :: Expr -> Infer UType
infer (Ref x) = do
  env <- localEnv
  case atMay env x of
    Just x' -> return x'
    Nothing -> error $ "unbound reference: " ++ show x
infer (Global x) = do
  env <- globalEnv
  case M.lookup x env of
    Just x' -> return $ unfreeze x'
    Nothing -> error $ "unbound variable: " ++ x
infer (Lam (Just t) x) = do
  let t' = unfreeze t
  tx <- localType t' $ infer x
  return $ (UTerm (LamT t' tx))
infer (Lam Nothing x) = do
  t  <- UVar <$> lift (lift freeVar)
  tx <- localType t $ infer x
  return $ (UTerm (LamT t tx))
infer (App x y) = do
  tx <- infer x
  ty <- infer y
  t  <- UVar <$> lift (lift freeVar)
  lift $ tx `unify` UTerm (LamT ty t)
infer (Num _)  = return $ UTerm NumT
infer (Bool _) = return $ UTerm BoolT