git.iko.soy/elk
on 3925082 help sign in
elk/compiler/Transform/SortDefinitions.hs
172 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
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 Markdown uid md es -> Markdown uid md <$> mapM reorder es
102 GLShader _ _ _ -> return expr
104 PortOut name st signal -> PortOut name st <$> reorder signal
106 PortIn name st -> return $ PortIn name st
108 -- Actually do some reordering
109 Let defs body ->
110 do body' <- reorder body
112 -- Sort defs into strongly connected components.This
113 -- allows the programmer to write definitions in whatever
114 -- order they please, we can still define things in order
115 -- and generalize polymorphic functions when appropriate.
116 sccs <- Graph.stronglyConnComp <$> buildDefDict defs
117 let defss = map Graph.flattenSCC sccs
119 -- remove let-bound variables from the context
120 forM_ defs $ \(Canonical.Definition pattern _ _) -> do
121 bound (P.boundVars pattern)
122 mapM free (ctors pattern)
124 let A _ let' = foldr (\ds bod -> A ann (Let ds bod)) body' defss
126 return let'
128bindingReorder :: (P.CanonicalPattern, Canonical.Expr)
129 -> State (Set.Set String) (P.CanonicalPattern, Canonical.Expr)
130bindingReorder (pattern,expr) =
131 do expr' <- reorder expr
132 bound (P.boundVars pattern)
133 mapM_ free (ctors pattern)
134 return (pattern, expr')
137reorderAndGetDependencies :: Canonical.Def -> State (Set.Set String) (Canonical.Def, [String])
138reorderAndGetDependencies (Canonical.Definition pattern expr mType) =
139 do globalFrees <- get
140 -- work in a fresh environment
141 put Set.empty
142 expr' <- reorder expr
143 localFrees <- get
144 -- merge with global frees
145 modify (Set.union globalFrees)
146 return (Canonical.Definition pattern expr' mType, Set.toList localFrees)
149-- This also reorders the all of the sub-expressions in the Def list.
150buildDefDict :: [Canonical.Def] -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
151buildDefDict defs =
152 do pdefsDeps <- mapM reorderAndGetDependencies defs
153 return $ realDeps (addKey pdefsDeps)
155 where
156 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
157 addKey = zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
159 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
160 variableToKey (Canonical.Definition pattern _ _, key, _) =
161 [ (var, key) | var <- Set.toList (P.boundVars pattern) ]
163 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
164 variableToKeyMap pdefsDeps =
165 Map.fromList (concatMap variableToKey pdefsDeps)
167 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
168 realDeps pdefsDeps = map convert pdefsDeps
169 where
170 varDict = variableToKeyMap pdefsDeps
171 convert (pdef, key, deps) =
172 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)