git.iko.soy/elk
on 1925e31 help sign in
elk/src/Optimize/Inline.hs
301 lines 7.0 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 AST.Variable as Var
14import qualified Optimize.Environment as Env
18inline :: Map.Map Text Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
19inline rawSubs expression =
20 let
21 counts =
22 count expression
24 annotatedSubs =
25 Map.toList (Map.intersectionWith (,) counts rawSubs)
26 in
27 do (Subs subs defs) <- foldM processSubs (Subs Map.empty []) annotatedSubs
28 let inlinedExpr = replace subs expression
29 return $
30 if null defs then
31 inlinedExpr
33 else
34 Let (map (\(name, expr) -> Opt.Def Opt.dummyFacts name expr) defs) inlinedExpr
37data Subs = 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.Var (Var.local 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
64getDefName :: Opt.Def -> Text
65getDefName def =
66 case def of
67 Opt.Def _ name _ ->
68 name
70 Opt.TailDef _ name _ _ ->
71 name
75-- COUNT VARIABLE APPEARANCES
78count :: Opt.Expr -> Map.Map Text Int
79count expression =
80 let
81 count2 a b =
82 Map.unionWith (+) (count a) (count b)
84 countMany xs =
85 Map.unionsWith (+) (map count xs)
86 in
87 case expression of
88 Literal _ ->
89 Map.empty
91 Var (Var.Canonical home name) ->
92 case home of
93 Var.Local ->
94 Map.singleton name 1
96 Var.BuiltIn ->
97 Map.empty
99 Var.Module _ ->
100 Map.empty
102 Var.TopLevel _ ->
103 Map.empty
105 List exprs ->
106 countMany exprs
108 Binop _op left right ->
109 count2 left right
111 Function args body ->
112 deleteBatch args (count body)
114 Call func args ->
115 Map.unionWith (+) (count func) (countMany args)
117 TailCall _name _argNames args ->
118 countMany args
120 If branches finally ->
121 Map.unionsWith (+) (count finally : map (uncurry count2) branches)
123 Let defs expr ->
124 let
125 countDef def =
126 case def of
127 Opt.Def _ name body ->
128 Map.delete name (count body)
130 Opt.TailDef _ name argNames body ->
131 deleteBatch (name:argNames) (count body)
133 exprCount =
134 deleteBatch (map getDefName defs) (count expr)
135 in
136 Map.unionsWith (+) (exprCount : map countDef defs)
138 Case _ decider jumps ->
139 let
140 countDecider dcdr =
141 case dcdr of
142 Leaf (Inline expr) ->
143 count expr
145 Leaf (Jump _) ->
146 Map.empty
148 Chain _ success failure ->
149 Map.unionWith (+) (countDecider success) (countDecider failure)
151 FanOut _path tests fallback ->
152 Map.unionsWith (+) (map countDecider (fallback : map snd tests))
153 in
154 Map.unionsWith (+) (countDecider decider : map (count . snd) jumps)
156 Ctor _tag values ->
157 countMany values
159 CtorAccess root _index ->
160 count root
162 Access record _field ->
163 count record
165 Update record fields ->
166 Map.unionWith (+) (count record) (countMany (map snd fields))
168 Record fields ->
169 countMany (map snd fields)
171 Cmd _ ->
172 Map.empty
174 Sub _ ->
175 Map.empty
177 OutgoingPort _ _ ->
178 Map.empty
180 IncomingPort _ _ ->
181 Map.empty
183 Program _ expr ->
184 count expr
186 GLShader _ _ _ ->
187 Map.empty
189 Crash _ _ maybeBranchProblem ->
190 maybe Map.empty count maybeBranchProblem
194-- REPLACE
197replace :: Map.Map Text Opt.Expr -> Opt.Expr -> Opt.Expr
198replace substitutions expression =
199 let
200 go = replace substitutions
201 in
202 case expression of
203 Literal _ ->
204 expression
206 Var (Var.Canonical Var.Local name) ->
207 maybe expression id (Map.lookup name substitutions)
209 Var _ ->
210 expression
212 List exprs ->
213 List (map go exprs)
215 Binop op left right ->
216 Binop op (go left) (go right)
218 Function args body ->
219 Function args (replace (deleteBatch args substitutions) body)
221 Call func args ->
222 Call (go func) (map go args)
224 TailCall name argNames args ->
225 TailCall name argNames (map go args)
227 If branches finally ->
228 If (map (go *** go) branches) (go finally)
230 Let defs expr ->
231 let
232 replaceDef def =
233 case def of
234 Opt.Def facts name body ->
235 Opt.Def facts name
236 (replace (Map.delete name substitutions) body)
238 Opt.TailDef facts name argNames body ->
239 Opt.TailDef facts name argNames
240 (replace (deleteBatch (name:argNames) substitutions) body)
242 boundNames =
243 map getDefName defs
244 in
245 Let
246 (map replaceDef defs)
247 (replace (deleteBatch boundNames substitutions) expr)
249 Case exprName decider jumps ->
250 let
251 goDecider dcdr =
252 case dcdr of
253 Leaf (Inline expr) ->
254 Leaf (Inline (go expr))
256 Leaf (Jump _target) ->
257 dcdr
259 Chain chain success failure ->
260 Chain chain (goDecider success) (goDecider failure)
262 FanOut path tests fallback ->
263 FanOut path (map (second goDecider) tests) (goDecider fallback)
264 in
265 Case exprName (goDecider decider) (map (second go) jumps)
267 Ctor tag values ->
268 Ctor tag (map go values)
270 CtorAccess root index ->
271 CtorAccess (go root) index
273 Access record field ->
274 Access (go record) field
276 Update record fields ->
277 Update (go record) (map (second go) fields)
279 Record fields ->
280 Record (map (second go) fields)
282 Cmd _ ->
283 expression
285 Sub _ ->
286 expression
288 OutgoingPort _ _ ->
289 expression
291 IncomingPort _ _ ->
292 expression
294 Program kind expr ->
295 Program kind (go expr)
297 GLShader _ _ _ ->
298 expression
300 Crash home region maybeBranchProblem ->
301 Crash home region (fmap go maybeBranchProblem)