git.iko.soy/elk
on ea77dcf help sign in
elk/src/Optimize/Inline.hs
283 lines 6.5 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Inline (inline) where
5import Control.Arrow (second, (***))
6import Control.Monad (foldM)
7import qualified Data.List as List
8import qualified Data.Map as Map
9import Data.Text (Text)
11import AST.Expression.Optimized (Expr(..), Decider(..), Choice(..))
12import qualified AST.Expression.Optimized as Opt
13import qualified Optimize.Environment as Env
17inline :: Map.Map Text Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
18inline rawSubs expression =
19 let
20 counts =
21 count expression
23 annotatedSubs =
24 Map.toList (Map.intersectionWith (,) counts rawSubs)
25 in
26 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
27 let inlinedExpr = replace subs expression
28 return $
29 if null defs then
30 inlinedExpr
32 else
33 Let (map (second Opt.Def) defs) inlinedExpr
36data Subs =
37 Subs
38 { _substitutions :: Map.Map Text Opt.Expr
39 , _definitions :: [(Text, Opt.Expr)]
40 }
43processSubs :: Subs -> (Text, (Int, Opt.Expr)) -> Env.Optimizer Subs
44processSubs (Subs subs defs) (name, (n, expr)) =
45 if n == 1 then
46 return $ Subs (Map.insert name expr subs) defs
48 else
49 do uniqueName <- Env.freshName
50 return $ Subs
51 (Map.insert name (Opt.VarLocal uniqueName) subs)
52 ((uniqueName, expr) : defs)
56-- HELPERS
59deleteBatch :: [Text] -> Map.Map Text a -> Map.Map Text a
60deleteBatch names dict =
61 List.foldl' (flip Map.delete) dict names
65-- COUNT VARIABLE APPEARANCES
68count :: Opt.Expr -> Map.Map Text Int
69count expression =
70 let
71 count2 a b =
72 Map.unionWith (+) (count a) (count b)
74 countMany xs =
75 Map.unionsWith (+) (map count xs)
76 in
77 case expression of
78 Literal _ ->
79 Map.empty
81 VarLocal name ->
82 Map.singleton name 1
84 VarGlobal _ ->
85 Map.empty
87 List exprs ->
88 countMany exprs
90 Binop _ left right ->
91 count2 left right
93 Function args body ->
94 deleteBatch args (count body)
96 Call func args ->
97 Map.unionWith (+) (count func) (countMany args)
99 TailCall _name _argNames args ->
100 countMany args
102 If branches finally ->
103 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
105 Let defs expr ->
106 let
107 countDef (name, def) =
108 case def of
109 Opt.Def body ->
110 Map.delete name (count body)
112 Opt.TailDef argNames body ->
113 deleteBatch (name:argNames) (count body)
115 exprCount =
116 deleteBatch (map fst defs) (count expr)
117 in
118 Map.unionsWith (+) (exprCount : map countDef defs)
120 Case _ decider jumps ->
121 let
122 countDecider dcdr =
123 case dcdr of
124 Leaf (Inline expr) ->
125 count expr
127 Leaf (Jump _) ->
128 Map.empty
130 Chain _ success failure ->
131 Map.unionWith (+) (countDecider success) (countDecider failure)
133 FanOut _path tests fallback ->
134 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
135 in
136 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
138 Ctor _tag values ->
139 countMany values
141 CtorAccess root _index ->
142 count root
144 Access record _field ->
145 count record
147 Update record fields ->
148 Map.unionWith (+) (count record) (countMany (map snd fields))
150 Record fields ->
151 countMany (map snd fields)
153 Cmd _ _ ->
154 Map.empty
156 Sub _ _ ->
157 Map.empty
159 OutgoingPort _ _ ->
160 Map.empty
162 IncomingPort _ _ ->
163 Map.empty
165 Program _ expr ->
166 count expr
168 GLShader _ ->
169 Map.empty
171 Crash _ _ maybeBranchProblem ->
172 maybe Map.empty count maybeBranchProblem
176-- REPLACE
179replace :: Map.Map Text Opt.Expr -> Opt.Expr -> Opt.Expr
180replace substitutions expression =
181 let
182 go = replace substitutions
183 in
184 case expression of
185 Literal _ ->
186 expression
188 VarLocal name ->
189 maybe expression id (Map.lookup name substitutions)
191 VarGlobal _ ->
192 expression
194 List exprs ->
195 List (map go exprs)
197 Binop op left right ->
198 Binop op (go left) (go right)
200 Function args body ->
201 Function args (replace (deleteBatch args substitutions) body)
203 Call func args ->
204 Call (go func) (map go args)
206 TailCall name argNames args ->
207 TailCall name argNames (map go args)
209 If branches finally ->
210 If (map (go *** go) branches) (go finally)
212 Let defs expr ->
213 let
214 replaceDef (name, def) =
215 case def of
216 Opt.Def body ->
217 (,) name $ Opt.Def $
218 replace (Map.delete name substitutions) body
220 Opt.TailDef argNames body ->
221 (,) name $ Opt.TailDef argNames $
222 replace (deleteBatch (name:argNames) substitutions) body
224 boundNames =
225 map fst defs
226 in
227 Let
228 (map replaceDef defs)
229 (replace (deleteBatch boundNames substitutions) expr)
231 Case exprName decider jumps ->
232 let
233 goDecider dcdr =
234 case dcdr of
235 Leaf (Inline expr) ->
236 Leaf (Inline (go expr))
238 Leaf (Jump _target) ->
239 dcdr
241 Chain chain success failure ->
242 Chain chain (goDecider success) (goDecider failure)
244 FanOut path tests fallback ->
245 FanOut path (map (second goDecider) tests) (goDecider fallback)
246 in
247 Case exprName (goDecider decider) (map (second go) jumps)
249 Ctor tag values ->
250 Ctor tag (map go values)
252 CtorAccess root index ->
253 CtorAccess (go root) index
255 Access record field ->
256 Access (go record) field
258 Update record fields ->
259 Update (go record) (map (second go) fields)
261 Record fields ->
262 Record (map (second go) fields)
264 Cmd _ _ ->
265 expression
267 Sub _ _ ->
268 expression
270 OutgoingPort _ _ ->
271 expression
273 IncomingPort _ _ ->
274 expression
276 Program kind expr ->
277 Program kind (go expr)
279 GLShader _ ->
280 expression
282 Crash home region maybeBranchProblem ->
283 Crash home region (fmap go maybeBranchProblem)