typechecker work

2 files changed, 179 insertions, 56 deletions

diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index 2b125b9..e25657b 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -95,6 +95,7 @@ data Type s
   | TFun (X Type s) (Type s) (Type s)
   | TCon (X Type s) Name
   | TVar (X Type s) Name
+  | TForall (X Type s) Name (Type s)  -- ^ implicit
 
   -- extension point
   | TExt (X Type s) !(E Type s)
@@ -156,6 +157,8 @@ instance Pretty (E Type s) => Pretty (Type s) where
     prettysPrec 0 a . showString " -> " . prettysPrec (-1) b
   prettysPrec _ (TCon _ n) = prettysPrec 11 n
   prettysPrec _ (TVar _ n) = prettysPrec 11 n
+  prettysPrec d (TForall _ n t) = showParen (d > -1) $
+    showString "forall " . prettysPrec 11 n . showString "." . prettysPrec (-1) t
   prettysPrec d (TExt _ e) = prettysPrec d e
 
 instance (Pretty (X Type s), Pretty (E Type s)) => Pretty (DataDef s) where
@@ -168,6 +171,17 @@ instance (Pretty (X Type s), Pretty (E Type s)) => Pretty (DataDef s) where
                       [prettysPrec 11 cname . foldr (.) id [showString " " . prettysPrec 11 t | t <- fields]
                       | (cname, fields) <- cons])
 
+instance Pretty Operator where
+  prettysPrec _ = showString . \case
+    OAdd -> "(+)"
+    OSub -> "(-)"
+    OMul -> "(*)"
+    ODiv -> "(/)"
+    OMod -> "(%)"
+    OEqu -> "(==)"
+    OPow -> "(^)"
+    OCons -> "(:)"
+
 instance HasExt DataDef where
   type HasXField DataDef = 'True
   type HasECon DataDef = 'False
@@ -225,6 +239,7 @@ instance HasExt Type where
   extOf (TFun x _ _) = x
   extOf (TCon x _) = x
   extOf (TVar x _) = x
+  extOf (TForall x _ _) = x
   extOf (TExt x _) = x
 
   extMap p ps1 ps2 f g (TApp x a b) = TApp (f x) (extMap p ps1 ps2 f g a) (map (extMap p ps1 ps2 f g) b)
@@ -233,6 +248,7 @@ instance HasExt Type where
   extMap p ps1 ps2 f g (TFun x a b) = TFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
   extMap _ _   _   f _ (TCon x n) = TCon (f x) n
   extMap _ _   _   f _ (TVar x n) = TVar (f x) n
+  extMap p ps1 ps2 f g (TForall x n a) = TForall (f x) n (extMap p ps1 ps2 f g a)
   extMap _ _   _   _ g (TExt x e) = g x e
 
 instance HasExt Pattern where
diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
index 0347e81..f62b097 100644
--- a/src/HSVIS/Typecheck.hs
+++ b/src/HSVIS/Typecheck.hs
@@ -16,11 +16,12 @@ module HSVIS.Typecheck (
 import Control.Monad
 import Data.Bifunctor (first, second)
 import Data.Foldable (toList)
-import Data.List (find)
+import Data.List (find, inits)
 import Data.Map.Strict (Map)
 import Data.Maybe (fromMaybe)
 import Data.Monoid (Ap(..))
 import qualified Data.Map.Strict as Map
+import Data.Tuple (swap)
 import Data.Set (Set)
 import qualified Data.Set as Set
 
@@ -37,9 +38,9 @@ import HSVIS.Typecheck.Solve
 
 data StageTC
 
-type instance X DataDef StageTC = ()
+type instance X DataDef StageTC = CKind
 type instance X FunDef  StageTC = CType
-type instance X FunEq   StageTC = CType
+type instance X FunEq   StageTC = ()
 type instance X Kind    StageTC = ()
 type instance X Type    StageTC = CKind
 type instance X Pattern StageTC = CType
@@ -64,7 +65,7 @@ data StageTyped
 
 type instance X DataDef StageTyped = TType
 type instance X FunDef  StageTyped = TType
-type instance X FunEq   StageTyped = TType
+type instance X FunEq   StageTyped = ()
 type instance X Kind    StageTyped = ()
 type instance X Type    StageTyped = TKind
 type instance X Pattern StageTyped = TType
@@ -95,7 +96,7 @@ instance Pretty (E Kind StageTC) where
 typecheck :: FilePath -> String -> PProgram -> ([Diagnostic], TProgram)
 typecheck fp source prog =
   let (ds1, cs, _, _, progtc) =
-        runTCM (tcProgram prog) (fp, source) 1 (Env mempty mempty)
+        runTCM (tcProgram prog) (fp, source) 1 (Env mempty mempty mempty)
       (ds2, subK, subT) = solveConstrs cs
   in (toList (ds1 <> ds2), doneProg subK subT progtc)
 
@@ -107,7 +108,14 @@ data Constr
   | CEqK CKind CKind Range
   deriving (Show)
 
-data Env = Env (Map Name CKind) (Map Name CType)
+data Env = Env (Map Name CKind)   -- ^ types in scope (including variables)
+               (Map Name CType)   -- ^ values in scope (constructors and variables)
+               (Map Name InvCon)  -- ^ patterns in scope (inverse constructors)
+  deriving (Show)
+
+data InvCon = InvCon (Map Name CKind)  -- ^ universally quantified type variables
+                     CType    -- ^ input type of the inverse constructor (result of the constructor)
+                     [CType]  -- ^ field types
   deriving (Show)
 
 newtype TCM a = TCM {
@@ -158,34 +166,44 @@ genId = TCM $ \_ i env -> (mempty, mempty, i + 1, env, i)
 
 getKind :: Name -> TCM (Maybe CKind)
 getKind name = do
-  Env tenv _ <- getFullEnv
+  Env tenv _ _ <- getFullEnv
   return (Map.lookup name tenv)
 
 getType :: Name -> TCM (Maybe CType)
 getType name = do
-  Env _ venv <- getFullEnv
+  Env _ venv _ <- getFullEnv
   return (Map.lookup name venv)
 
+getInvCon :: Name -> TCM (Maybe InvCon)
+getInvCon name = do
+  Env _ _ icenv <- getFullEnv
+  return (Map.lookup name icenv)
+
 modifyTEnv :: (Map Name CKind -> Map Name CKind) -> TCM ()
 modifyTEnv f = do
-  Env tenv venv <- getFullEnv
-  putFullEnv (Env (f tenv) venv)
+  Env tenv venv icenv <- getFullEnv
+  putFullEnv (Env (f tenv) venv icenv)
 
 modifyVEnv :: (Map Name CType -> Map Name CType) -> TCM ()
 modifyVEnv f = do
-  Env tenv venv <- getFullEnv
-  putFullEnv (Env tenv (f venv))
+  Env tenv venv icenv <- getFullEnv
+  putFullEnv (Env tenv (f venv) icenv)
+
+modifyICEnv :: (Map Name InvCon -> Map Name InvCon) -> TCM ()
+modifyICEnv f = do
+  Env tenv venv icenv <- getFullEnv
+  putFullEnv (Env tenv venv (f icenv))
 
 scopeTEnv :: TCM a -> TCM a
 scopeTEnv m = do
-  Env origtenv _ <- getFullEnv
+  Env origtenv _ _ <- getFullEnv
   res <- m
   modifyTEnv (\_ -> origtenv)
   return res
 
 scopeVEnv :: TCM a -> TCM a
 scopeVEnv m = do
-  Env _ origvenv <- getFullEnv
+  Env _ origvenv _ <- getFullEnv
   res <- m
   modifyVEnv (\_ -> origvenv)
   return res
@@ -213,11 +231,15 @@ getType' rng name = getType name >>= \case
 tcProgram :: PProgram -> TCM CProgram
 tcProgram (Program ddefs1 fdefs1) = do
   (kconstrs, ddefs2) <- collectConstraints isCEqK $ do
-    mapM_ prepareDataDef ddefs1
-    mapM tcDataDef ddefs1
+    ks <- mapM prepareDataDef ddefs1
+    zipWithM kcDataDef ks ddefs1
 
   kinduvars <- solveKindVars kconstrs
   let ddefs3 = map (substDdef kinduvars mempty) ddefs2
+  modifyTEnv (Map.map (substKind kinduvars))
+
+  forM_ ddefs3 $ \ddef ->
+    modifyICEnv (Map.fromList (generateInvCons ddef) <>)
 
   traceM (unlines (map pretty ddefs3))
 
@@ -225,32 +247,46 @@ tcProgram (Program ddefs1 fdefs1) = do
 
   return (Program ddefs3 fdefs2)
 
-prepareDataDef :: PDataDef -> TCM ()
+-- Bring data type name in scope with a kind of the specified arity
+prepareDataDef :: PDataDef -> TCM (CKind, [CKind])
 prepareDataDef (DataDef _ name params _) = do
   parkinds <- mapM (\_ -> genKUniVar) params
   let k = foldr (KFun ()) (KType ()) parkinds
   modifyTEnv (Map.insert name k)
+  return (k, parkinds)
 
 -- Assumes that the kind of the name itself has already been registered with
 -- the correct arity (this is done by prepareDataDef).
-tcDataDef :: PDataDef -> TCM CDataDef
-tcDataDef (DataDef rng name params cons) = do
-  kd <- getKind' rng name
-  let (pkinds, kret) = splitKind kd
+kcDataDef :: (CKind, [CKind]) -> PDataDef -> TCM CDataDef
+kcDataDef (kd, parkinds) (DataDef _ name params cons) = do
+  -- ensure unicity of type param names
+  params' <-
+    let prenames = Set.fromList (map snd params)
+        namegen = filter (`Set.notMember` prenames) [Name ('t' : show i) | i <- [1::Int ..]]
+    in forM (zip3 params (inits (map snd params)) namegen) $ \((rng, pname), previous, replname) ->
+         if pname `elem` previous
+           then do raise SError rng $ "Duplicate type parameter: " ++ pretty pname
+                   return replname
+           else return pname
+
+  -- kind-check the constructors
+  cons' <- scopeTEnv $ do
+    modifyTEnv (Map.fromList (zip params' parkinds) <>)
+    forM cons $ \(cname, fieldtys) -> do
+      fieldtys' <- mapM (kcType (Just (KType ()))) fieldtys
+      return (cname, fieldtys')
 
-  -- sanity checking; would be nicer to store these in prepareDataDef already
-  when (length pkinds /= length params) $ error "tcDataDef: Invalid param kind list length"
-  case kret of Right () -> return ()
-               _ -> error "tcDataDef: Invalid ret kind"
+  return (DataDef kd name (zip parkinds params') cons')
 
-  cons' <- scopeTEnv $ do
-    modifyTEnv (Map.fromList (zip (map snd params) pkinds) <>)
-    mapM (\(cname, fieldtys) -> (cname,) <$> mapM (kcType (Just (KType ()))) fieldtys) cons
-  return (DataDef () name (zip pkinds (map snd params)) cons')
+generateInvCons :: CDataDef -> [(Name, InvCon)]
+generateInvCons (DataDef k tname params cons) =
+  let tyvars = Map.fromList (map swap params)
+      resty = TApp (KType ()) (TCon k tname) (map (uncurry TVar) params)
+  in [(cname, InvCon tyvars resty fields) | (cname, fields) <- cons]
 
-promoteDown :: Maybe CKind -> TCM CKind
-promoteDown Nothing = genKUniVar
-promoteDown (Just k) = return k
+promoteDownK :: Maybe CKind -> TCM CKind
+promoteDownK Nothing = genKUniVar
+promoteDownK (Just k) = return k
 
 downEqK :: Range -> Maybe CKind -> CKind -> TCM ()
 downEqK _ Nothing _ = return ()
@@ -263,7 +299,7 @@ kcType mdown = \case
   TApp rng t ts -> do
     t' <- kcType Nothing t
     ts' <- mapM (kcType Nothing) ts
-    retk <- promoteDown mdown
+    retk <- promoteDownK mdown
     let expected = foldr (KFun ()) retk (map extOf ts')
     emit $ CEqK (extOf t') expected rng
     return (TApp retk t' ts')
@@ -299,6 +335,15 @@ kcType mdown = \case
     downEqK rng mdown k
     return (TVar k n)
 
+  TForall rng n t -> do  -- implicit forall
+    k1 <- genKUniVar
+    k2 <- genKUniVar
+    downEqK rng mdown k2
+    t' <- scopeTEnv $ do
+      modifyTEnv (Map.insert n k1)
+      kcType (Just k2) t
+    return (TForall k2 n t')  -- not 'k1 -> k2' because the forall is implicit
+
 tcFunDef :: PFunDef -> TCM CFunDef
 tcFunDef (FunDef rng name msig eqs) = do
   when (not $ allEq (fmap (length . funeqPats) eqs)) $
@@ -313,7 +358,66 @@ tcFunDef (FunDef rng name msig eqs) = do
   return (FunDef typ name (TypeSig typ) eqs')
 
 tcFunEq :: CType -> PFunEq -> TCM CFunEq
-tcFunEq down (FunEq rng name pats rhs) = error "tcFunEq"
+tcFunEq down (FunEq rng name pats rhs) = do
+  -- getFullEnv >>= \env -> traceM $ "[tcFunEq] Env = " ++ show env
+  (argtys, rhsty) <- unfoldFunTy rng (length pats) down
+  scopeVEnv $ do
+    pats' <- zipWithM tcPattern argtys pats
+    rhs' <- tcRHS rhsty rhs
+    return (FunEq () name pats' rhs')
+
+tcPattern :: CType -> PPattern -> TCM CPattern
+tcPattern down = \case
+  PWildcard _ -> return $ PWildcard down
+  PVar _ n -> modifyVEnv (Map.insert n down) >> return (PVar down n)
+  PAs _ n p -> modifyVEnv (Map.insert n down) >> tcPattern down p
+  PCon rng n ps ->
+    getInvCon n >>= \case
+      Just (InvCon tyvars match fields) -> do
+        unisub <- mapM genUniVar tyvars  -- substitution for the universally quantified variables
+        let match' = substType mempty mempty unisub match
+            fields' = map (substType mempty mempty unisub) fields
+        emit $ CEq down match' rng
+        PCon match' n <$> zipWithM tcPattern fields' ps
+      Nothing -> do
+        raise SError rng $ "Constructor not in scope: " ++ pretty n
+        return (PWildcard down)
+  POp rng p1 op p2 ->
+    case op of
+      OCons -> do
+        eltty <- genUniVar (KType ())
+        let listty = TList (KType ()) eltty
+        emit $ CEq down listty rng
+        p1' <- tcPattern eltty p1
+        p2' <- tcPattern listty p2
+        return (POp listty p1' OCons p2')
+      _ -> do
+        raise SError rng $ "Operator is not a constructor: " ++ pretty op
+        return (PWildcard down)
+  PList rng ps -> do
+    eltty <- genUniVar (KType ())
+    let listty = TList (KType ()) eltty
+    emit $ CEq down listty rng
+    PList listty <$> mapM (tcPattern eltty) ps
+  PTup rng ps -> do
+    ts <- mapM (\_ -> genUniVar (KType ())) ps
+    emit $ CEq down (TTup (KType ()) ts) rng
+    PTup (TTup (KType ()) ts) <$> zipWithM tcPattern ts ps
+
+tcRHS :: CType -> PRHS -> TCM CRHS
+tcRHS = error "tcRHS"
+
+unfoldFunTy :: Range -> Int -> CType -> TCM ([CType], CType)
+unfoldFunTy _ n t | n <= 0 = return ([], t)
+unfoldFunTy rng n (TFun _ t1 t2) = do
+  (params, core) <- unfoldFunTy rng (n - 1) t2
+  return (t1 : params, core)
+unfoldFunTy rng n t = do
+  vars <- replicateM n (genUniVar (KType ()))
+  core <- genUniVar (KType ())
+  let expected = foldr (TFun (KType ())) core vars
+  emit $ CEq expected t rng
+  return (vars, core)
 
 solveKindVars :: Bag (CKind, CKind, Range) -> TCM (Map Int CKind)
 solveKindVars cs = do
@@ -366,36 +470,44 @@ solveKindVars cs = do
 solveConstrs :: Bag Constr -> (Bag Diagnostic, Map Int TKind, Map Int TType)
 solveConstrs = error "solveConstrs"
 
-substProg :: Map Int CKind  -- ^ Kind variable instantiations
-          -> Map Int CType  -- ^ Type variable instantiations
+-- substitute unification variables
+substProg :: Map Int CKind  -- ^ Kind unification variable instantiations
+          -> Map Int CType  -- ^ Type unification variable instantiations
           -> CProgram
           -> CProgram
 substProg = error "substProg"
 
+-- substitute unification variables
 substDdef :: Map Int CKind -> Map Int CType -> CDataDef -> CDataDef
-substDdef mk mt (DataDef () name pars cons) =
-  DataDef () name
+substDdef mk mt (DataDef k name pars cons) =
+  DataDef (substKind mk k) name
           (map (first (substKind mk)) pars)
-          (map (second (map (substType mk mt))) cons)
-
-substType :: Map Int CKind -> Map Int CType -> CType -> CType
-substType mk mt = \case
-  TApp k t ts -> TApp (substKind mk k) (substType mk mt t) (map (substType mk mt) ts)
-  TTup k ts -> TTup (substKind mk k) (map (substType mk mt) ts)
-  TList k t -> TList (substKind mk k) (substType mk mt t)
-  TFun k t1 t2 -> TFun (substKind mk k) (substType mk mt t1) (substType mk mt t2)
-  TCon k n -> TCon (substKind mk k) n
-  TVar k n -> TVar (substKind mk k) n
-  t@(TExt _ (TUniVar v)) -> fromMaybe t (Map.lookup v mt)
+          (map (second (map (substType mk mt mempty))) cons)
 
+substType :: Map Int CKind  -- ^ kind uvars
+          -> Map Int CType  -- ^ type uvars
+          -> Map Name CType  -- ^ type variables
+          -> CType -> CType
+substType mk mt mtv = go
+  where
+    go (TApp k t ts) = TApp (substKind mk k) (go t) (map go ts)
+    go (TTup k ts) = TTup (substKind mk k) (map go ts)
+    go (TList k t) = TList (substKind mk k) (go t)
+    go (TFun k t1 t2) = TFun (substKind mk k) (go t1) (go t2)
+    go (TCon k n) = TCon (substKind mk k) n
+    go t@(TVar _ n) = fromMaybe t (Map.lookup n mtv)
+    go (TForall k n t) = TForall (substKind mk k) n (go t)
+    go t@(TExt _ (TUniVar v)) = fromMaybe t (Map.lookup v mt)
+
+-- substitute unification variables
 substKind :: Map Int CKind -> CKind -> CKind
 substKind m = \case
   KType () -> KType ()
   KFun () k1 k2 -> KFun () (substKind m k1) (substKind m k2)
   k@(KExt () (KUniVar v)) -> fromMaybe k (Map.lookup v m)
 
-doneProg :: Map Int TKind  -- ^ Kind variable instantiations
-         -> Map Int TType  -- ^ Type variable instantiations
+doneProg :: Map Int TKind  -- ^ Kind unification variable instantiations
+         -> Map Int TType  -- ^ Type unification variable instantiations
          -> CProgram
          -> TProgram
 doneProg = error "doneProg"
@@ -414,11 +526,6 @@ allEq l = case toList l of
 funeqPats :: FunEq t -> [Pattern t]
 funeqPats (FunEq _ _ pats _) = pats
 
-splitKind :: Kind s -> ([Kind s], Either (E Kind s) (X Kind s))
-splitKind (KType x) = ([], Right x)
-splitKind (KFun _ k1 k2) = first (k1:) (splitKind k2)
-splitKind (KExt _ e) = ([], Left e)
-
 isCEqK :: Constr -> Maybe (CKind, CKind, Range)
 isCEqK (CEqK k1 k2 rng) = Just (k1, k2, rng)
 isCEqK _ = Nothing