git.iko.soy/elk
on e714001 help sign in
elk/src/Transform/SortDefinitions.hs
170 lines 5.6 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
17ctors :: P.CanonicalPattern -> [String]
18ctors pattern =
19 case pattern of
20 P.Var _ -> []
21 P.Alias _ p -> ctors p
22 P.Record _ -> []
23 P.Anything -> []
24 P.Literal _ -> []
25 P.Data (V.Canonical home name) ps ->
26 case home of
27 V.Local -> name : rest
28 V.BuiltIn -> rest
29 V.Module _ -> rest
30 where
31 rest = concatMap ctors ps
33free :: String -> State (Set.Set String) ()
34free x = modify (Set.insert x)
36freeIfLocal :: V.Canonical -> State (Set.Set String) ()
37freeIfLocal (V.Canonical home name) =
38 do case home of
39 V.Local -> free name
40 V.BuiltIn -> return ()
41 V.Module _ -> return ()
43bound :: Set.Set String -> State (Set.Set String) ()
44bound boundVars = modify (\freeVars -> Set.difference freeVars boundVars)
46sortDefs :: Canonical.Expr -> Canonical.Expr
47sortDefs expr = evalState (reorder expr) Set.empty
49reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
50reorder (A ann expr) =
51 A ann <$>
52 case expr of
53 -- Be careful adding and restricting freeVars
54 Var var -> freeIfLocal var >> return expr
56 Lambda p e ->
57 uncurry Lambda <$> bindingReorder (p,e)
59 Binop op e1 e2 ->
60 do freeIfLocal op
61 Binop op <$> reorder e1 <*> reorder e2
63 Case e cases ->
64 Case <$> reorder e <*> mapM bindingReorder cases
66 Data name es ->
67 do free name
68 Data name <$> mapM reorder es
70 -- Just pipe the reorder though
71 Literal _ -> return expr
73 Range e1 e2 ->
74 Range <$> reorder e1 <*> reorder e2
76 ExplicitList es ->
77 ExplicitList <$> mapM reorder es
79 App e1 e2 ->
80 App <$> reorder e1 <*> reorder e2
82 MultiIf branches ->
83 MultiIf <$> mapM (\(e1,e2) -> (,) <$> reorder e1 <*> reorder e2) branches
85 Access e lbl ->
86 Access <$> reorder e <*> return lbl
88 Remove e lbl ->
89 Remove <$> reorder e <*> return lbl
91 Insert e lbl v ->
92 Insert <$> reorder e <*> return lbl <*> reorder v
94 Modify e fields ->
95 Modify <$> reorder e <*> mapM (\(k,v) -> (,) k <$> reorder v) fields
97 Record fields ->
98 Record <$> mapM (\(k,v) -> (,) k <$> reorder v) fields
100 GLShader _ _ _ -> return expr
102 PortOut name st signal -> PortOut name st <$> reorder signal
104 PortIn name st -> return $ PortIn name st
106 -- Actually do some reordering
107 Let defs body ->
108 do body' <- reorder body
110 -- Sort defs into strongly connected components.This
111 -- allows the programmer to write definitions in whatever
112 -- order they please, we can still define things in order
113 -- and generalize polymorphic functions when appropriate.
114 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
115 let defss = map Graph.flattenSCC sccs
117 -- remove let-bound variables from the context
118 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
119 bound (P.boundVars pattern)
120 mapM free (ctors pattern)
122 let A _ let' = foldr (\ds bod -> A ann (Let ds bod)) body' defss
124 return let'
126bindingReorder :: (P.CanonicalPattern, Canonical.Expr)
127 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
128bindingReorder (pattern,expr) =
129 do expr' <- reorder expr
130 bound (P.boundVars pattern)
131 mapM_ free (ctors pattern)
132 return (pattern, expr')
135reorderAndGetDependencies :: Canonical.Def -> State (Set.Set String) (Canonical.Def, [String])
136reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
137 do globalFrees <- get
138 -- work in a fresh environment
139 put Set.empty
140 expr' <- reorder expr
141 localFrees <- get
142 -- merge with global frees
143 modify (Set.union globalFrees)
144 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
147-- This also reorders the all of the sub-expressions in the Def list.
148buildDefDict :: [Canonical.Def] -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
149buildDefDict defs =
150 do pdefsDeps <- mapM reorderAndGetDependencies defs
151 return $ realDeps (addKey pdefsDeps)
153 where
154 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
155 addKey = zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
157 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
158 variableToKey (Canonical.Definition pattern _ _, key, _) =
159 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
161 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
162 variableToKeyMap pdefsDeps =
163 Map.fromList (concatMap variableToKey pdefsDeps)
165 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
166 realDeps pdefsDeps = map convert pdefsDeps
167 where
168 varDict = variableToKeyMap pdefsDeps
169 convert (pdef, key, deps) =
170 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)