git.iko.soy/elk
on 54dfb21 help sign in
elk/src/Optimize/Inline.hs
291 lines 7.0 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Inline (inline) where
4import Control.Arrow (second, (***))
5import Control.Monad (foldM)
6import qualified Data.List as List
7import qualified Data.Map as Map
9import AST.Expression.General (PortImpl(..))
10import AST.Expression.Optimized (Expr(..), Decider(..), Choice(..))
11import qualified AST.Expression.Optimized as Opt
12import qualified AST.Variable as Var
13import qualified Optimize.Environment as Env
16inline :: Map.Map String Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
17inline rawSubs expression =
18 let
19 counts =
20 count expression
22 annotatedSubs =
23 Map.toList (Map.intersectionWith (,) counts rawSubs)
24 in
25 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
26 let inlinedExpr = replace subs expression
27 return $
28 if null defs then
29 inlinedExpr
31 else
32 Let (map (\(name, expr) -> Opt.Def Opt.dummyFacts name expr) defs) inlinedExpr
35data Subs = Subs
36 { _substitutions :: Map.Map String Opt.Expr
37 , _definitions :: [(String, Opt.Expr)]
38 }
41processSubs :: Subs -> (String, (Int, Opt.Expr)) -> Env.Optimizer Subs
42processSubs (Subs subs defs) (name, (n, expr)) =
43 if n == 1 then
44 return $ Subs (Map.insert name expr subs) defs
46 else
47 do uniqueName <- Env.freshName
48 return $ Subs
49 (Map.insert name (Opt.Var (Var.local uniqueName)) subs)
50 ((uniqueName, expr) : defs)
52-- HELPERS
54deleteBatch :: [String] -> Map.Map String a -> Map.Map String a
55deleteBatch names dict =
56 List.foldl' (flip Map.delete) dict names
59getDefName :: Opt.Def -> String
60getDefName def =
61 case def of
62 Opt.Def _ name _ ->
63 name
65 Opt.TailDef _ name _ _ ->
66 name
69-- COUNT VARIABLE APPEARANCES
71count :: Opt.Expr -> Map.Map String Int
72count expression =
73 let
74 count2 a b =
75 Map.unionWith (+) (count a) (count b)
77 countMany xs =
78 Map.unionsWith (+) (map count xs)
79 in
80 case expression of
81 Literal _ ->
82 Map.empty
84 Var (Var.Canonical home name) ->
85 case home of
86 Var.Local ->
87 Map.singleton name 1
89 Var.BuiltIn ->
90 Map.empty
92 Var.Module _ ->
93 Map.empty
95 Var.TopLevel _ ->
96 Map.empty
98 Range lo hi ->
99 count2 lo hi
101 ExplicitList exprs ->
102 countMany exprs
104 Binop _op left right ->
105 count2 left right
107 Function args body ->
108 deleteBatch args (count body)
110 Call func args ->
111 Map.unionWith (+) (count func) (countMany args)
113 TailCall _name _argNames args ->
114 countMany args
116 If branches finally ->
117 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
119 Let defs expr ->
120 let
121 countDef def =
122 case def of
123 Opt.Def _ name body ->
124 Map.delete name (count body)
126 Opt.TailDef _ name argNames body ->
127 deleteBatch (name:argNames) (count body)
129 exprCount =
130 deleteBatch (map getDefName defs) (count expr)
131 in
132 Map.unionsWith (+) (exprCount : map countDef defs)
134 Case _ decider jumps ->
135 let
136 countDecider dcdr =
137 case dcdr of
138 Leaf (Inline expr) ->
139 count expr
141 Leaf (Jump _) ->
142 Map.empty
144 Chain _ success failure ->
145 Map.unionWith (+) (countDecider success) (countDecider failure)
147 FanOut _path tests fallback ->
148 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
149 in
150 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
152 Data _tag values ->
153 countMany values
155 DataAccess root _index ->
156 count root
158 Access record _field ->
159 count record
161 Update record fields ->
162 Map.unionWith (+) (count record) (countMany (map snd fields))
164 Record fields ->
165 countMany (map snd fields)
167 Port impl ->
168 case impl of
169 In _ _ ->
170 Map.empty
172 Out _name expr _tipe ->
173 count expr
175 Task _name expr _tipe ->
176 count expr
178 GLShader _ _ _ ->
179 Map.empty
181 Crash _ _ ->
182 Map.empty
185-- REPLACE
187replace :: Map.Map String Opt.Expr -> Opt.Expr -> Opt.Expr
188replace substitutions expression =
189 let
190 go = replace substitutions
191 in
192 case expression of
193 Literal _ ->
194 expression
196 Var (Var.Canonical Var.Local name) ->
197 maybe expression id (Map.lookup name substitutions)
199 Var _ ->
200 expression
202 Range lo hi ->
203 Range (go lo) (go hi)
205 ExplicitList exprs ->
206 ExplicitList (map go exprs)
208 Binop op left right ->
209 Binop op (go left) (go right)
211 Function args body ->
212 Function args (replace (deleteBatch args substitutions) body)
214 Call func args ->
215 Call (go func) (map go args)
217 TailCall name argNames args ->
218 TailCall name argNames (map go args)
220 If branches finally ->
221 If (map (go *** go) branches) (go finally)
223 Let defs expr ->
224 let
225 replaceDef def =
226 case def of
227 Opt.Def facts name body ->
228 Opt.Def facts name
229 (replace (Map.delete name substitutions) body)
231 Opt.TailDef facts name argNames body ->
232 Opt.TailDef facts name argNames
233 (replace (deleteBatch (name:argNames) substitutions) body)
235 boundNames =
236 map getDefName defs
237 in
238 Let
239 (map replaceDef defs)
240 (replace (deleteBatch boundNames substitutions) expr)
242 Case exprName decider jumps ->
243 let
244 goDecider dcdr =
245 case dcdr of
246 Leaf (Inline expr) ->
247 Leaf (Inline (go expr))
249 Leaf (Jump _target) ->
250 dcdr
252 Chain chain success failure ->
253 Chain chain (goDecider success) (goDecider failure)
255 FanOut path tests fallback ->
256 FanOut path (map (second goDecider) tests) (goDecider fallback)
257 in
258 Case exprName (goDecider decider) (map (second go) jumps)
260 Data tag values ->
261 Data tag (map go values)
263 DataAccess root index ->
264 DataAccess (go root) index
266 Access record field ->
267 Access (go record) field
269 Update record fields ->
270 Update (go record) (map (second go) fields)
272 Record fields ->
273 Record (map (second go) fields)
275 Port impl ->
276 Port $
277 case impl of
278 In _ _ ->
279 impl
281 Out name expr tipe ->
282 Out name (go expr) tipe
284 Task name expr tipe ->
285 Task name (go expr) tipe
287 GLShader _ _ _ ->
288 expression
290 Crash _ _ ->
291 expression