add minimal evaluator and type inference

1 files changed, 97 insertions, 0 deletions

diff --git a/lambda-calcul/haskell/src/Minilang/Lambda/Infer.hs b/lambda-calcul/haskell/src/Minilang/Lambda/Infer.hs
new file mode 100644
index 0000000..415849a
--- /dev/null
+++ b/lambda-calcul/haskell/src/Minilang/Lambda/Infer.hs
@@ -0,0 +1,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 Env = Env {nextVar :: Int, typeEnv :: Map.Map Text Type}
+  deriving (Eq, Show)
+
+annotate :: Eval.Term -> Either TypeError (ATerm Type)
+annotate t =
+  evalStateT (go (fromTerm t) []) newEnv
+  where
+    go :: ATerm a -> [(Text, Type)] -> StateT Env (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
+
+newEnv :: Env
+newEnv = Env 0 Map.empty
+
+bind :: Text -> Type -> StateT Env (Either TypeError) ()
+bind x ty = do
+  modify (\(Env n env) -> Env n (Map.insert x ty env))
+
+lookup :: Text -> StateT Env (Either TypeError) (Maybe Type)
+lookup x = do
+  Env _ env <- get
+  pure $ Map.lookup x env
+
+freshVar :: (Monad m) => StateT Env m Type
+freshVar = do
+  get >>= \(Env n _) ->
+    modify (\(Env n' env') -> Env (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