1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
{-# LANGUAGE LambdaCase #-}
module Minilang.Lambda.Infer where
import Control.Monad.State (MonadState (..), StateT, evalStateT, modify)
import Data.Bifunctor (first)
import Data.Char (chr, ord)
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Text (Text, pack)
import qualified Minilang.Lambda.Eval as Eval
import Minilang.Lambda.Unify (Type (..), UnifyError, apply, unify)
import Prelude hiding (lookup)
data ATerm ann
= Var Text ann
| Lam Text (ATerm ann) ann
| App (ATerm ann) (ATerm ann) ann
deriving (Show, Eq)
data TypeError
= UnifyError UnifyError
deriving (Show, Eq)
fromTerm :: Eval.Term -> ATerm ()
fromTerm (Eval.Var x) = Var x ()
fromTerm (Eval.Lam x body) = Lam x (fromTerm body) ()
fromTerm (Eval.App a b) = App (fromTerm a) (fromTerm b) ()
infer :: Eval.Term -> Either TypeError Type
infer t = do
annotated <- annotate t
constraints <- collect [annotated] []
subs <- first UnifyError $ unify constraints
pure $ apply subs (typeOf annotated)
collect :: [ATerm Type] -> [(Type, Type)] -> Either TypeError [(Type, Type)]
collect terms constraints = case terms of
[] -> Right constraints
(Var _ _) : rest -> collect rest constraints
(Lam _ body _) : rest -> collect (body : rest) constraints
(App x y t) : rest ->
let (f, b) = (typeOf x, typeOf y)
in collect (x : y : rest) ((f, b :-> t) : constraints)
typeOf :: ATerm Type -> Type
typeOf (Var _ t) = t
typeOf (Lam _ _ t) = t
typeOf (App _ _ t) = t
data Context = Context {nextVar :: Int, typeContext :: Map.Map Text Type}
deriving (Eq, Show)
annotate :: Eval.Term -> Either TypeError (ATerm Type)
annotate t =
evalStateT (go (fromTerm t) []) newContext
where
go :: ATerm a -> [(Text, Type)] -> StateT Context (Either TypeError) (ATerm Type)
go t bounds = case t of
Var x _ -> case List.lookup x bounds of
Just ty -> pure $ Var x ty
Nothing ->
lookup x >>= \case
Just ty -> pure $ Var x ty
Nothing -> do
v <- freshVar
bind x v
pure $ Var x v
Lam x body _ -> do
ty <- freshVar
body' <- go body ((x, ty) : bounds)
pure $ Lam x body' (ty :-> typeOf body')
App a b _ ->
App <$> go a bounds <*> go b bounds <*> freshVar
newContext :: Context
newContext = Context 0 Map.empty
bind :: Text -> Type -> StateT Context (Either TypeError) ()
bind x ty = do
modify (\(Context n env) -> Context n (Map.insert x ty env))
lookup :: Text -> StateT Context (Either TypeError) (Maybe Type)
lookup x = do
Context _ env <- get
pure $ Map.lookup x env
freshVar :: (Monad m) => StateT Context m Type
freshVar = do
get >>= \(Context n _) ->
modify (\(Context n' env') -> Context (n' + 1) env') >> pure (TyVar $ nameOf n)
where
nameOf n
| n < 26 = pack [chr (n + ord 'a')]
| otherwise =
let (d, r) = n `divMod` 26
in pack $ chr (r + ord 'a') : show d
|
