git.iko.soy/elk
on df41ae8 help sign in
elk/src/Canonicalize/Sort.hs
226 lines 5.8 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Canonicalize.Sort (definitions) where
4import Control.Monad.State
5import qualified Data.Graph as Graph
6import qualified Data.Map as Map
7import qualified Data.Maybe as Maybe
8import qualified Data.Set as Set
10import AST.Expression.General (Expr'(..))
11import qualified AST.Expression.Canonical as Canonical
12import qualified AST.Pattern as P
13import qualified AST.Variable as V
14import qualified Reporting.Annotation as A
18-- STARTING POINT
21definitions :: Canonical.Expr -> Canonical.Expr
22definitions expression =
23 evalState (reorder expression) Set.empty
26free :: String -> State (Set.Set String) ()
27free x =
28 modify (Set.insert x)
31freeIfLocal :: V.Canonical -> State (Set.Set String) ()
32freeIfLocal (V.Canonical home name) =
33 case home of
34 V.Local -> free name
35 V.TopLevel _ -> free name
36 V.BuiltIn -> return ()
37 V.Module _ -> return ()
41-- REORDER EXPRESSIONS
44reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
45reorder (A.A ann expression) =
46 A.A ann <$>
47 case expression of
48 -- Be careful adding and restricting freeVars
49 Var var ->
50 do freeIfLocal var
51 return expression
53 Lambda pattern body ->
54 uncurry Lambda <$> bindingReorder (pattern,body)
56 Binop op leftExpr rightExpr ->
57 do freeIfLocal op
58 Binop op <$> reorder leftExpr <*> reorder rightExpr
60 Case expr cases ->
61 Case <$> reorder expr <*> mapM bindingReorder cases
63 Data name exprs ->
64 do free name
65 Data name <$> mapM reorder exprs
67 -- Just pipe the reorder though
68 Literal _ ->
69 return expression
71 Range lowExpr highExpr ->
72 Range <$> reorder lowExpr <*> reorder highExpr
74 ExplicitList es ->
75 ExplicitList <$> mapM reorder es
77 App func arg ->
78 App <$> reorder func <*> reorder arg
80 If branches finally ->
81 If
82 <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
83 <*> reorder finally
85 Access record field ->
86 Access
87 <$> reorder record
88 <*> return field
90 Update record fields ->
91 Update
92 <$> reorder record
93 <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
95 Record fields ->
96 Record
97 <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields
99 Cmd _ _ ->
100 return expression
102 Sub _ _ ->
103 return expression
105 SaveEnv _ _ ->
106 return expression
108 GLShader _ _ _ ->
109 return expression
111 -- Actually do some reordering
112 Let defs body ->
113 do body' <- reorder body
115 -- Sort defs into strongly connected components.This
116 -- allows the programmer to write definitions in whatever
117 -- order they please, we can still define things in order
118 -- and generalize polymorphic functions when appropriate.
119 sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
120 let defss = map Graph.flattenSCC sccs
122 -- remove let-bound variables from the context
123 forM_ defs $ \(Canonical.Definition _ pattern _ _) -> do
124 bound pattern
125 mapM free (ctors pattern)
127 let (A.A _ let') =
128 foldr (\ds bod -> A.A ann (Let ds bod)) body' defss
130 return let'
133ctors :: P.Canonical -> [String]
134ctors (A.A _ pattern) =
135 case pattern of
136 P.Var _ ->
137 []
139 P.Alias _ p ->
140 ctors p
142 P.Record _ ->
143 []
145 P.Anything ->
146 []
148 P.Literal _ ->
149 []
151 P.Data (V.Canonical home name) ps ->
152 case home of
153 V.Local -> name : rest
154 V.TopLevel _ -> name : rest
155 V.BuiltIn -> rest
156 V.Module _ -> rest
157 where
158 rest = concatMap ctors ps
161bound :: P.Canonical -> State (Set.Set String) ()
162bound pattern =
163 let
164 boundVars =
165 P.boundVarSet pattern
166 in
167 modify (\freeVars -> Set.difference freeVars boundVars)
170bindingReorder
171 :: (P.Canonical, Canonical.Expr)
172 -> State (Set.Set String) (P.Canonical, Canonical.Expr)
173bindingReorder (pattern, expr) =
174 do expr' <- reorder expr
175 bound pattern
176 mapM_ free (ctors pattern)
177 return (pattern, expr')
181-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS
184-- This also reorders the all of the sub-expressions in the Def list.
185buildDefGraph
186 :: [Canonical.Def]
187 -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
188buildDefGraph defs =
189 do pdefsDeps <- mapM reorderAndGetDependencies defs
190 return $ realDeps (addKey pdefsDeps)
191 where
192 addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
193 addKey =
194 zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]
196 variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
197 variableToKey (Canonical.Definition _ pattern _ _, key, _) =
198 [ (var, key) | var <- P.boundVarList pattern ]
200 variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
201 variableToKeyMap pdefsDeps =
202 Map.fromList (concatMap variableToKey pdefsDeps)
204 realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
205 realDeps pdefsDeps =
206 map convert pdefsDeps
207 where
208 varDict =
209 variableToKeyMap pdefsDeps
211 convert (pdef, key, deps) =
212 (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)
215reorderAndGetDependencies
216 :: Canonical.Def
217 -> State (Set.Set String) (Canonical.Def, [String])
218reorderAndGetDependencies (Canonical.Definition facts pattern expr mType) =
219 do globalFrees <- get
220 -- work in a fresh environment
221 put Set.empty
222 expr' <- reorder expr
223 localFrees <- get
224 -- merge with global frees
225 modify (Set.union globalFrees)
226 return (Canonical.Definition facts pattern expr' mType, Set.toList localFrees)