Some typechecker work

1 files changed, 92 insertions, 20 deletions

diff --git a/src/HSVIS/Typecheck.hs b/src/HSVIS/Typecheck.hs
index de9d7db..ba853a0 100644
--- a/src/HSVIS/Typecheck.hs
+++ b/src/HSVIS/Typecheck.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE EmptyDataDeriving #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeFamilies #-}
@@ -7,11 +8,16 @@
 module HSVIS.Typecheck where
 
 import Control.Monad
-import Data.Bifunctor (first, second)
+import Data.Bifunctor (first)
 import Data.Foldable (toList)
 import Data.Map.Strict (Map)
-import Data.Monoid (First(..))
+import Data.Maybe (fromMaybe)
+import Data.Monoid (Ap(..))
 import qualified Data.Map.Strict as Map
+import Data.Set (Set)
+import qualified Data.Set as Set
+
+import Debug.Trace
 
 import Data.Bag
 import Data.List.NonEmpty.Util
@@ -19,6 +25,7 @@ import HSVIS.AST
 import HSVIS.Parser
 import HSVIS.Diagnostic
 import HSVIS.Pretty
+import HSVIS.Typecheck.Solve
 
 
 data StageTC
@@ -33,7 +40,8 @@ type instance X RHS     StageTC = CType
 type instance X Expr    StageTC = CType
 
 data instance E Type StageTC = TUniVar Int deriving (Show)
-data instance E Kind StageTC = KUniVar Int deriving (Show)
+data instance E Kind StageTC = KUniVar Int deriving (Show, Eq, Ord)
+data instance E TypeSig StageTC deriving (Show)
 
 type CProgram = Program StageTC
 type CDataDef = DataDef StageTC
@@ -58,6 +66,7 @@ type instance X Expr    StageTyped = TType
 
 data instance E Type StageTyped deriving (Show)
 data instance E Kind StageTyped deriving (Show)
+data instance E TypeSig StageTyped deriving (Show)
 
 type TProgram = Program StageTyped
 type TDataDef = DataDef StageTyped
@@ -69,8 +78,11 @@ type TPattern = Pattern StageTyped
 type TRHS     = RHS     StageTyped
 type TExpr    = Expr    StageTyped
 
+instance Pretty (E Kind StageTC) where
+  prettysPrec _ (KUniVar n) = showString ("?k" ++ show n)
+
 
-typecheck :: FilePath -> String -> PProgram -> ([Diagnostic], Program TType)
+typecheck :: FilePath -> String -> PProgram -> ([Diagnostic], TProgram)
 typecheck fp source prog =
   let (ds1, cs, _, _, progtc) =
         runTCM (tcProgram prog) (fp, source) 1 (Env mempty mempty)
@@ -269,29 +281,86 @@ tcFunDef (FunDef _ name msig eqs) = do
   return (FunDef typ name (TypeSig typ) eqs')
 
 tcFunEq :: CType -> PFunEq -> TCM CFunEq
-tcFunEq = _
+tcFunEq = error "tcFunEq"
+
+newtype SolveM v t m a = SolveM (Map v (Bag t) -> Map v t -> m (a, Map v (Bag t), Map v t))
+instance Monad m => Functor (SolveM v t m) where
+  fmap f (SolveM g) = SolveM $ \m r -> do (x, m', r') <- g m r
+                                          return (f x, m', r')
+instance Monad m => Applicative (SolveM v t m) where
+  pure x = SolveM $ \m r -> return (x, m, r)
+  (<*>) = ap
+instance Monad m => Monad (SolveM v t m) where
+  SolveM f >>= g = SolveM $ \m r -> do (x, m1, r1) <- f m r
+                                       let SolveM h = g x
+                                       h m1 r1
+
+solvemStateGet :: Monad m => SolveM v t m (Map v (Bag t))
+solvemStateGet = SolveM $ \m r -> return (m, m, r)
+
+solvemStateUpdate :: Monad m => (Map v (Bag t) -> Map v (Bag t)) -> SolveM v t m ()
+solvemStateUpdate f = SolveM $ \m r -> return ((), f m, r)
+
+solvemLogUpdate :: Monad m => (Map v t -> Map v t) -> SolveM v t m ()
+solvemLogUpdate f = SolveM $ \m r -> return ((), m, f r)
+
+solvemStateVars :: Monad m => SolveM v t m [v]
+solvemStateVars = Map.keys <$> solvemStateGet
+
+solvemStateRHS :: (Ord v, Monad m) => v -> SolveM v t m (Bag t)
+solvemStateRHS v = fromMaybe mempty . Map.lookup v <$> solvemStateGet
+
+solvemStateSet :: (Ord v, Monad m) => v -> Bag t -> SolveM v t m ()
+solvemStateSet v b = solvemStateUpdate (Map.insert v b)
+
+solvemLogEq :: (Ord v, Monad m) => v -> t -> SolveM v t m ()
+solvemLogEq v t = solvemLogUpdate (Map.insert v t)
 
 solveKindVars :: Bag (CKind, CKind, Range) -> TCM ()
-solveKindVars =
-  mapM_ $ \(a, b, rng) -> do
-    let (subst, First merr) = reduce a b
-    forM_ merr $ \(erra, errb) ->
-      raise rng $
-        "Kind mismatch:\n\
-        \- Expected: " ++ pretty a ++ "\n\
-        \- Observed: " ++ pretty b ++ "\n\
-        \because '" ++ pretty erra ++ "' and '" ++ pretty errb ++ "' don't match"
-    let collected :: [(Int, Bag CKind)]
-        collected = Map.assocs $ Map.fromListWith (<>) (fmap (second pure) (toList subst))
-    _
+solveKindVars cs = do
+  traceShowM cs
+  traceShowM $ solveConstraints
+                 reduce
+                 (foldMap pure . kindUniVars)
+                 (\v repl -> substKind (Map.singleton v repl))
+                 (\case KExt () (KUniVar v) -> Just v
+                        _ -> Nothing)
+                 kindSize
+                 (map (\(a, b, _) -> (a, b)) (toList cs))
   where
-    reduce :: CKind -> CKind -> (Bag (Int, CKind), First (CKind, CKind))
-    reduce (KType ()) (KType ()) = mempty
-    reduce (KFun () a b) (KFun () c d) = reduce a c <> reduce b d
+    reduce :: CKind -> CKind -> (Bag (Int, CKind), Bag (CKind, CKind))
+    -- unification variables produce constraints on a unification variable
+    reduce (KExt () (KUniVar i)) (KExt () (KUniVar j)) | i == j = mempty
     reduce (KExt () (KUniVar i)) k = (pure (i, k), mempty)
     reduce k (KExt () (KUniVar i)) = (pure (i, k), mempty)
+    -- if lhs and rhs have equal prefixes, recurse
+    reduce (KType ()) (KType ()) = mempty
+    reduce (KFun () a b) (KFun () c d) = reduce a c <> reduce b d
+    -- otherwise, this is a kind mismatch
     reduce k1 k2 = (mempty, pure (k1, k2))
 
+    kindSize :: CKind -> Int
+    kindSize KType{} = 1
+    kindSize (KFun () a b) = 1 + kindSize a + kindSize b
+    kindSize (KExt () KUniVar{}) = 1
+
+solveConstrs :: Bag Constr -> (Bag Diagnostic, Map Name TType)
+solveConstrs = error "solveConstrs"
+
+substProg :: Map Name TType -> CProgram -> TProgram
+substProg = error "substProg"
+
+substKind :: Map Int CKind -> CKind -> CKind
+substKind _ k@KType{} = k
+substKind m (KFun () k1 k2) = KFun () (substKind m k1) (substKind m k2)
+substKind m k@(KExt () (KUniVar v)) = fromMaybe k (Map.lookup v m)
+
+kindUniVars :: CKind -> Set Int
+kindUniVars = \case
+  KType{} -> mempty
+  KFun () a b -> kindUniVars a <> kindUniVars b
+  KExt () (KUniVar v) -> Set.singleton v
+
 allEq :: (Eq a, Foldable t) => t a -> Bool
 allEq l = case toList l of
             [] -> True
@@ -308,3 +377,6 @@ splitKind (KExt _ e) = ([], Left e)
 isCEqK :: Constr -> Maybe (CKind, CKind, Range)
 isCEqK (CEqK k1 k2 rng) = Just (k1, k2, rng)
 isCEqK _ = Nothing
+
+foldMapM :: (Applicative f, Monoid m, Foldable t) => (a -> f m) -> t a -> f m
+foldMapM f = getAp . foldMap (Ap . f)