git.iko.soy/elk
on 2a6e92b help sign in
elk/src/Transform/SortDefinitions.hs
196 lines 5.8 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 Input name st ->
115 return $ Input name st
117 Output name st signal ->
118 Output name st <$> reorder signal
120 Loopback name st maybeExpr ->
121 Loopback name st <$>
122 maybe (return Nothing) (fmap Just . reorder) maybeExpr
124 -- Actually do some reordering
125 Let defs body ->
126 do body' <- reorder body
128 -- Sort defs into strongly connected components.This
129 -- allows the programmer to write definitions in whatever
130 -- order they please, we can still define things in order
131 -- and generalize polymorphic functions when appropriate.
132 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
133 let defss = map Graph.flattenSCC sccs
135 -- remove let-bound variables from the context
136 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
137 bound (P.boundVars pattern)
138 mapM free (ctors pattern)
140 let A _ let' =
141 foldr (\ds bod -> A ann (Let ds bod)) body' defss
143 return let'
146bindingReorder
147 :: (P.CanonicalPattern, Canonical.Expr)
148 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
149bindingReorder (pattern,expr) =
150 do expr' <- reorder expr
151 bound (P.boundVars pattern)
152 mapM_ free (ctors pattern)
153 return (pattern, expr')
156reorderAndGetDependencies
157 :: Canonical.Def
158 -> State (Set.Set String) (Canonical.Def, [String])
159reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
160 do globalFrees <- get
161 -- work in a fresh environment
162 put Set.empty
163 expr' <- reorder expr
164 localFrees <- get
165 -- merge with global frees
166 modify (Set.union globalFrees)
167 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
170-- This also reorders the all of the sub-expressions in the Def list.
171buildDefDict
172 :: [Canonical.Def]
173 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
174buildDefDict defs =
175 do pdefsDeps <- mapM reorderAndGetDependencies defs
176 return $ realDeps (addKey pdefsDeps)
177 where
178 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
179 addKey =
180 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
182 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
183 variableToKey (Canonical.Definition pattern _ _, key, _) =
184 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
186 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
187 variableToKeyMap pdefsDeps =
188 Map.fromList (concatMap variableToKey pdefsDeps)
190 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
191 realDeps pdefsDeps =
192 map convert pdefsDeps
193 where
194 varDict = variableToKeyMap pdefsDeps
195 convert (pdef, key, deps) =
196 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)