git.iko.soy/elk
on 267f708 help sign in
elk/src/Optimize/Inline.hs
299 lines 6.9 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.Optimized (Expr(..), Decider(..), Choice(..))
10import qualified AST.Expression.Optimized as Opt
11import qualified AST.Variable as Var
12import 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)
54-- HELPERS
57deleteBatch :: [String] -> Map.Map String a -> Map.Map String a
58deleteBatch names dict =
59 List.foldl' (flip Map.delete) dict names
62getDefName :: Opt.Def -> String
63getDefName def =
64 case def of
65 Opt.Def _ name _ ->
66 name
68 Opt.TailDef _ name _ _ ->
69 name
73-- COUNT VARIABLE APPEARANCES
76count :: Opt.Expr -> Map.Map String Int
77count expression =
78 let
79 count2 a b =
80 Map.unionWith (+) (count a) (count b)
82 countMany xs =
83 Map.unionsWith (+) (map count xs)
84 in
85 case expression of
86 Literal _ ->
87 Map.empty
89 Var (Var.Canonical home name) ->
90 case home of
91 Var.Local ->
92 Map.singleton name 1
94 Var.BuiltIn ->
95 Map.empty
97 Var.Module _ ->
98 Map.empty
100 Var.TopLevel _ ->
101 Map.empty
103 ExplicitList exprs ->
104 countMany exprs
106 Binop _op left right ->
107 count2 left right
109 Function args body ->
110 deleteBatch args (count body)
112 Call func args ->
113 Map.unionWith (+) (count func) (countMany args)
115 TailCall _name _argNames args ->
116 countMany args
118 If branches finally ->
119 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
121 Let defs expr ->
122 let
123 countDef def =
124 case def of
125 Opt.Def _ name body ->
126 Map.delete name (count body)
128 Opt.TailDef _ name argNames body ->
129 deleteBatch (name:argNames) (count body)
131 exprCount =
132 deleteBatch (map getDefName defs) (count expr)
133 in
134 Map.unionsWith (+) (exprCount : map countDef defs)
136 Case _ decider jumps ->
137 let
138 countDecider dcdr =
139 case dcdr of
140 Leaf (Inline expr) ->
141 count expr
143 Leaf (Jump _) ->
144 Map.empty
146 Chain _ success failure ->
147 Map.unionWith (+) (countDecider success) (countDecider failure)
149 FanOut _path tests fallback ->
150 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
151 in
152 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
154 Data _tag values ->
155 countMany values
157 DataAccess root _index ->
158 count root
160 Access record _field ->
161 count record
163 Update record fields ->
164 Map.unionWith (+) (count record) (countMany (map snd fields))
166 Record fields ->
167 countMany (map snd fields)
169 Cmd _ ->
170 Map.empty
172 Sub _ ->
173 Map.empty
175 OutgoingPort _ _ ->
176 Map.empty
178 IncomingPort _ _ ->
179 Map.empty
181 Program _ expr ->
182 count expr
184 GLShader _ _ _ ->
185 Map.empty
187 Crash _ _ maybeBranchProblem ->
188 maybe Map.empty count maybeBranchProblem
192-- REPLACE
195replace :: Map.Map String Opt.Expr -> Opt.Expr -> Opt.Expr
196replace substitutions expression =
197 let
198 go = replace substitutions
199 in
200 case expression of
201 Literal _ ->
202 expression
204 Var (Var.Canonical Var.Local name) ->
205 maybe expression id (Map.lookup name substitutions)
207 Var _ ->
208 expression
210 ExplicitList exprs ->
211 ExplicitList (map go exprs)
213 Binop op left right ->
214 Binop op (go left) (go right)
216 Function args body ->
217 Function args (replace (deleteBatch args substitutions) body)
219 Call func args ->
220 Call (go func) (map go args)
222 TailCall name argNames args ->
223 TailCall name argNames (map go args)
225 If branches finally ->
226 If (map (go *** go) branches) (go finally)
228 Let defs expr ->
229 let
230 replaceDef def =
231 case def of
232 Opt.Def facts name body ->
233 Opt.Def facts name
234 (replace (Map.delete name substitutions) body)
236 Opt.TailDef facts name argNames body ->
237 Opt.TailDef facts name argNames
238 (replace (deleteBatch (name:argNames) substitutions) body)
240 boundNames =
241 map getDefName defs
242 in
243 Let
244 (map replaceDef defs)
245 (replace (deleteBatch boundNames substitutions) expr)
247 Case exprName decider jumps ->
248 let
249 goDecider dcdr =
250 case dcdr of
251 Leaf (Inline expr) ->
252 Leaf (Inline (go expr))
254 Leaf (Jump _target) ->
255 dcdr
257 Chain chain success failure ->
258 Chain chain (goDecider success) (goDecider failure)
260 FanOut path tests fallback ->
261 FanOut path (map (second goDecider) tests) (goDecider fallback)
262 in
263 Case exprName (goDecider decider) (map (second go) jumps)
265 Data tag values ->
266 Data tag (map go values)
268 DataAccess root index ->
269 DataAccess (go root) index
271 Access record field ->
272 Access (go record) field
274 Update record fields ->
275 Update (go record) (map (second go) fields)
277 Record fields ->
278 Record (map (second go) fields)
280 Cmd _ ->
281 expression
283 Sub _ ->
284 expression
286 OutgoingPort _ _ ->
287 expression
289 IncomingPort _ _ ->
290 expression
292 Program kind expr ->
293 Program kind (go expr)
295 GLShader _ _ _ ->
296 expression
298 Crash home region maybeBranchProblem ->
299 Crash home region (fmap go maybeBranchProblem)