Start working on a type checkertypes

7 files changed, 135 insertions, 57 deletions

diff --git a/src/Haskell/AST.hs b/src/Haskell/AST.hs
index f8b5c72..54d19d3 100644
--- a/src/Haskell/AST.hs
+++ b/src/Haskell/AST.hs
@@ -20,12 +20,12 @@ data Toplevel = TopDef Def
 data Def = Def Name Expr
   deriving (Show, Eq)
 
-data Expr = App Expr [Expr]
-          | Ref Name
-          | Num Integer
-          | Tup [Expr]
-          | Lam [Name] Expr
-          | Case Expr [(Pat, Expr)]
+data Expr = App Expr [Expr] (Maybe Type)
+          | Ref Name (Maybe Type)
+          | Num Integer (Maybe Type)
+          | Tup [Expr] (Maybe Type)
+          | Lam [Name] Expr (Maybe Type)
+          | Case Expr [(Pat, Expr)] (Maybe Type)
   deriving (Show, Eq)
 
 data Pat = PatAny
@@ -69,12 +69,12 @@ instance Pretty Def where
     pretty (Def n e) = Node (n ++ " =") [pretty e]
 
 instance Pretty Expr where
-    pretty (App e as) = Bracket "(" ")" "" (map pretty (e:as))
-    pretty (Ref n) = Leaf n
-    pretty (Num n) = Leaf (show n)
-    pretty (Tup es) = Bracket "(" ")" "," (map pretty es)
-    pretty (Lam as e) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
-    pretty (Case e arms) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
+    pretty (App e as _) = Bracket "(" ")" "" (map pretty (e:as))
+    pretty (Ref n _) = Leaf n
+    pretty (Num n _) = Leaf (show n)
+    pretty (Tup es _) = Bracket "(" ")" "," (map pretty es)
+    pretty (Lam as e _) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
+    pretty (Case e arms _) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
       where go (p, e') = Node (pprintOneline p ++ " ->") [pretty e']
 
 instance Pretty Pat where
@@ -122,12 +122,36 @@ instance AllRefs Def where
     allRefs (Def _ e) = allRefs e
 
 instance AllRefs Expr where
-    allRefs (App e es) = nub $ concatMap allRefs (e : es)
-    allRefs (Ref n) = [n]
-    allRefs (Num _) = []
-    allRefs (Tup es) = nub $ concatMap allRefs es
-    allRefs (Lam ns e) = allRefs e \\ ns
-    allRefs (Case e pairs) = nub $ allRefs e ++ concatMap (allRefs . snd) pairs
+    allRefs (App e es _) = nub $ concatMap allRefs (e : es)
+    allRefs (Ref n _) = [n]
+    allRefs (Num _ _) = []
+    allRefs (Tup es _) = nub $ concatMap allRefs es
+    allRefs (Lam ns e _) = allRefs e \\ ns
+    allRefs (Case e pairs _) = nub $ allRefs e ++ concatMap (allRefs . snd) pairs
 
 instance AllRefs Inst where
     allRefs (Inst _ _ ds) = nub $ concatMap allRefs ds
+
+
+typeOf :: Expr -> Type
+typeOf (App _ _ (Just ty)) = ty
+typeOf (Ref _ (Just ty)) = ty
+typeOf (Num _ (Just ty)) = ty
+typeOf (Tup _ (Just ty)) = ty
+typeOf (Lam _ _ (Just ty)) = ty
+typeOf (Case _ _ (Just ty)) = ty
+typeOf e = error $ "Unresolved type in typeOf in expression: " ++ show e
+
+typeApply :: Type -> Type -> Maybe Type
+typeApply (TyFun t1 t2) t3
+    | t1 == t3  = Just t2
+    | otherwise = Nothing
+typeApply _ _ = Nothing
+
+recurse :: (Pat -> Pat) -> (Expr -> Expr) -> Expr -> Expr
+recurse _  f (App e as ty) = App (f e) (map f as) ty
+recurse _  _ (Ref n ty) = Ref n ty
+recurse _  _ (Num k ty) = Num k ty
+recurse _  f (Tup es ty) = Tup (map f es) ty
+recurse _  f (Lam ns e ty) = Lam ns (f e) ty
+recurse fp f (Case e as ty) = Case (f e) (map (\(p, e') -> (fp p, f e')) as) ty
diff --git a/src/Haskell/Env.hs b/src/Haskell/Env.hs
index aca2367..636f298 100644
--- a/src/Haskell/Env.hs
+++ b/src/Haskell/Env.hs
@@ -11,7 +11,7 @@ data Env = Env { eDefs :: Map.Map Name Expr }
 
 instance Pretty Env where
     pretty (Env defs) =
-        Node "" [Node (n ++ " =") [pretty e] | (n, e) <- Map.assocs defs]
+        Node "" [Node (n ++ " =") [pretty e, pretty (typeOf e)] | (n, e) <- Map.assocs defs]
 
 emptyEnv :: Env
 emptyEnv = Env Map.empty
diff --git a/src/Haskell/Parser.hs b/src/Haskell/Parser.hs
index 41216d5..0536fde 100644
--- a/src/Haskell/Parser.hs
+++ b/src/Haskell/Parser.hs
@@ -29,35 +29,35 @@ pDef = do
     symbolO ";"
     case args of
         [] -> return $ Def n ex
-        _  -> return $ Def n (Lam args ex)
+        _  -> return $ Def n (Lam args ex Nothing)
 
 pExpr :: Parser Expr
 pExpr = pLam <|> pCase <|> pApp
   where
-    pSimpleExpr = choice [Num <$> pNum
-                         ,Ref <$> pVariable
+    pSimpleExpr = choice [Num <$> pNum <*> return Nothing
+                         ,Ref <$> pVariable <*> return Nothing
                          ,parens (pExpr `sepBy` symbolO ",") >>= \case
                             [ex] -> return ex
-                            exs  -> return $ Tup exs]
+                            exs  -> return $ Tup exs Nothing]
 
     pLam = do
         symbolO "\\"
         args <- many1 pNameV
         symbolO "->"
         body <- pExpr
-        return $ Lam args body
+        return $ Lam args body Nothing
 
     pApp = many1 pSimpleExpr >>= \case
             []     -> undefined
             [e]    -> return e
-            (e:es) -> return $ App e es
+            (e:es) -> return $ App e es Nothing
 
     pCase = do
         symbolW "case"
         e <- pExpr
         symbolW "of"
         arms <- braces (pCaseArm `sepBy` symbolO ";")
-        return $ Case e arms
+        return $ Case e arms Nothing
 
     pCaseArm = do
         pat <- pLargePat
diff --git a/src/Haskell/Rewrite.hs b/src/Haskell/Rewrite.hs
index bcec6f7..9812231 100644
--- a/src/Haskell/Rewrite.hs
+++ b/src/Haskell/Rewrite.hs
@@ -15,12 +15,12 @@ import Util
 
 rewrite :: Name -> Expr -> Expr -> Expr
 rewrite name repl = \case
-    App e as -> App (rewrite name repl e) (map (rewrite name repl) as)
-    Ref n -> if n == name then repl else Ref n
-    Num k -> Num k
-    Tup es -> Tup (map (rewrite name repl) es)
-    Lam ns e -> if name `elem` ns then Lam ns e else Lam ns (rewrite name repl e)
-    Case e as -> Case (rewrite name repl e) (map caseArm as)
+    App e as ty -> App (rewrite name repl e) (map (rewrite name repl) as) ty
+    Ref n ty -> if n == name then repl else Ref n ty
+    Num k ty -> Num k ty
+    Tup es ty -> Tup (map (rewrite name repl) es) ty
+    Lam ns e ty -> if name `elem` ns then Lam ns e ty else Lam ns (rewrite name repl e) ty
+    Case e as ty -> Case (rewrite name repl e) (map caseArm as) ty
   where
     caseArm (p, e') =
         if name `elem` boundVars p then (p, e') else (p, rewrite name repl e')
@@ -32,18 +32,21 @@ boundVars (PatCon n ps) = nub $ [n] ++ concatMap boundVars ps
 boundVars (PatTup ps) = nub $ concatMap boundVars ps
 
 betared :: Expr -> Expr
-betared (App (Lam (n:as) bd) (arg:args)) =
-    App (Lam as (rewrite n arg bd)) args
+betared (App (Lam (n:as) bd lty) (arg:args) ty) =
+    App (Lam as
+             (rewrite n arg bd)
+             (Just $ fromJust $ typeApply (fromJust lty) (typeOf arg)))
+        args ty
 betared e = recurse id betared e
 
 etared :: Expr -> Expr
-etared (Lam (n:as) (App e es@(_:_)))
-    | last es == Ref n =
-        Lam as (App e (init es))
+etared (Lam (n:as) (App e es@(_:_) aty) ty)
+    | Ref n' _ <- last es, n == n' =
+        Lam as (App e (init es) aty) ty
 etared e = recurse id etared e
 
 casered :: Bool -> Expr -> Expr
-casered ambig orig@(Case subj arms) =
+casered ambig orig@(Case subj arms _) =
     case catMaybes [(,rhs) <$> unify p subj | (p, rhs) <- arms] of
         [] -> recurse id (casered ambig) orig
         ((mp, rhs):rest) ->
@@ -55,12 +58,12 @@ casered ambig orig@(Case subj arms) =
 casered ambig e = recurse id (casered ambig) e
 
 etacase :: Expr -> Expr
-etacase (Case subj arms) | all (uncurry eqPE) arms = subj
+etacase (Case subj arms _) | all (uncurry eqPE) arms = subj
 etacase e = recurse id etacase e
 
 casecase :: Expr -> Expr
-casecase (Case (Case subj arms1) arms2) =
-    Case subj [(p, Case e arms2) | (p, e) <- arms1]
+casecase (Case (Case subj arms1 _) arms2 ty) =
+    Case subj [(p, Case e arms2 ty) | (p, e) <- arms1] ty
 casecase e = recurse id casecase e
 
 autoSimp :: Expr -> Expr
@@ -73,15 +76,16 @@ eqPE pat expr = case unify pat expr of
     Nothing -> False
     Just mp -> all (uncurry isIdmap) (Map.assocs mp)
   where
-    isIdmap n = (== Ref n)
+    isIdmap n (Ref n' _) = n == n'
+    isIdmap _ _ = False
 
 unify :: Pat -> Expr -> Maybe (Map.Map Name Expr)
 unify PatAny _ = Just Map.empty
 unify (PatVar n) e = Just (Map.singleton n e)
-unify (PatCon n ps) (App n' es)
-    | n' == Ref n, length ps == length es =
+unify (PatCon n ps) (App (Ref n' _) es _)
+    | n == n', length ps == length es =
         foldM (\m (p, e) -> unify p e >>= reconcile m) Map.empty (zip ps es)
-unify (PatTup ps) (Tup es)
+unify (PatTup ps) (Tup es _)
     | length ps == length es =
         foldM (\m (p, e) -> unify p e >>= reconcile m) Map.empty (zip ps es)
 unify _ _ = Nothing
@@ -94,14 +98,6 @@ reconcile m1 m2 = foldM func m1 (Map.assocs m2)
                     | otherwise -> Nothing
 
 normalise :: Expr -> Expr
-normalise (App e []) = normalise e
-normalise (Lam [] e) = normalise e
+normalise (App e [] _) = normalise e
+normalise (Lam [] e _) = normalise e
 normalise e = recurse id normalise e
-
-recurse :: (Pat -> Pat) -> (Expr -> Expr) -> Expr -> Expr
-recurse _  f (App e as) = App (f e) (map f as)
-recurse _  _ (Ref n) = Ref n
-recurse _  _ (Num k) = Num k
-recurse _  f (Tup es) = Tup (map f es)
-recurse _  f (Lam ns e) = Lam ns (f e)
-recurse fp f (Case e as) = Case (f e) (map (\(p, e') -> (fp p, f e')) as)
diff --git a/src/Haskell/TypeCheck.hs b/src/Haskell/TypeCheck.hs
new file mode 100644
index 0000000..52e26b7
--- /dev/null
+++ b/src/Haskell/TypeCheck.hs
@@ -0,0 +1,55 @@
+module Haskell.TypeCheck (typeCheck) where
+
+import qualified Data.Map.Strict as Map
+import Haskell.AST
+
+
+typeCheck :: (Eq a, TypeCheck a) => a -> Either String a
+typeCheck initVal = go initVal Map.empty
+  where
+    go value mapping = do
+        value2 <- annotate mapping value
+        if value2 == value then return value else go value2 mapping
+
+class TypeCheck a where
+    annotate :: Map.Map String Type -> a -> Either String a
+
+instance TypeCheck AST where
+    annotate mapping (AST tops) = AST <$> mapM (annotate mapping) tops
+
+instance TypeCheck Toplevel where
+    annotate mapping (TopDef x)   = TopDef <$> annotate mapping x
+    annotate _       (TopDecl x)  = return $ TopDecl x
+    annotate _       (TopData x)  = return $ TopData x
+    annotate _       (TopClass x) = return $ TopClass x
+    annotate mapping (TopInst x)  = TopInst <$> annotate mapping x
+
+instance TypeCheck Def where
+    annotate mapping (Def name expr) = Def name <$> annotate mapping expr
+
+instance TypeCheck Inst where
+    annotate mapping (Inst name ty defs) = Inst name ty <$> mapM (annotate mapping) defs
+
+-- TODO: Just passing mappings around is not sufficient; the typechecking
+-- needs to be done on the whole environment at once, and dynamically
+-- updating the mapping as it goes. In fact, I'd probably do better just
+-- looking up how the simplest possible implementation of Hindley-Miller
+-- goes, stripping away most classes, and then implementing that.
+instance TypeCheck Expr where
+    annotate mapping (App func' [] _) = annotate mapping func'
+    annotate mapping (App func' [arg'] _) = do
+        func <- annotate mapping func'
+        arg <- annotate mapping arg'
+        case typeApply (typeOf func) (typeOf arg) of
+            Just ty -> return $ App func [arg] (Just ty)
+            Nothing -> Left $ "Cannot apply argument of type " ++ show (typeOf arg) ++
+                              " to function of type " ++ show (typeOf func)
+    annotate mapping (App func' (arg':args') _) = 
+        combineApps <$> annotate mapping (App (App func' [arg'] Nothing) args' Nothing)
+
+    -- annotate mapping (Ref )
+
+
+combineApps :: Expr -> Expr
+combineApps (App (App f as1 _) as2 ty) = combineApps (App f (as1 ++ as2) ty)
+combineApps e = recurse id combineApps e
diff --git a/src/Main.hs b/src/Main.hs
index a6d80b3..1b1f17a 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -7,8 +7,9 @@ import Haskell.AST
 import Haskell.Env
 import Haskell.Env.Cmd
 import Haskell.Env.Context
-import Haskell.Rewrite
 import Haskell.Parser
+import Haskell.Rewrite
+import Haskell.TypeCheck
 import Pretty
 import System.Environment
 import System.Exit
@@ -221,7 +222,8 @@ main = do
         _       -> die "Usage: verify-hs <file.hs>"
 
     source <- readFile fname
-    ast <- tryEither' (parseAST fname source)
+    ast' <- tryEither' (parseAST fname source)
+    ast <- tryEither' (typeCheck ast')
     -- print ast
     env0 <- tryEither (envFromAST ast)
     -- print env0
diff --git a/verify-hs.cabal b/verify-hs.cabal
index 68dec3d..08c69fb 100644
--- a/verify-hs.cabal
+++ b/verify-hs.cabal
@@ -18,6 +18,7 @@ executable verify-hs
     Haskell.Env.Context
     Haskell.Parser
     Haskell.Rewrite
+    Haskell.TypeCheck
     Pretty
     Util
   build-depends: