Pass failure in assignTypeVars out of inferenceStep
[darcs-mirror-sem_syn.git] / Type.hs
1 module Type where
2
3 import AST 
4
5 import Data.Graph 
6 import Control.Monad.State
7 import Control.Monad.Error
8 import Util
9 import Data.Maybe
10 import Data.List (nub,nubBy,union)
11
12
13 import Data.Map (Map)
14 import qualified Data.Map as Map
15
16 -- type erasure
17 eraseType (AST decls) =
18     AST $ map (\(Decl f ftype ps e) ->
19              Decl f FTUndet (map eraseTypeP ps) (eraseTypeE e)) decls 
20
21 eraseTypeP (PVar id t varname)   
22     = PVar id TUndet varname
23 eraseTypeP (PCon id t conname ps)
24     = PCon id TUndet conname (map eraseTypeP ps)
25
26 eraseTypeE (EVar id t varname) 
27     = EVar id TUndet varname 
28 eraseTypeE (ECon id t conname es)
29     = ECon id TUndet conname (map eraseTypeE es)
30 eraseTypeE (EFun id t funname es)
31     = EFun id TUndet funname (map eraseTypeE es)
32
33 eraseTypeT (TAST decls) = 
34     TAST $ map (\(TDecl f ps es bs) -> 
35                     TDecl f (map eraseTypeP ps) (map eraseTypeE es)bs) decls
36
37 -- type inference
38
39 initTMap :: [ (Name, FType) ]
40 initTMap =
41     [ (Name "Z",   TFun [] [] (TCon (Name "Nat") [])),
42       (Name "S",   TFun [] [TCon (Name "Nat") []] (TCon (Name "Nat") [])),
43       (Name "Nil",  TFun [0] [] (TCon (Name "List") [TVar 0])),
44       (Name "Cons", TFun [0] [TVar 0, TCon (Name "List") [TVar 0]] 
45                 (TCon (Name "List") [TVar 0])) ]
46
47
48 typeInference (AST decls) = 
49     let mAst = do { (decls',_,_) <- 
50                         foldr (\decls m -> 
51                                    do (rdecls, tMap,  icount)  <- m
52                                       (decls', tMap', icount') <- inferenceStep decls tMap icount
53                                       return $ (decls'++rdecls, tMap', icount')
54                               ) (return ([],initTMap,initIcount)) declss
55                   ; return $ AST decls' } 
56     in case mAst of 
57          Left s  -> error s 
58          Right a ->  a 
59     where
60       initIcount = 100 -- FIXME 
61       declss = 
62           let scc = stronglyConnComp callGraph 
63           in  reverse $ map (\x -> case x of 
64                            AcyclicSCC f  -> 
65                                filter (\(Decl g _ _ _) -> f == g) decls
66                            CyclicSCC  fs -> 
67                                filter (\(Decl g _ _ _) -> g `elem` fs) decls) scc
68 --      callGraph = map (\f -> (f,f,snub $ f:funCallsE e)) $
69 --                     grupBy $ map (\(Decl f _ _ _) -> f) decls
70       callGraph = 
71           let fMap  = Map.fromListWith union $ 
72                        map (\(Decl f _ _ e) -> (f,f:funCallsE e)) decls 
73               fMap' = Map.map (snub) fMap 
74           in map (\(f,fs) -> (f,f,fs)) $ Map.toList fMap'
75       funCallsE (EVar _ _ v)    = []
76       funCallsE (EFun _ _ f es) = f:concatMap funCallsE es 
77       funCallsE (ECon _ _ _ es) = concatMap funCallsE es 
78
79
80 inferenceStep decls tmap icount = 
81       do { (decls0,  (tmpMap, icount0)) <- runStateT (makeInitConstr tmap decls) ([],icount)
82          ; (decls' , (constr, icount')) <- runStateT (mapM (assignTypeVars tmpMap tmap) decls0) ([],icount0)
83          ; (tmpMap', etypeMap') <- solveConstr tmpMap constr
84          ; let decls'' = map (repl tmpMap' etypeMap') decls'
85          ; return (decls'', tmpMap' ++ tmap, icount') }
86         where 
87           repl tM cM (Decl f ftype ps e) =
88               Decl f (fromJust $ lookup f tM) (map replP ps) (replE e)
89               where
90                 replP (PVar id (TVar i) v)    
91                     = PVar id (fromJust $ lookup i cM) v
92                 replP (PCon id (TVar i) c ps)
93                     = PCon id (fromJust $ lookup i cM) c (map replP ps)
94                 replE (EVar id (TVar i) v)
95                     = EVar id (fromJust $ lookup i cM) v
96                 replE (ECon id (TVar i) c es)
97                     = ECon id (fromJust $ lookup i cM) c (map replE es)
98                 replE (EFun id (TVar i) c es)
99                     = EFun id (fromJust $ lookup i cM) c (map replE es)
100           extractConstr ds = map (\(Decl f t _ _) -> (f,t)) $
101                                 nubBy isSameFunc ds
102
103
104
105 solveConstr tmpMap constr 
106     = substStep constr (tmpMap, rearrange constr)
107     where 
108       introForAll (k,TFun _ ts t) =
109           let vs = snub $ varsT t ++ concatMap varsT ts 
110           in (k,TFun vs ts t)
111       rearrange constr = 
112           let vs = nub $ concatMap (\(t1,t2) -> varsT t1 ++ varsT t2) constr 
113           in map (\x -> (x,TVar x)) vs                
114       varsT (TVar i)    = [i]
115       varsT (TCon _ ts) = concatMap varsT ts 
116       varsT (TUndet)    = []
117       substStep [] (tM,cM) = return (map introForAll tM,cM)
118       substStep ((t,t'):cs) (tM,cM) =
119           do { subs <- unify t t'
120              ; substStep
121                   (performSubstC subs cs)
122                   (performSubstTM subs tM, performSubstCM subs cM) }
123       performSubstC subs cs
124           = map (\(t1,t2) -> (performSubstT subs t1, performSubstT subs t2)) cs
125       performSubstTM subs tM 
126           = map (\(k,v) -> (k, performSubstFT subs v)) tM
127       performSubstCM subs cM
128           = map (\(k,v) -> (k, performSubstT subs v)) cM
129       performSubstFT subs (TFun is ts t) 
130           = TFun [] (map (performSubstT subs) ts) (performSubstT subs t)
131       performSubstT subs (TUndet) = TUndet 
132       performSubstT subs (TVar i) = 
133           case lookup (TVar i) subs of 
134             Just t' -> t'
135             _       -> TVar i
136       performSubstT subs (TCon c ts) =
137           TCon c (map (performSubstT subs) ts)
138       unify :: Type -> Type -> Either String [ (Type, Type) ]
139       unify (TVar i) t | not (i `elem` varsT t) = return [ (TVar i, t) ]
140       unify t (TVar i) | not (i `elem` varsT t) = return [ (TVar i, t) ]
141       unify (TVar i) (TVar j) | i == j = return []
142       unify (TCon c ts) (TCon c' ts') | c == c' 
143           = do { ss <- mapM (uncurry unify) $ zip ts ts'
144                ; return $ concat ss }
145       unify t t' = throwError $ "Can't unify types: " ++ show ( ppr (t,t'))
146                  
147     
148                
149
150 makeInitConstr tmap decls =
151     do { mapM_ (\(Decl f t ps e) ->
152                       do { tmpMap <- getTmpMap 
153                          ; case t of
154                              FTUndet -> 
155                                  case lookup f tmpMap of 
156                                    Just t' -> 
157                                        return ()
158                                    _ -> 
159                                        do { i  <- newTypeVar 
160                                           ; is <- mapM (\_->newTypeVar) ps 
161                                           ; let t' = TFun [] (map TVar is) (TVar i) 
162                                           ; putTmpMap ((f,t'):tmpMap)
163                                           ; return ()  }
164                              _ -> 
165                                  putTmpMap ((f,t):tmpMap)}) $ 
166          (nubBy isSameFunc decls)
167        ; tmpMap <- getTmpMap
168        ; let decls' = map (\(Decl f t ps e) -> 
169                              Decl f (fromJust $ lookup f tmpMap) ps e) decls
170        ; return decls' }
171     where getTmpMap    = do { (tmpMap,i) <- get; return tmpMap }
172           putTmpMap tm = do { (_,i) <- get; put (tm,i) }
173           newTypeVar   = do { (tm,i) <- get; put (tm,i+1); return i}
174
175     
176                
177
178 assignTypeVars tmpMap typeMap (Decl fname ftype ps e) =
179     do ps' <- mapM assignTypeVarsP ps
180        e'  <- assignTypeVarsE      e
181        unifyFT ftype (TFun [] (map typeofP ps') (typeofE e'))
182        let vtp = concatMap vtMapP ps'
183        let vte = vtMapE e'
184        mapM_ (\(x,t) -> case (lookup x vte) of 
185                           { Just t' -> unifyT t t'; _ -> return ()}) vtp 
186        mapM_ (\(x,t) -> case (lookup x vte) of 
187                           { Just t' -> unifyT t t' }) vte 
188        return $ Decl fname ftype ps' e'
189     where
190       vtMapP (PVar _ t x)    = [(x,t)]
191       vtMapP (PCon _ _ c ps) = concatMap vtMapP ps 
192       vtMapE (EVar _ t x)    = [(x,t)]
193       vtMapE (ECon _ _ c es) = concatMap vtMapE es
194       vtMapE (EFun _ _ c es) = concatMap vtMapE es
195 --      newTypeVar :: State ( [(Type,Type)], Int ) Int
196       newTypeVar = do { (constr, icount) <- get
197                       ; put (constr, icount+1)
198                       ; return icount }
199       addConstr s t = do { (constr, icount) <- get
200                            ; put ((s,t):constr, icount) }
201       assignTypeVarsP (PVar id t v) = 
202           do { i <- newTypeVar
203              ; unifyT t (TVar i) 
204              ; return $ PVar id (TVar i) v } 
205       assignTypeVarsP (PCon id t c ps) = 
206           do { i <- newTypeVar
207              ; case lookup c typeMap of
208                  Just t' -> 
209                      do { ps' <- mapM assignTypeVarsP ps 
210                         ; unifyFT t' (TFun [] (map typeofP ps') (TVar i))
211                         ; unifyT  t  (TVar i)
212                         ; return $ PCon id (TVar i) c ps' }}
213       assignTypeVarsE (EVar id t v) = 
214           do { i <- newTypeVar 
215              ; unifyT t (TVar i)
216              ; return $ EVar id (TVar i) v }
217       assignTypeVarsE (ECon id t c es) =
218           do { i <- newTypeVar
219              ; case lookup c typeMap of
220                  Just t' -> 
221                      do { es' <- mapM assignTypeVarsE es 
222                         ; unifyFT t' (TFun [] (map typeofE es') (TVar i))
223                         ; unifyT  t  (TVar i)
224                         ; return $ ECon id (TVar i) c es' }
225                  Nothing -> fail $ "No type " ++ show c ++ " in type map"
226              }
227       assignTypeVarsE (EFun id t f es) =
228           do { i <- newTypeVar
229              ; case lookup f (typeMap ++ tmpMap)  of
230                  Just t' -> 
231                      do { es' <- mapM assignTypeVarsE es 
232                         ; unifyFT t' (TFun [] (map typeofE es') (TVar i))
233                         ; unifyT  t  (TVar i)
234                         ; return $ EFun id (TVar i) f es' }
235                  _ ->
236                      fail $ (show f ++ " is not in " ++ show (typeMap ++ tmpMap))
237              }
238 --      unifyT :: Type -> Type -> State ([(Type,Type)],Int) ()
239       unifyT (TUndet) _ = return ()
240       unifyT _ (TUndet) = return ()
241       unifyT (TVar i) (TVar j) | i == j = return ()
242       unifyT t t'       = addConstr t t'
243       unifyFT (FTUndet) _ = return ()
244       unifyFT _ (FTUndet) = return ()
245       unifyFT t t' = 
246           do { s  <- escapeForAll t 
247              ; s' <- escapeForAll t'
248              ; case (s,s') of 
249                  (TFun _ ts t, TFun _ ts' t') ->
250                      mapM_ (uncurry unifyT) $ zip (t:ts) (t':ts') }
251       escapeForAll (TFun is ts t) =
252           do { is' <- mapM (\_ -> newTypeVar) is 
253              ; let ts' = map (replaceTVar (zip is is')) ts
254              ; let t'   = replaceTVar (zip is is') t 
255              ; return $ TFun [] ts' t'}
256       replaceTVar table TUndet = TUndet
257       replaceTVar table (TVar i) =
258           case lookup i table of
259             Just j -> TVar j 
260             _      -> TVar i
261       replaceTVar table (TCon t ts) =
262           TCon t (map (replaceTVar table) ts)
263
264                      
265