work

2 files changed, 124 insertions, 9 deletions

diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index e25657b..058f5ac 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -126,6 +126,7 @@ data Expr s
   | EOp   (X Expr s) (Expr s) Operator (Expr s)
   | EIf   (X Expr s) (Expr s) (Expr s) (Expr s)
   | ECase (X Expr s) (Expr s) [(Pattern s, RHS s)]
+  -- TODO: pattern bindings?
   | ELet  (X Expr s) [FunDef s] (Expr s)
   | EError (X Expr s)
 deriving instance (Show (X Expr s), Show (Pattern s), Show (RHS s), Show (FunDef s)) => Show (Expr s)
diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
index f62b097..ad754cf 100644
--- a/src/HSVIS/Typecheck.hs
+++ b/src/HSVIS/Typecheck.hs
@@ -6,6 +6,14 @@
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TupleSections #-}
 {-# LANGUAGE GADTs #-}
+
+
+
+{-# OPTIONS -Wno-unused-top-binds #-}
+{-# OPTIONS -Wno-unused-imports #-}
+
+
+
 module HSVIS.Typecheck (
   StageTyped,
   typecheck,
@@ -95,10 +103,12 @@ 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 mempty)
-      (ds2, subK, subT) = solveConstrs cs
-  in (toList (ds1 <> ds2), doneProg subK subT progtc)
+  let (ds, cs, _, _, resprog) =
+        runTCM (tcTop prog) (fp, source) 1 (Env mempty mempty mempty)
+  in trace ("[tc] resprog = " ++ show resprog) $
+     if not (null cs)
+       then error $ "Constraints left after typechecker completion: " ++ show cs
+       else (toList ds, resprog)
 
 data Constr
   -- Equality constraints: "left" must be equal to "right" because of the thing
@@ -228,6 +238,12 @@ getType' rng name = getType name >>= \case
     genUniVar (KType ())
   Just k -> return k
 
+tcTop :: PProgram -> TCM TProgram
+tcTop prog = do
+  (cs, prog') <- collectConstraints Just (tcProgram prog)
+  (subK, subT) <- solveConstrs cs
+  return $ doneProg subK subT prog'
+
 tcProgram :: PProgram -> TCM CProgram
 tcProgram (Program ddefs1 fdefs1) = do
   (kconstrs, ddefs2) <- collectConstraints isCEqK $ do
@@ -366,6 +382,7 @@ tcFunEq down (FunEq rng name pats rhs) = do
     rhs' <- tcRHS rhsty rhs
     return (FunEq () name pats' rhs')
 
+-- | Brings the bound variables in scope
 tcPattern :: CType -> PPattern -> TCM CPattern
 tcPattern down = \case
   PWildcard _ -> return $ PWildcard down
@@ -405,7 +422,96 @@ tcPattern down = \case
     PTup (TTup (KType ()) ts) <$> zipWithM tcPattern ts ps
 
 tcRHS :: CType -> PRHS -> TCM CRHS
-tcRHS = error "tcRHS"
+tcRHS _ (Guarded _ _) = error "typecheck: Guards not yet supported"
+tcRHS down (Plain _ e) = do
+  e' <- tcExpr down e
+  return $ Plain (extOf e') e'
+
+tcExpr :: CType -> PExpr -> TCM CExpr
+tcExpr down = \case
+  ELit rng lit -> do
+    let ty = case lit of
+               LInt{} -> TCon (KType ()) (Name "Int")
+               LFloat{} -> TCon (KType ()) (Name "Double")
+               LChar{} -> TCon (KType ()) (Name "Char")
+               LString{} -> TList (KType ()) (TCon (KType ()) (Name "Char"))
+    emit $ CEq down ty rng
+    return (ELit ty lit)
+
+  EVar rng n -> EVar <$> getType' rng n <*> pure n
+
+  ECon rng n -> ECon <$> getType' rng n <*> pure n
+
+  EList rng es -> do
+    eltty <- genUniVar (KType ())
+    let listty = TList (KType ()) eltty
+    emit $ CEq down listty rng
+    EList listty <$> mapM (tcExpr listty) es
+
+  ETup rng es -> do
+    ts <- mapM (\_ -> genUniVar (KType ())) es
+    emit $ CEq down (TTup (KType ()) ts) rng
+    ETup (TTup (KType ()) ts) <$> zipWithM tcExpr ts es
+
+  EApp _ e1 es -> do
+    argtys <- mapM (\_ -> genUniVar (KType ())) es
+    let funty = foldr (TFun (KType ())) down argtys
+    EApp funty <$> tcExpr funty e1
+               <*> zipWithM tcExpr argtys es
+
+  -- TODO: these types are way too monomorphic and in any case these
+  -- ~operators~ functions should not be built-in
+  EOp rng e1 op e2 -> do
+    let int = TCon (KType ()) (Name "Int")
+        bool = TCon (KType ()) (Name "Bool")
+    (rty, aty1, aty2) <- case op of
+      OAdd -> return (int, int, int)
+      OSub -> return (int, int, int)
+      OMul -> return (int, int, int)
+      ODiv -> return (int, int, int)
+      OMod -> return (int, int, int)
+      OEqu -> return (bool, int, int)
+      OPow -> return (int, int, int)
+      OCons -> do
+        eltty <- genUniVar (KType ())
+        let listty = TList (KType ()) eltty
+        return (listty, eltty, listty)
+    emit $ CEq down rty rng
+    e1' <- tcExpr aty1 e1
+    e2' <- tcExpr aty2 e2
+    return (EOp rty e1' op e2')
+
+  EIf _ e1 e2 e3 -> do
+    e1' <- tcExpr (TCon (KType ()) (Name "Bool")) e1
+    e2' <- tcExpr down e2
+    e3' <- tcExpr down e3
+    return (EIf down e1' e2' e3')
+
+  ECase _ e1 alts -> do
+    ty <- genUniVar (KType ())
+    e1' <- tcExpr ty e1
+    alts' <- forM alts $ \(pat, rhs) ->
+      scopeVEnv $
+        (,) <$> tcPattern ty pat <*> tcRHS down rhs
+    return $ ECase down e1' alts'
+
+  ELet rng defs body -> do
+    bound <- mapM (\(FunDef _ n _ _) -> (n,) <$> genUniVar (KType ())) defs
+    defs' <- forM defs $ \def@(FunDef _ name _ _) ->
+      scopeVEnv $ do
+        modifyVEnv (Map.fromList [(n, t) | (n, t) <- bound, n /= name] <>)
+        tcFunDef def
+    -- take the actual found types for typechecking the body (and linking them
+    -- to the variables generated above)
+    let bound2 = map (\(FunDef ty n _ _) -> (n, ty)) defs'
+    forM_ (zip bound bound2) $ \((_, tvar), (_, ty)) ->
+      emit $ CEq ty tvar rng  -- in which order? which range? /shrug/
+    scopeVEnv $ do
+      modifyVEnv (Map.fromList bound2 <>)
+      body' <- tcExpr down body
+      return $ ELet down defs' body'
+
+  EError _ -> return $ EError down
 
 unfoldFunTy :: Range -> Int -> CType -> TCM ([CType], CType)
 unfoldFunTy _ n t | n <= 0 = return ([], t)
@@ -419,6 +525,13 @@ unfoldFunTy rng n t = do
   emit $ CEq expected t rng
   return (vars, core)
 
+solveConstrs :: Bag Constr -> TCM (Map Int CKind, Map Int CType)
+solveConstrs constrs = do
+  let (tcs, kcs) = partitionConstrs constrs
+  subK <- solveKindVars kcs
+  subT <- solveTypeVars tcs
+  return (subK, subT)
+
 solveKindVars :: Bag (CKind, CKind, Range) -> TCM (Map Int CKind)
 solveKindVars cs = do
   let (asg, errs) =
@@ -467,8 +580,9 @@ solveKindVars cs = do
     kindSize (KFun () a b) = 1 + kindSize a + kindSize b
     kindSize (KExt () KUniVar{}) = 2
 
-solveConstrs :: Bag Constr -> (Bag Diagnostic, Map Int TKind, Map Int TType)
-solveConstrs = error "solveConstrs"
+partitionConstrs :: Foldable t => t Constr -> (Bag (CType, CType, Range), Bag (CKind, CKind, Range))
+partitionConstrs = foldMap $ \case CEq t1 t2 r -> (pure (t1, t2, r), mempty)
+                                   CEqK k1 k2 r -> (mempty, pure (k1, k2, r))
 
 -- substitute unification variables
 substProg :: Map Int CKind  -- ^ Kind unification variable instantiations
@@ -495,9 +609,9 @@ substType mk mt mtv = go
     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 (TVar k n) = fromMaybe (TVar (substKind mk k) n) (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)
+    go (TExt k (TUniVar v)) = fromMaybe (TExt (substKind mk k) (TUniVar v)) (Map.lookup v mt)
 
 -- substitute unification variables
 substKind :: Map Int CKind -> CKind -> CKind