git.iko.soy/elk
on 022925a help sign in
elk/src/Optimize/Expression.hs
403 lines 11 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Expression
4 ( Cycle
5 , optimize
6 , destructArgs
7 , optimizePotentialTailCall
8 )
9 where
12import Prelude hiding (cycle)
13import Control.Monad (foldM)
14import qualified Data.Set as Set
16import qualified AST.Canonical as Can
17import qualified AST.Optimized as Opt
18import qualified AST.Module.Name as ModuleName
19import qualified Data.Index as Index
20import qualified Elm.Name as N
21import qualified Optimize.Case as Case
22import qualified Optimize.Names as Names
23import qualified Reporting.Annotation as A
27-- OPTIMIZE
30type Cycle =
31 Set.Set N.Name
34optimize :: Cycle -> Can.Expr -> Names.Tracker Opt.Expr
35optimize cycle (A.At region expression) =
36 case expression of
37 Can.VarLocal name ->
38 pure (Opt.VarLocal name)
40 Can.VarTopLevel home name ->
41 if Set.member name cycle then
42 pure (Opt.VarCycle home name)
43 else
44 Names.registerGlobal home name
46 Can.VarKernel home name ->
47 Names.registerKernel home (Opt.VarKernel home name)
49 Can.VarForeign home name _ ->
50 Names.registerGlobal home name
52 Can.VarDebug home name _ ->
53 Names.registerKernel N.debug (Opt.VarDebug name home region Nothing)
55 Can.VarOperator _ home name _ ->
56 Names.registerGlobal home name
58 Can.Chr chr ->
59 pure (Opt.Chr chr)
61 Can.Str str ->
62 pure (Opt.Str str)
64 Can.Int int ->
65 pure (Opt.Int int)
67 Can.Float float ->
68 pure (Opt.Float float)
70 Can.List entries ->
71 Opt.List <$> traverse (optimize cycle) entries
73 Can.Negate expr ->
74 do func <- Names.registerGlobal ModuleName.basics N.negate
75 arg <- optimize cycle expr
76 pure $ Opt.Call func [arg]
78 Can.Binop _ home name _ left right ->
79 do optFunc <- Names.registerGlobal home name
80 optLeft <- optimize cycle left
81 optRight <- optimize cycle right
82 return (Opt.Call optFunc [optLeft, optRight])
84 Can.Lambda args body ->
85 do (argNames, defss) <- destructArgs args
86 obody <- optimize cycle body
87 pure $ Opt.Function argNames (foldr Opt.Let obody (concat defss))
89 Can.Call func args ->
90 Opt.Call
91 <$> optimize cycle func
92 <*> traverse (optimize cycle) args
94 Can.If branches finally ->
95 let
96 optimizeBranch (condition, branch) =
97 (,)
98 <$> optimize cycle condition
99 <*> optimize cycle branch
100 in
101 Opt.If
102 <$> traverse optimizeBranch branches
103 <*> optimize cycle finally
105 Can.Let def body ->
106 optimizeDef cycle def =<< optimize cycle body
108 Can.LetRec defs body ->
109 case defs of
110 [def] ->
111 Opt.Let
112 <$> optimizePotentialTailCall cycle def
113 <*> optimize cycle body
115 _ ->
116 do obody <- optimize cycle body
117 foldM (\bod def -> optimizeDef cycle def bod) obody defs
119 Can.LetDestruct pattern expr body ->
120 do (name, defs) <- destruct pattern
121 oexpr <- optimize cycle expr
122 obody <- optimize cycle body
123 pure $
124 Opt.Let (Opt.Def name oexpr) (foldr Opt.Let obody defs)
126 Can.Case expr branches ->
127 let
128 optimizeBranch root (Can.CaseBranch pattern branch) =
129 do defs <- destructHelp (Opt.VarLocal root) pattern []
130 obranch <- optimize cycle branch
131 pure (pattern, foldr Opt.Let obranch defs)
132 in
133 do oexpr <- optimize cycle expr
134 case oexpr of
135 Opt.VarLocal root ->
136 Case.optimize root <$> traverse (optimizeBranch root) branches
138 _ ->
139 do root <- Names.generate
140 Case.optimize root <$> traverse (optimizeBranch root) branches
142 Can.Accessor field ->
143 Names.registerField field (Opt.Accessor field)
145 Can.Access record field ->
146 do optRecord <- optimize cycle record
147 Names.registerField field (Opt.Access optRecord field)
149 Can.Update record updates ->
150 Names.registerFieldDict updates Opt.Update
151 <*> optimize cycle record
152 <*> traverse (optimize cycle) updates
154 Can.Record fields ->
155 Names.registerFieldDict fields Opt.Record
156 <*> traverse (optimize cycle) fields
158 Can.Unit ->
159 Names.registerKernel N.utils Opt.Unit
161 Can.Tuple a b maybeC ->
162 Names.registerKernel N.utils Opt.Tuple
163 <*> optimize cycle a
164 <*> optimize cycle b
165 <*> traverse (optimize cycle) maybeC
167 Can.Shader _ src _ ->
168 pure (Opt.Shader src)
172-- DEFINITION
175optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
176optimizeDef cycle def body =
177 case def of
178 Can.Def (A.At _ name) args expr ->
179 optimizeDefHelp cycle name args expr body
181 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
182 optimizeDefHelp cycle name (map fst typedArgs) expr body
185optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
186optimizeDefHelp cycle name args expr body =
187 case args of
188 [] ->
189 do oexpr <- optimize cycle expr
190 pure $ Opt.Let (Opt.Def name oexpr) body
192 _ ->
193 do (argNames, defss) <- destructArgs args
194 let oexpr = Opt.Function argNames (foldr Opt.Let body (concat defss))
195 pure $ Opt.Let (Opt.Def name oexpr) body
199-- DESTRUCTURING
202destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [[Opt.Def]])
203destructArgs args =
204 unzip <$> traverse destruct args
207destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Def])
208destruct pattern@(A.At _ ptrn) =
209 case ptrn of
210 Can.PVar name ->
211 pure (name, [])
213 Can.PAlias subPattern name ->
214 do revDefs <- destructHelp (Opt.VarLocal name) subPattern []
215 pure (name, reverse revDefs)
217 _ ->
218 do name <- Names.generate
219 revDefs <- destructHelp (Opt.VarLocal name) pattern []
220 pure (name, reverse revDefs)
223destructHelp :: Opt.Expr -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
224destructHelp root (A.At region pattern) revDefs =
225 case pattern of
226 Can.PAnything ->
227 pure revDefs
229 Can.PVar name ->
230 pure (Opt.Def name root : revDefs)
232 Can.PRecord fields ->
233 let toDef name = Opt.Def name (Opt.Access root name) in
234 Names.registerFieldList fields (map toDef fields ++ revDefs)
236 Can.PAlias subPattern name ->
237 destructHelp (Opt.VarLocal name) subPattern $
238 Opt.Def name root : revDefs
240 Can.PUnit ->
241 pure revDefs
243 Can.PTuple a b maybeC ->
244 do name <- Names.generate
245 let tuple = Opt.VarLocal name
246 let newRevDefs = Opt.Def name root : revDefs
247 case maybeC of
248 Nothing ->
249 destructHelp (Opt.CtorAccess tuple Index.second) b =<<
250 destructHelp (Opt.CtorAccess tuple Index.first) a newRevDefs
252 Just c ->
253 destructHelp (Opt.CtorAccess tuple Index.third) c =<<
254 destructHelp (Opt.CtorAccess tuple Index.second) b =<<
255 destructHelp (Opt.CtorAccess tuple Index.first) a newRevDefs
257 Can.PList [] ->
258 pure revDefs
260 Can.PList (hd:tl) ->
261 destructCons root hd (A.At region (Can.PList tl)) revDefs
263 Can.PCons hd tl ->
264 destructCons root hd tl revDefs
266 Can.PChr _ ->
267 pure revDefs
269 Can.PStr _ ->
270 pure revDefs
272 Can.PInt _ ->
273 pure revDefs
275 Can.PCtor _ _ _ _ _ args ->
276 do name <- Names.generate
277 foldM
278 (destructCtorArg (Opt.VarLocal name))
279 (Opt.Def name root : revDefs)
280 args
283destructCons :: Opt.Expr -> Can.Pattern -> Can.Pattern -> [Opt.Def] -> Names.Tracker [Opt.Def]
284destructCons root hd tl revDefs =
285 do name <- Names.generate
286 let list = Opt.VarLocal name
287 destructHelp (Opt.CtorAccess list Index.first) hd =<<
288 destructHelp (Opt.CtorAccess list Index.second) tl (Opt.Def name root : revDefs)
291destructCtorArg :: Opt.Expr -> [Opt.Def] -> Can.PatternCtorArg -> Names.Tracker [Opt.Def]
292destructCtorArg root revDefs (Can.PatternCtorArg index _ arg) =
293 destructHelp (Opt.CtorAccess root index) arg revDefs
297-- TAIL CALL
300optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
301optimizePotentialTailCall cycle def =
302 case def of
303 Can.Def (A.At _ name) args expr ->
304 do (argNames, defss) <- destructArgs args
305 toTailDef name argNames defss <$>
306 optimizeTail cycle name argNames expr
308 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
309 do (argNames, defss) <- destructArgs (map fst typedArgs)
310 toTailDef name argNames defss <$>
311 optimizeTail cycle name argNames expr
314optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
315optimizeTail cycle name argNames locExpr@(A.At _ expression) =
316 case expression of
317 Can.Call (A.At _ (Can.VarTopLevel _ funcName)) args ->
318 if funcName == name && length args == length argNames then
319 Opt.TailCall name argNames <$> traverse (optimize cycle) args
320 else
321 optimize cycle locExpr
323 Can.If branches finally ->
324 let
325 optimizeBranch (condition, branch) =
326 (,)
327 <$> optimize cycle condition
328 <*> optimizeTail cycle name argNames branch
329 in
330 Opt.If
331 <$> traverse optimizeBranch branches
332 <*> optimizeTail cycle name argNames finally
334 Can.Let def body ->
335 optimizeDef cycle def =<< optimizeTail cycle name argNames body
337 Can.LetRec defs body ->
338 case defs of
339 [def] ->
340 Opt.Let
341 <$> optimizePotentialTailCall cycle def
342 <*> optimizeTail cycle name argNames body
344 _ ->
345 do obody <- optimizeTail cycle name argNames body
346 foldM (\bod def -> optimizeDef cycle def bod) obody defs
348 Can.LetDestruct pattern expr body ->
349 do (dname, defs) <- destruct pattern
350 oexpr <- optimize cycle expr
351 obody <- optimizeTail cycle dname argNames body
352 pure $
353 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Let obody defs)
355 Can.Case expr branches ->
356 let
357 optimizeBranch root (Can.CaseBranch pattern branch) =
358 do defs <- destructHelp (Opt.VarLocal root) pattern []
359 obranch <- optimizeTail cycle name argNames branch
360 pure (pattern, foldr Opt.Let obranch defs)
361 in
362 do oexpr <- optimize cycle expr
363 case oexpr of
364 Opt.VarLocal root ->
365 Case.optimize root <$> traverse (optimizeBranch root) branches
367 _ ->
368 do root <- Names.generate
369 Case.optimize root <$> traverse (optimizeBranch root) branches
371 _ ->
372 optimize cycle locExpr
376-- DETECT TAIL CALLS
379toTailDef :: N.Name -> [N.Name] -> [[Opt.Def]] -> Opt.Expr -> Opt.Def
380toTailDef name argNames defss body =
381 if hasTailCall body then
382 Opt.TailDef name argNames (foldr Opt.Let body (concat defss))
383 else
384 Opt.Def name (Opt.Function argNames (foldr Opt.Let body (concat defss)))
387hasTailCall :: Opt.Expr -> Bool
388hasTailCall expression =
389 case expression of
390 Opt.TailCall _ _ _ ->
391 True
393 Opt.If branches finally ->
394 hasTailCall finally || any (hasTailCall . snd) branches
396 Opt.Let _ body ->
397 hasTailCall body
399 Opt.Case _ _ branches ->
400 any (hasTailCall . snd) branches
402 _ ->
403 False