aboutsummaryrefslogtreecommitdiff
path: root/src/Haskell/AST.hs
blob: 6d25153201e7a021c5d8df9c38e5df2d6a349014 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
module Haskell.AST where

import Data.List
import qualified Data.Set as Set
import Pretty


type Name = String
type TyVar = String

data AST = AST [Toplevel]
  deriving (Show, Eq)

data Toplevel = TopDef Def
              | TopDecl Decl
              | TopData Data
              | TopClass Class
              | TopInst Inst
  deriving (Show, Eq)

data Def = Def Name Expr
  deriving (Show, Eq)

data Expr = App Expr [Expr]
          | Ref Name
          | Con Name
          | Num Integer
          | Tup [Expr]
          | Lam [Name] Expr
          | Case Expr [(Pat, Expr)]
  deriving (Show, Eq)

data Pat = PatAny
         | PatVar Name
         | PatCon Name [Pat]
         | PatTup [Pat]
  deriving (Show, Eq)

data Decl = Decl Name Type
  deriving (Show, Eq)

data Type = TyTup [Type]
          | TyInt
          | TyFun Type Type
          | TyRef Name [Type]
          | TyVar Name
          | TyVoid
  deriving (Show, Eq)

data Data = Data Name [TyVar] [(Name, [Type])]
  deriving (Show, Eq)

data Class = Class Name [TyVar] [Decl]
  deriving (Show, Eq)

data Inst = Inst Name Type [Def]
  deriving (Show, Eq)


instance Pretty AST where
    pretty (AST tops) = Node "" (map pretty tops)

instance Pretty Toplevel where
    pretty (TopDef x) = pretty x
    pretty (TopDecl x) = pretty x
    pretty (TopData x) = pretty x
    pretty (TopClass x) = pretty x
    pretty (TopInst x) = pretty x

instance Pretty Def where
    pretty (Def n e) = Node (n ++ " =") [pretty e]

instance Pretty Expr where
    pretty (App e as) = Bracket "(" ")" "" (map pretty (e:as))
    pretty (Ref n) = Leaf n
    pretty (Con n) = Leaf n
    pretty (Num n) = Leaf (show n)
    pretty (Tup es) = Bracket "(" ")" "," (map pretty es)
    pretty (Lam as e) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
    pretty (Case e arms) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
      where go (p, e') = Node (pprintOneline p ++ " ->") [pretty e']

instance Pretty Pat where
    pretty PatAny = Leaf "_"
    pretty (PatVar n) = Leaf n
    pretty (PatCon n ps) = Bracket "(" ")" "" (Leaf n : map pretty ps)
    pretty (PatTup ps) = Bracket "(" ")" "," (map pretty ps)

instance Pretty Decl where
    pretty (Decl n t) = Node (n ++ " :: ") [pretty t]

instance Pretty Type where
    pretty (TyTup ts) = Bracket "(" ")" "," (map pretty ts)
    pretty TyInt = Leaf "Int"
    pretty (TyFun t u) = Leaf $ pprintOneline t ++ " -> " ++ pprintOneline u
    pretty (TyRef n as) = Bracket "(" ")" "" (Leaf n : map pretty as)
    pretty (TyVar n) = Leaf n
    pretty TyVoid = Leaf "#Void"

instance Pretty Data where
    pretty (Data n as cs) = Node ("data " ++ n ++ " " ++ intercalate " " as ++ " =") [Bracket "" "" "|" (map go cs)]
      where go (m, ts) = Node m (map pretty ts)

instance Pretty Class where
    pretty (Class n as ds) = Node ("class " ++ n ++ " " ++ intercalate " " as ++ " where") [Bracket "{" "}" ";" (map pretty ds)]

instance Pretty Inst where
    pretty (Inst n t ds) = Node ("instance " ++ n ++ " " ++ pprintOneline t ++ " where") [Bracket "{" "}" ";" (map pretty ds)]


-- This excludes constructor names, since those are not variables. This _does_
-- include bound variables; if you don't want that, use freeVariables.
class AllVars a where
    allVars :: a -> Set.Set Name

instance AllVars AST where
    allVars (AST tops) = Set.unions (map allVars tops)

instance AllVars Toplevel where
    allVars (TopDef def) = allVars def
    allVars (TopDecl _) = mempty
    allVars (TopData _) = mempty
    allVars (TopClass _) = mempty
    allVars (TopInst inst) = allVars inst

instance AllVars Def where
    allVars (Def n e) = Set.insert n (allVars e)

instance AllVars Inst where
    allVars (Inst _ _ ds) = Set.unions (map allVars ds)

instance AllVars Expr where
    allVars (App e es) = Set.unions (map allVars (e : es))
    allVars (Ref n) = Set.singleton n
    allVars (Con _) = mempty
    allVars (Num _) = mempty
    allVars (Tup es) = Set.unions (map allVars es)
    allVars (Lam ns e) = Set.fromList ns <> allVars e
    allVars (Case e pairs) =
        allVars e <> Set.unions [allVars p <> allVars e' | (p, e') <- pairs]

instance AllVars Pat where
    allVars PatAny = mempty
    allVars (PatVar n) = Set.singleton n
    allVars (PatCon _ ps) = Set.unions (map allVars ps)
    allVars (PatTup ps) = Set.unions (map allVars ps)


boundVars :: Pat -> Set.Set Name
boundVars PatAny = mempty
boundVars (PatVar n) = Set.singleton n
boundVars (PatCon _ ps) = Set.unions (map boundVars ps)
boundVars (PatTup ps) = Set.unions (map boundVars ps)

freeVariables :: Expr -> Set.Set Name
freeVariables (App e es) = freeVariables e <> Set.unions (map freeVariables es)
freeVariables (Ref n) = Set.singleton n
freeVariables (Con _) = mempty
freeVariables (Num _) = mempty
freeVariables (Tup es) = Set.unions (map freeVariables es)
freeVariables (Lam ns e) = freeVariables e Set.\\ Set.fromList ns
freeVariables (Case e pairs) =
    freeVariables e <> Set.unions [freeVariables e' Set.\\ boundVars p
                                  | (p, e') <- pairs]