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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DeriveAnyClass #-}
module Minilang.Lambda.Infer where
import Control.Monad.State (MonadState (..), StateT, evalStateT, modify)
import Data.Aeson (FromJSON, ToJSON)
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 GHC.Generics (Generic)
import qualified Minilang.Lambda.Eval as Eval
import Minilang.Lambda.Unify (Type (..), UnifyError, apply, unify)
import Prelude hiding (lookup)
import qualified Minilang.Lambda.Unify as Unify
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)
render :: TypeError -> Text
render (UnifyError ue) = Unify.render ue
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, Generic, ToJSON, FromJSON)
extendContext :: Text -> Type -> Context -> Context
extendContext x ty (Context n env) =
Context n (Map.insert x ty env)
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
|
