Haskell doodle made while studying unification-fd's source code.

Paper.hs

-- This doesn't work correctly.

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeSynonymInstances #-}

-- https://ro-che.info/articles/2017-06-17-generic-unification
-- http://hackage.haskell.org/package/unification-fd-0.10.0.1/docs/Control-Unification.html
-- https://www.schoolofhaskell.com/user/bartosz/understanding-algebras
-- https://www.michaelpj.com/blog/2018/04/08/catamorphic-lc-interpreter.html

module Paper where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State

import qualified Data.IntMap as IM
import qualified Data.Map as M

import Safe (atMay)

newtype Fix a = Fix { unFix :: a (Fix a) }

cata :: Functor f => (f a -> a) -> Fix f -> a
cata alg = alg . fmap (cata alg) . unFix

type Name = String

data ExprF a
  = Ref Int
  | Global Name
  | Lam (Maybe Type) a
  | App a a
  | Num Integer
  | Bool Bool
  deriving Functor

type Expr = Fix ExprF

ref :: Int -> Expr
ref = Fix . Ref

global :: Name -> Expr
global = Fix . Global

lam :: Expr -> Expr
lam = Fix . Lam Nothing

lamt :: Type -> Expr -> Expr
lamt t = Fix . Lam (Just t)

app :: Expr -> Expr -> Expr
app x y = Fix $ App x y

num :: Integer -> Expr
num = Fix . Num

bool :: Bool -> Expr
bool = Fix . Bool

data TypeF a
  = TArr a a
  | TNum
  | TBool
  deriving (Foldable, Functor, Show, Traversable)

zipMatch :: TypeF a -> TypeF a -> Maybe (TypeF (Either a (a, a)))
zipMatch (TArr x x') (TArr y y') =
  Just (TArr (Right (x, y)) (Right (x', y')))
zipMatch TNum TNum   = Just TNum
zipMatch TBool TBool = Just TBool
zipMatch _ _         = Nothing

type Type = Fix TypeF

deriving instance Show Type

tarr :: Type -> Type -> Type
tarr x y = Fix $ TArr x y

tnum :: Type
tnum = Fix TNum

tbool :: Type
tbool = Fix TBool

data UTerm a
  = UVar Int
  | UTerm (a (UTerm a))

type UType = UTerm TypeF

deriving instance Show (UTerm TypeF)

unfreeze :: Functor a => Fix a -> UTerm a
unfreeze = UTerm . fmap unfreeze . unFix

type Environment = M.Map Name UType
type Constraint  = (UType, UType)

data Binding = Binding
  { next :: Int
  , bind :: IM.IntMap UType
  } deriving Show

type Infer = StateT Binding (ReaderT ([UType], Environment) (Except String))

localEnv :: ([UType], Environment) -> [UType]
localEnv = fst

globalEnv :: ([UType], Environment) -> Environment
globalEnv = snd

localType :: UType -> Infer a -> Infer a
localType x = local (\(xs, ys) -> ((x:xs), ys))

newType :: Infer UType
newType = do
  Binding i bs <- get
  put $ Binding (i + 1) bs
  return $ UVar i

bindType :: Int -> UType -> Infer UType
bindType v x = do
  Binding i bs <- get
  put $ Binding i (IM.insert v x bs)
  return x

unify :: UType -> UType -> Infer UType
unify (UTerm x) (UTerm y) = match x y
unify v@(UVar _) x = unify x v
unify x (UVar v) = occurs v x >>= \case
  True  -> throwError "infinite type"
  False -> bindType v x

occurs :: Int -> UType -> Infer Bool
occurs x (UTerm y)           = or <$> mapM (occurs x) y
occurs x (UVar x') | x == x' = return True
occurs _ _                   = return False

match :: TypeF UType -> TypeF UType -> Infer UType
match x y =
  case zipMatch x y of
    Nothing -> throwError $ "cannot match type `" ++ show x ++ "` with `" ++ show y
    Just xy -> UTerm <$> flip mapM xy \case
      Left x'        -> return x'
      Right (x', y') -> unify x' y'

algebra :: ExprF (Infer UType) -> Infer UType
algebra (Ref x) = do
  env <- localEnv <$> ask
  case atMay env x of
    Just x' -> return x'
    Nothing -> throwError $ "unbound reference: " ++ show x
algebra (Global x) = do
  env <- globalEnv <$> ask
  case M.lookup x env of
    Just x' -> return x'
    Nothing -> throwError $ "unbound variable: " ++ x
algebra (Lam (Just t) x) = do
  let t' = unfreeze t
  tx <- localType t' x
  return $ (UTerm (TArr t' tx))
algebra (Lam Nothing x) = do
  t  <- newType
  tx <- localType t x
  return $ (UTerm (TArr t tx))
algebra (App x y) = do
  tx <- x
  ty <- y
  t  <- newType
  unify tx (UTerm (TArr ty t))
algebra (Num _)  = return $ UTerm TNum
algebra (Bool _) = return $ UTerm TBool

applyBindings :: UType -> Infer UType
applyBindings (UVar v) = do
  x <- IM.lookup v <$> bind <$> get
  case x of
    Just x' -> return x'
    Nothing -> return $ UVar v
applyBindings (UTerm x) = UTerm <$> mapM applyBindings x

gather :: Expr -> Except String UType
gather x =
  fst <$>
    runReaderT (runStateT (cata algebra x) (Binding 0 IM.empty)) ([], M.empty)

check :: Expr -> Except String (UType, Binding)
check x = runReaderT (runStateT result (Binding 0 IM.empty)) ([], M.empty)
  where result = cata algebra x >>= applyBindings

-- gather $ lamt tnum (ref 0)