on ee72256 help sign in
elk/src/Type/State.hs
200 lines 6.6 kB history raw
1module Type.State where
3import Control.Applicative ( (<$>), (<*>), Applicative, (<|>) )
4import Control.Monad.State
5import qualified Data.Map as Map
6import qualified Data.Traversable as Traversable
7import qualified Data.UnionFind.IO as UF
9import qualified AST.Annotation as A
10import AST.PrettyPrint
11import Text.PrettyPrint as P
12import Type.Type
14-- Pool
15-- Holds a bunch of variables
16-- The rank of each variable is less than or equal to the pool's "maxRank"
17-- The young pool exists to make it possible to identify these vars in constant time.
19data Pool = Pool
20 { maxRank :: Int
21 , inhabitants :: [Variable]
22 } deriving Show
24emptyPool = Pool { maxRank = outermostRank, inhabitants = [] }
26type Env = Map.Map String Variable
28-- Keeps track of the environment, type variable pool, and a list of errors
29data SolverState = SS {
30 sEnv :: Env,
31 sSavedEnv :: Env,
32 sPool :: Pool,
33 sMark :: Int,
34 sErrors :: [P.Doc]
37initialState = SS {
38 sEnv = Map.empty,
39 sSavedEnv = Map.empty,
40 sPool = emptyPool,
41 sMark = noMark + 1, -- The mark must never be equal to noMark!
42 sErrors = []
45modifyEnv f = modify $ \state -> state { sEnv = f (sEnv state) }
46modifyPool f = modify $ \state -> state { sPool = f (sPool state) }
48addError :: A.Region -> Maybe String -> UF.Point Descriptor -> UF.Point Descriptor
49 -> StateT SolverState IO ()
50addError region hint t1 t2 =
51 do err <- liftIO makeError
52 modify $ \state -> state { sErrors = err : sErrors state }
53 where
54 makeError = do
55 t1' <- pretty <$> toSrcType t1
56 t2' <- pretty <$> toSrcType t2
57 return . P.vcat $
58 [ P.text "Type error" <+> pretty region <> P.colon
59 , maybe P.empty P.text hint
60 , P.text ""
61 , P.nest 8 $ A.getRegionDocs region
62 , P.text ""
63 , P.text " Expected Type:" <+> t1'
64 , P.text " Actual Type:" <+> t2'
65 ]
67switchToPool pool = modifyPool (\_ -> pool)
69getPool :: StateT SolverState IO Pool
70getPool = sPool <$> get
72getEnv :: StateT SolverState IO Env
73getEnv = sEnv <$> get
75saveLocalEnv :: StateT SolverState IO ()
76saveLocalEnv = do
77 env <- sEnv <$> get
78 modify $ \state -> state { sSavedEnv = env }
80uniqueMark :: StateT SolverState IO Int
81uniqueMark = do
82 state <- get
83 let mark = sMark state
84 put $ state { sMark = mark + 1 }
85 return mark
87nextRankPool :: StateT SolverState IO Pool
88nextRankPool = do
89 pool <- getPool
90 return $ Pool { maxRank = maxRank pool + 1, inhabitants = [] }
92register :: Variable -> StateT SolverState IO Variable
93register variable = do
94 modifyPool $ \pool -> pool { inhabitants = variable : inhabitants pool }
95 return variable
97introduce :: Variable -> StateT SolverState IO Variable
98introduce variable = do
99 pool <- getPool
100 liftIO $ UF.modifyDescriptor variable (\desc -> desc { rank = maxRank pool })
101 register variable
103flatten :: Type -> StateT SolverState IO Variable
104flatten term =
105 case term of
106 VarN maybeAlias v -> do
107 liftIO $ UF.modifyDescriptor v $ \desc -> desc { alias = maybeAlias <|> alias desc }
108 return v
109 TermN maybeAlias t -> do
110 flatStructure <- traverseTerm flatten t
111 pool <- getPool
112 var <- liftIO . UF.fresh $ Descriptor
113 { structure = Just flatStructure
114 , rank = maxRank pool
115 , flex = Flexible
116 , name = Nothing
117 , copy = Nothing
118 , mark = noMark
119 , alias = maybeAlias
120 }
121 register var
123makeInstance :: Variable -> StateT SolverState IO Variable
124makeInstance var = do
125 alreadyCopied <- uniqueMark
126 freshVar <- makeCopy alreadyCopied var
127 restore alreadyCopied var
128 return freshVar
130makeCopy :: Int -> Variable -> StateT SolverState IO Variable
131makeCopy alreadyCopied variable = do
132 desc <- liftIO $ UF.descriptor variable
133 case () of
134 () | mark desc == alreadyCopied ->
135 case copy desc of
136 Just v -> return v
137 Nothing -> error $ "Error copying type variable. This should be impossible." ++
138 " Please report an error to the github repo!"
140 | rank desc /= noRank || flex desc == Constant ->
141 return variable
143 | otherwise -> do
144 pool <- getPool
145 newVar <- liftIO $ UF.fresh $ Descriptor
146 { structure = Nothing
147 , rank = maxRank pool
148 , mark = noMark
149 , flex = case flex desc of
150 Is s -> Is s
151 _ -> Flexible
152 , copy = Nothing
153 , name = case flex desc of
154 Rigid -> Nothing
155 _ -> name desc
156 , alias = Nothing
157 }
158 register newVar
160 -- Link the original variable to the new variable. This lets us
161 -- avoid making multiple copies of the variable we are instantiating.
162 --
163 -- Need to do this before recursively copying the structure of
164 -- the variable to avoid looping on cyclic terms.
165 liftIO $ UF.modifyDescriptor variable $ \desc ->
166 desc { mark = alreadyCopied, copy = Just newVar }
168 -- Now we recursively copy the structure of the variable.
169 -- We have already marked the variable as copied, so we
170 -- will not repeat this work or crawl this variable again.
171 case structure desc of
172 Nothing -> return newVar
173 Just term -> do
174 newTerm <- traverseTerm (makeCopy alreadyCopied) term
175 liftIO $ UF.modifyDescriptor newVar $ \desc ->
176 desc { structure = Just newTerm }
177 return newVar
179restore :: Int -> Variable -> StateT SolverState IO Variable
180restore alreadyCopied variable = do
181 desc <- liftIO $ UF.descriptor variable
182 if mark desc /= alreadyCopied
183 then return variable
184 else do
185 restoredStructure <-
186 Traversable.traverse (traverseTerm (restore alreadyCopied)) (structure desc)
187 liftIO $ UF.modifyDescriptor variable $ \desc ->
188 desc { mark = noMark, rank = noRank, structure = restoredStructure }
189 return variable
191traverseTerm :: (Monad f, Applicative f) => (a -> f b) -> Term1 a -> f (Term1 b)
192traverseTerm f term =
193 case term of
194 App1 a b -> App1 <$> f a <*> f b
195 Fun1 a b -> Fun1 <$> f a <*> f b
196 Var1 x -> Var1 <$> f x
197 EmptyRecord1 -> return EmptyRecord1
198 Record1 fields ext ->
199 Record1 <$> Traversable.traverse (mapM f) fields <*> f ext