git.iko.soy/elk
on 69a4366 help sign in
elk/src/Transform/SortDefinitions.hs
192 lines 5.7 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Transform.SortDefinitions (sortDefs) where
4import Control.Monad.State
5import Control.Applicative ((<$>),(<*>))
6import qualified Data.Graph as Graph
7import qualified Data.Map as Map
8import qualified Data.Maybe as Maybe
9import qualified Data.Set as Set
11import AST.Annotation
12import AST.Expression.General (Expr'(..))
13import qualified AST.Expression.Canonical as Canonical
14import qualified AST.Pattern as P
15import qualified AST.Variable as V
18ctors :: P.CanonicalPattern -> [String]
19ctors pattern =
20 case pattern of
21 P.Var _ -> []
22 P.Alias _ p -> ctors p
23 P.Record _ -> []
24 P.Anything -> []
25 P.Literal _ -> []
26 P.Data (V.Canonical home name) ps ->
27 case home of
28 V.Local -> name : rest
29 V.BuiltIn -> rest
30 V.Module _ -> rest
31 where
32 rest = concatMap ctors ps
35free :: String -> State (Set.Set String) ()
36free x =
37 modify (Set.insert x)
40freeIfLocal :: V.Canonical -> State (Set.Set String) ()
41freeIfLocal (V.Canonical home name) =
42 case home of
43 V.Local -> free name
44 V.BuiltIn -> return ()
45 V.Module _ -> return ()
48bound :: Set.Set String -> State (Set.Set String) ()
49bound boundVars =
50 modify (\freeVars -> Set.difference freeVars boundVars)
53sortDefs :: Canonical.Expr -> Canonical.Expr
54sortDefs expr =
55 evalState (reorder expr) Set.empty
58reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
59reorder (A ann expr) =
60 A ann <$>
61 case expr of
62 -- Be careful adding and restricting freeVars
63 Var var ->
64 freeIfLocal var >> return expr
66 Lambda p e ->
67 uncurry Lambda <$> bindingReorder (p,e)
69 Binop op e1 e2 ->
70 do freeIfLocal op
71 Binop op <$> reorder e1 <*> reorder e2
73 Case e cases ->
74 Case <$> reorder e <*> mapM bindingReorder cases
76 Data name es ->
77 do free name
78 Data name <$> mapM reorder es
80 -- Just pipe the reorder though
81 Literal _ ->
82 return expr
84 Range e1 e2 ->
85 Range <$> reorder e1 <*> reorder e2
87 ExplicitList es ->
88 ExplicitList <$> mapM reorder es
90 App e1 e2 ->
91 App <$> reorder e1 <*> reorder e2
93 MultiIf branches ->
94 MultiIf <$> mapM (\(e1,e2) -> (,) <$> reorder e1 <*> reorder e2) branches
96 Access e lbl ->
97 Access <$> reorder e <*> return lbl
99 Remove e lbl ->
100 Remove <$> reorder e <*> return lbl
102 Insert e lbl v ->
103 Insert <$> reorder e <*> return lbl <*> reorder v
105 Modify e fields ->
106 Modify <$> reorder e <*> mapM (\(k,v) -> (,) k <$> reorder v) fields
108 Record fields ->
109 Record <$> mapM (\(k,v) -> (,) k <$> reorder v) fields
111 GLShader _ _ _ ->
112 return expr
114 PortOut name st signal ->
115 PortOut name st <$> reorder signal
117 PortIn name st ->
118 return $ PortIn name st
120 -- Actually do some reordering
121 Let defs body ->
122 do body' <- reorder body
124 -- Sort defs into strongly connected components.This
125 -- allows the programmer to write definitions in whatever
126 -- order they please, we can still define things in order
127 -- and generalize polymorphic functions when appropriate.
128 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
129 let defss = map Graph.flattenSCC sccs
131 -- remove let-bound variables from the context
132 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
133 bound (P.boundVars pattern)
134 mapM free (ctors pattern)
136 let A _ let' =
137 foldr (\ds bod -> A ann (Let ds bod)) body' defss
139 return let'
142bindingReorder
143 :: (P.CanonicalPattern, Canonical.Expr)
144 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
145bindingReorder (pattern,expr) =
146 do expr' <- reorder expr
147 bound (P.boundVars pattern)
148 mapM_ free (ctors pattern)
149 return (pattern, expr')
152reorderAndGetDependencies
153 :: Canonical.Def
154 -> State (Set.Set String) (Canonical.Def, [String])
155reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
156 do globalFrees <- get
157 -- work in a fresh environment
158 put Set.empty
159 expr' <- reorder expr
160 localFrees <- get
161 -- merge with global frees
162 modify (Set.union globalFrees)
163 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
166-- This also reorders the all of the sub-expressions in the Def list.
167buildDefDict
168 :: [Canonical.Def]
169 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
170buildDefDict defs =
171 do pdefsDeps <- mapM reorderAndGetDependencies defs
172 return $ realDeps (addKey pdefsDeps)
173 where
174 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
175 addKey =
176 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
178 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
179 variableToKey (Canonical.Definition pattern _ _, key, _) =
180 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
182 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
183 variableToKeyMap pdefsDeps =
184 Map.fromList (concatMap variableToKey pdefsDeps)
186 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
187 realDeps pdefsDeps =
188 map convert pdefsDeps
189 where
190 varDict = variableToKeyMap pdefsDeps
191 convert (pdef, key, deps) =
192 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)