Type inference using Bluefin

SoundEager.hs

{-# OPTIONS_GHC -Wno-orphans #-}

module SoundEager where

import Bluefin.Compound (Handle, mapHandle, useImpl, useImplIn)
import Bluefin.Eff (Eff, Effects, runPureEff, (:&), (:>))
import Bluefin.Internal (StateSource (StateSource))
import Bluefin.State (State, evalState, get, modify, put)
import Bluefin.StateSource (newState, withStateSource)
import Data.Functor ((<&>))
import Data.List (find)

instance Handle StateSource where
  mapHandle StateSource = StateSource

type Varname = String

type Qname = String

type Level = Int

data Exp
  = Var Varname
  | App Exp Exp
  | Lam Varname Exp
  | Let Varname Exp Exp

data Tv e = Unbound String Level | Link (Typ e)

data Typ' (e :: Effects)
  = TVar (Int, State (Tv e) e)
  | QVar Qname
  | TArrow (Typ e) (Typ e)

type Typ e = (Int, Typ' e)

showTv :: (e :> es) => Tv e -> Eff es String
showTv (Unbound name level) = pure $ "Unbound " ++ name ++ " " ++ show level
showTv (Link ty) = do
  ty' <- showTyp ty
  pure $ "Link (" ++ ty' ++ ")"

showTyp' :: (e :> es) => Typ' e -> Eff es String
showTyp' (TVar (_, s)) = do
  tv <- showTv =<< get s
  pure $ "TVar (" ++ tv ++ ")"
showTyp' (QVar name) = pure $ "QVar " ++ name
showTyp' (TArrow typ1 typ2) = do
  typ1' <- showTyp typ1
  typ2' <- showTyp typ2
  pure $ "TArrow (" ++ typ1' ++ ") (" ++ typ2' ++ ")"

showTyp :: (e :> es) => Typ e -> Eff es String
showTyp (_, ty) = showTyp' ty

data LevelOps e = MkLevelOps
  { currentLevelImpl :: Eff e Level,
    enterLevelImpl :: Eff e (),
    leaveLevelImpl :: Eff e ()
  }

runLevelOps :: (forall e. LevelOps e -> Eff (e :& es) r) -> Eff es r
runLevelOps k =
  evalState 0 $ \s ->
    k
      MkLevelOps
        { currentLevelImpl = get s,
          enterLevelImpl = modify s (+ 1),
          leaveLevelImpl = modify s (subtract 1)
        }

instance Handle LevelOps where
  mapHandle MkLevelOps {currentLevelImpl, enterLevelImpl, leaveLevelImpl} =
    MkLevelOps
      { currentLevelImpl = useImpl currentLevelImpl,
        enterLevelImpl = useImpl enterLevelImpl,
        leaveLevelImpl = useImpl leaveLevelImpl
      }

currentLevel :: (e :> es) => LevelOps e -> Eff es Level
currentLevel e = useImpl (currentLevelImpl e)

enterLevel, leaveLevel :: (e :> es) => LevelOps e -> Eff es ()
enterLevel e = useImpl (enterLevelImpl e)
leaveLevel e = useImpl (leaveLevelImpl e)

data Effs e = MkEffs (LevelOps e) (State Int e) (StateSource e)

runEffs :: (forall e. Effs e -> Eff (e :& es) r) -> Eff es r
runEffs k =
  runLevelOps $ \ops ->
    evalState 0 $ \cnt ->
      withStateSource $ \source ->
        useImplIn k (MkEffs (mapHandle ops) (mapHandle cnt) (mapHandle source))

unify :: (e :> es) => Typ e -> Typ e -> Eff es ()
unify (tag1, _) (tag2, _) | tag1 == tag2 = pure ()
unify (_, TVar tv@(_, ref)) t2 =
  get ref >>= \case
    Link t1 -> unify t1 t2
    Unbound {} -> occurs tv t2 >> put ref (Link t2)
unify t1 (_, TVar tv@(_, ref)) =
  get ref >>= \case
    Link t2 -> unify t1 t2
    Unbound {} -> occurs tv t1 >> put ref (Link t1)
unify (_, TArrow tyl1 tyl2) (_, TArrow tyr1 tyr2) = do
  unify tyl1 tyr1
  unify tyl2 tyr2
unify _ _ = error "Cannot unify types"

fresh :: (e :> es) => State Int e -> Eff es Int
fresh c = get c <* modify c (+ 1)

occurs :: (e :> es) => (Int, State (Tv e) e) -> Typ e -> Eff es ()
occurs tvr@(tag1, _) = \case
  (_, TVar (tag2, _)) | tag1 == tag2 -> error "occurs check"
  (_, TVar (_, ref)) ->
    get ref >>= \case
      Unbound name level -> do
        minLevel <-
          get (snd tvr) <&> \case
            Unbound _ level' -> min level level'
            _ -> level
        put ref $ Unbound name minLevel
      Link ty -> occurs tvr ty
  (_, TArrow t1 t2) -> occurs tvr t1 >> occurs tvr t2
  _ -> pure ()

gen :: (e :> es) => Effs e -> Typ e -> Eff es (Typ e)
gen effs@(MkEffs lvl cnt _) ty@(_, TVar (_, ref)) =
  get ref >>= \case
    Unbound name level -> do
      currLevel <- currentLevel lvl
      if level > currLevel
        then fresh cnt <&> (,QVar name)
        else pure ty
    Link ty' -> gen effs ty'
gen effs (tag, TArrow ty1 ty2) = fmap (tag,) . TArrow <$> gen effs ty1 <*> gen effs ty2
gen _ ty = pure ty

newvar :: (e :> es) => Effs e -> Eff es (Typ e)
newvar (MkEffs lvl cnt source) = do
  currLevel <- currentLevel lvl
  tag <- fresh cnt
  s <- newState source $ Unbound ("t" ++ show tag) currLevel
  fresh cnt <&> (,TVar (tag, s))

inst :: (e :> es) => Effs e -> Typ e -> Eff es (Typ e)
inst effs ty0 = fst <$> go [] ty0
  where
    go subst (_, QVar name) =
      case find ((== name) . fst) subst of
        Just (_, ty) -> pure (ty, subst)
        Nothing -> do
          ty' <- newvar effs
          pure (ty', (name, ty') : subst)
    go subst (tag, TArrow ty1 ty2) = do
      (ty1', subst') <- go subst ty1
      (ty2', subst'') <- go subst' ty2
      pure ((tag, TArrow ty1' ty2'), subst'')
    go subst ty@(_, TVar (_, s)) =
      get s >>= \case
        Link ty' -> go subst ty'
        Unbound {} -> pure (ty, subst)

typeof :: (e :> es) => Effs e -> [(String, Typ e)] -> Exp -> Eff es (Typ e)
typeof effs@(MkEffs lvl cnt _) env e0 =
  case e0 of
    Var x ->
      inst effs $ case find ((== x) . fst) env of
        Just (_, typ) -> typ
        Nothing -> error $ "Could not look up variable with name: " ++ x
    App e1 e2 -> do
      tyFun <- typeof effs env e1
      tyArg <- typeof effs env e2
      tyRes <- newvar effs
      tag <- fresh cnt
      unify tyFun (tag, TArrow tyArg tyRes)
      pure tyRes
    Lam x e -> do
      tyX <- newvar effs
      tyE <- typeof effs ((x, tyX) : env) e
      fresh cnt <&> (,TArrow tyX tyE)
    Let x e rest -> do
      tyE <- enterLevel lvl *> typeof effs env e <* leaveLevel lvl
      tyE' <- gen effs tyE
      typeof effs ((x, tyE') : env) rest

idExp, c1Exp, c2Exp, c3Exp :: Exp
idExp = Lam "x" (Var "x")
c1Exp = Lam "x" (Lam "y" (App (Var "x") (Var "y")))
c2Exp = Let "y" (Lam "z" (Var "z")) (Var "y")
c3Exp = Lam "x" (Let "y" (Lam "z" (Var "z")) (App (Var "y") (Var "x")))

idTy, c1Ty, c2Ty, c3Ty :: String
idTy = runPureEff $ runEffs $ \effs -> showTyp =<< typeof effs [] idExp
c1Ty = runPureEff $ runEffs $ \effs -> showTyp =<< typeof effs [] c1Exp
c2Ty = runPureEff $ runEffs $ \effs -> showTyp =<< typeof effs [] c2Exp
c3Ty = runPureEff $ runEffs $ \effs -> showTyp =<< typeof effs [] c3Exp

main :: IO ()
main = mapM_ print [idTy, c1Ty, c2Ty, c3Ty]