WIP refactor and union types, type variables

2 files changed, 71 insertions, 14 deletions

diff --git a/ast/CC/AST/Source.hs b/ast/CC/AST/Source.hs
index e648759..e64e058 100644
--- a/ast/CC/AST/Source.hs
+++ b/ast/CC/AST/Source.hs
@@ -3,6 +3,8 @@ module CC.AST.Source(
     module CC.Types
 ) where
 
+import qualified Data.Set as Set
+import Data.Set (Set)
 import Data.List
 
 import CC.Pretty
@@ -12,19 +14,32 @@ import CC.Types
 data Program = Program [Decl]
   deriving (Show, Read)
 
-data Decl = Def Def  -- import?
+data Decl = DeclFunc FuncDef
+          | DeclType TypeDef
+          | DeclAlias AliasDef
   deriving (Show, Read)
 
-data Def = Function (Maybe Type)
-                    (Name, SourceRange)
-                    [(Name, SourceRange)]
-                    Expr
+data FuncDef =
+    FuncDef (Maybe Type)
+            (Name, SourceRange)
+            [(Name, SourceRange)]
+            Expr
+  deriving (Show, Read)
+
+-- Named type with named arguments
+data TypeDef = TypeDef (Name, SourceRange) [(Name, SourceRange)] Type
+  deriving (Show, Read)
+
+data AliasDef = AliasDef (Name, SourceRange) [(Name, SourceRange)] Type
   deriving (Show, Read)
 
 data Type = TFun Type Type
           | TInt
           | TTup [Type]
-  deriving (Show, Read)
+          | TNamed Name [Type]  -- named type with type arguments
+          | TUnion (Set Type)
+          | TyVar Name
+  deriving (Eq, Ord, Show, Read)
 
 data Expr = Lam SourceRange [(Name, SourceRange)] Expr
           | Let SourceRange (Name, SourceRange) Expr Expr
@@ -32,15 +47,24 @@ data Expr = Lam SourceRange [(Name, SourceRange)] Expr
           | Int SourceRange Int
           | Tup SourceRange [Expr]
           | Var SourceRange Name
+          | Constr SourceRange Name  -- type constructor
           | Annot SourceRange Expr Type
   deriving (Show, Read)
 
+instance Semigroup Program where Program p1 <> Program p2 = Program (p1 <> p2)
+instance Monoid    Program where mempty                   = Program mempty
+
 instance Pretty Type where
     prettyPrec _ TInt = "Int"
     prettyPrec p (TFun a b) =
         precParens p 2 (prettyPrec 3 a ++ " -> " ++ prettyPrec 2 b)
     prettyPrec _ (TTup ts) =
         "(" ++ intercalate ", " (map pretty ts) ++ ")"
+    prettyPrec _ (TNamed n ts) =
+        n ++ "[" ++ intercalate ", " (map pretty ts) ++ "]"
+    prettyPrec _ (TUnion ts) =
+        "{" ++ intercalate " | " (map pretty (Set.toList ts)) ++ "}"
+    prettyPrec _ (TyVar n) = "<" ++ n ++ ">"
 
 instance HasRange Expr where
     range (Lam sr _ _) = sr
@@ -49,4 +73,5 @@ instance HasRange Expr where
     range (Int sr _) = sr
     range (Tup sr _) = sr
     range (Var sr _) = sr
+    range (Constr sr _) = sr
     range (Annot sr _ _) = sr
diff --git a/ast/CC/AST/Typed.hs b/ast/CC/AST/Typed.hs
index b12b30a..cf67575 100644
--- a/ast/CC/AST/Typed.hs
+++ b/ast/CC/AST/Typed.hs
@@ -3,23 +3,36 @@ module CC.AST.Typed(
     module CC.Types
 ) where
 
+import qualified Data.Map.Strict as Map
+import Data.Map.Strict (Map)
+import qualified Data.Set as Set
+import Data.Set (Set)
 import Data.List
 
 import CC.Pretty
 import CC.Types
 
 
-data Program = Program [Def]
+data Program = Program [Def] (Map Name TypeDef)
   deriving (Show, Read)
 
 data Def = Def Name Expr
   deriving (Show, Read)
 
+-- Named type with type parameters
+data TypeDef = TypeDef Name [Int] Type
+  deriving (Show, Read)
+
 data Type = TFun Type Type
           | TInt
           | TTup [Type]
-          | TyVar Int
-  deriving (Show, Read)
+          | TNamed Name [Type]  -- named type with type arguments
+          | TUnion (Set Type)
+          | TyVar Rigidity Int
+  deriving (Eq, Ord, Show, Read)
+
+data Rigidity = Rigid | Instantiable
+  deriving (Eq, Ord, Show, Read)
 
 data TypeScheme = TypeScheme [Int] Type
   deriving (Show, Read)
@@ -30,6 +43,7 @@ data Expr = Lam Type Occ Expr
           | Int Int
           | Tup [Expr]
           | Var Occ
+          | Constr Type Name  -- Type is 'argument -> TNamed'
   deriving (Show, Read)
 
 data Occ = Occ Name Type
@@ -42,6 +56,7 @@ exprType (Call typ _ _) = typ
 exprType (Int _) = TInt
 exprType (Tup es) = TTup (map exprType es)
 exprType (Var (Occ _ typ)) = typ
+exprType (Constr typ _) = typ
 
 instance Pretty Type where
     prettyPrec _ TInt = "Int"
@@ -49,13 +64,18 @@ instance Pretty Type where
         precParens p 2 (prettyPrec 3 a ++ " -> " ++ prettyPrec 2 b)
     prettyPrec _ (TTup ts) =
         "(" ++ intercalate ", " (map pretty ts) ++ ")"
-    prettyPrec _ (TyVar i) = 't' : show i
+    prettyPrec _ (TNamed n ts) =
+        n ++ "[" ++ intercalate ", " (map pretty ts) ++ "]"
+    prettyPrec _ (TUnion ts) =
+        "{ " ++ intercalate " | " (map pretty (Set.toList ts)) ++ " }"
+    prettyPrec _ (TyVar Rigid i) = 't' : show i ++ "R"
+    prettyPrec _ (TyVar Instantiable i) = 't' : show i
 
 instance Pretty TypeScheme where
     prettyPrec p (TypeScheme bnds ty) =
         precParens p 2
-            ("forall " ++ intercalate " " (map (pretty . TyVar) bnds) ++ ". " ++
-                prettyPrec 2 ty)
+            ("forall " ++ intercalate " " (map (pretty . TyVar Instantiable) bnds) ++
+                ". " ++ prettyPrec 2 ty)
 
 instance Pretty Expr where
     prettyPrec p (Lam ty (Occ n t) e) =
@@ -75,10 +95,22 @@ instance Pretty Expr where
     prettyPrec _ (Tup es) = "(" ++ intercalate ", " (map pretty es) ++ ")"
     prettyPrec p (Var (Occ n t)) =
         precParens p 2 $
-            show n ++ " :: " ++ pretty t
+            n ++ " :: " ++ pretty t
+    prettyPrec p (Constr t n) =
+        precParens p 2 $
+            n ++ " :: " ++ pretty t
 
 instance Pretty Def where
     pretty (Def n e) = n ++ " = " ++ pretty e
 
+instance Pretty TypeDef where
+    pretty (TypeDef n [] t) =
+        "type " ++ n ++ " = " ++ pretty t
+    pretty (TypeDef n vs t) =
+        "type " ++ n ++ " " ++ intercalate " " (map (pretty . TyVar Instantiable) vs) ++
+            " = " ++ pretty t
+
 instance Pretty Program where
-    pretty (Program defs) = concatMap ((++ "\n") . pretty) defs
+    pretty (Program defs tdefs) =
+        concatMap (++ "\n") $
+            map pretty (Map.elems tdefs) ++ map pretty defs