on 2a930fd help sign in
elk/src/AST/Type.hs
138 lines 4.2 kB history raw
1{-# OPTIONS_GHC -W #-}
2module AST.Type where
4import Control.Applicative ((<$>), (<*>))
5import Data.Binary
6import qualified Data.Map as Map
8import qualified AST.Variable as Var
9import AST.PrettyPrint
10import qualified AST.Helpers as Help
11import Text.PrettyPrint as P
13data Type var
14 = Lambda (Type var) (Type var)
15 | Var String
16 | Type var
17 | App (Type var) [Type var]
18 | Record [(String, Type var)] (Maybe (Type var))
19 | Aliased Var.Canonical (Type var)
20 deriving (Eq,Show)
22type RawType = Type Var.Raw
23type CanonicalType = Type Var.Canonical
25fieldMap :: [(String,a)] -> Map.Map String [a]
26fieldMap fields =
27 foldl (\r (x,t) -> Map.insertWith (++) x [t] r) Map.empty fields
29recordOf :: [(String, Type var)] -> Type var
30recordOf fields = Record fields Nothing
32listOf :: RawType -> RawType
33listOf t = App (Type (Var.Raw "_List")) [t]
35tupleOf :: [RawType] -> RawType
36tupleOf ts = App (Type t) ts
37 where
38 t = Var.Raw ("_Tuple" ++ show (length ts))
40instance (Var.ToString var, Pretty var) => Pretty (Type var) where
41 pretty tipe =
42 case tipe of
43 Lambda _ _ -> P.sep [ t, P.sep (map (P.text "->" <+>) ts) ]
44 where
45 t:ts = map prettyLambda (collectLambdas tipe)
46 prettyLambda t = case t of
47 Lambda _ _ -> P.parens (pretty t)
48 _ -> pretty t
50 Var x -> P.text x
52 Type var ->
53 let v = Var.toString var in
54 P.text (if v == "_Tuple0" then "()" else v)
56 App f args ->
57 case (f,args) of
58 (Type name, [t])
59 | Var.toString name == "_List" -> P.brackets (pretty t)
61 (Type name, _)
62 | Help.isTuple (Var.toString name) ->
63 P.parens . P.sep . P.punctuate P.comma $ map pretty args
65 _ -> P.hang (pretty f) 2 (P.sep $ map prettyParens args)
67 Record _ _ ->
68 P.braces $ case flattenRecord tipe of
69 (fields, Nothing) -> prettyFields fields
70 (fields, Just x) ->
71 P.hang (P.text x <+> P.text "|") 4 (prettyFields fields)
72 where
73 prettyField (f,t) = P.text f <+> P.text ":" <+> pretty t
74 prettyFields fields = commaSep (map prettyField fields)
76 Aliased name t ->
77 let t' = pretty t in
78 if show t' `elem` ["Int", "Float", "String", "Char", "Bool"]
79 then t'
80 else pretty name
82collectLambdas :: Type var -> [Type var]
83collectLambdas tipe =
84 case tipe of
85 Lambda arg body -> arg : collectLambdas body
86 _ -> [tipe]
88prettyParens :: (Var.ToString var, Pretty var) => Type var -> Doc
89prettyParens tipe = parensIf (needed tipe) (pretty tipe)
90 where
91 needed t =
92 case t of
93 Aliased _ t' -> needed t'
95 Lambda _ _ -> True
97 App (Type name) [_] | Var.toString name == "_List" -> False
98 App (Type name) _ | Help.isTuple (Var.toString name) -> False
99 App t' [] -> needed t'
100 App _ _ -> True
102 _ -> False
104flattenRecord :: Type var -> ( [(String, Type var)], Maybe String )
105flattenRecord tipe =
106 case tipe of
107 Var x -> ([], Just x)
109 Record fields Nothing -> (fields, Nothing)
111 Record fields (Just ext) ->
112 let (fields',ext') = flattenRecord ext
113 in (fields' ++ fields, ext')
115 Aliased _ tipe' -> flattenRecord tipe'
117 _ -> error "Trying to flatten ill-formed record."
119instance Binary var => Binary (Type var) where
120 put tipe =
121 case tipe of
122 Lambda t1 t2 -> putWord8 0 >> put t1 >> put t2
123 Var x -> putWord8 1 >> put x
124 Type name -> putWord8 2 >> put name
125 App t1 t2 -> putWord8 3 >> put t1 >> put t2
126 Record fs ext -> putWord8 4 >> put fs >> put ext
127 Aliased var t -> putWord8 5 >> put var >> put t
129 get = do
130 n <- getWord8
131 case n of
132 0 -> Lambda <$> get <*> get
133 1 -> Var <$> get
134 2 -> Type <$> get
135 3 -> App <$> get <*> get
136 4 -> Record <$> get <*> get
137 5 -> Aliased <$> get <*> get
138 _ -> error "Error reading a valid type from serialized string"