git.iko.soy/elk
on 14fc72d help sign in
elk/src/Optimize/Expression.hs
460 lines 13 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.VarCtor opts home name index _ ->
53 Names.registerCtor home name index opts
55 Can.VarDebug home name _ ->
56 Names.registerKernel N.debug (Opt.VarDebug name home region Nothing)
58 Can.VarOperator _ home name _ ->
59 Names.registerGlobal home name
61 Can.Chr chr ->
62 Names.registerKernel N.utils (Opt.Chr chr)
64 Can.Str str ->
65 pure (Opt.Str str)
67 Can.Int int ->
68 pure (Opt.Int int)
70 Can.Float float ->
71 pure (Opt.Float float)
73 Can.List entries ->
74 Opt.List <$> traverse (optimize cycle) entries
76 Can.Negate expr ->
77 do func <- Names.registerGlobal ModuleName.basics N.negate
78 arg <- optimize cycle expr
79 pure $ Opt.Call func [arg]
81 Can.Binop _ home name _ left right ->
82 do optFunc <- Names.registerGlobal home name
83 optLeft <- optimize cycle left
84 optRight <- optimize cycle right
85 return (Opt.Call optFunc [optLeft, optRight])
87 Can.Lambda args body ->
88 do (argNames, destructors) <- destructArgs args
89 obody <- optimize cycle body
90 pure $ Opt.Function argNames (foldr Opt.Destruct obody destructors)
92 Can.Call func args ->
93 Opt.Call
94 <$> optimize cycle func
95 <*> traverse (optimize cycle) args
97 Can.If branches finally ->
98 let
99 optimizeBranch (condition, branch) =
100 (,)
101 <$> optimize cycle condition
102 <*> optimize cycle branch
103 in
104 Opt.If
105 <$> traverse optimizeBranch branches
106 <*> optimize cycle finally
108 Can.Let def body ->
109 optimizeDef cycle def =<< optimize cycle body
111 Can.LetRec defs body ->
112 case defs of
113 [def] ->
114 Opt.Let
115 <$> optimizePotentialTailCall cycle def
116 <*> optimize cycle body
118 _ ->
119 do obody <- optimize cycle body
120 foldM (\bod def -> optimizeDef cycle def bod) obody defs
122 Can.LetDestruct pattern expr body ->
123 do (name, destructs) <- destruct pattern
124 oexpr <- optimize cycle expr
125 obody <- optimize cycle body
126 pure $
127 Opt.Let (Opt.Def name oexpr) (foldr Opt.Destruct obody destructs)
129 Can.Case expr branches ->
130 let
131 optimizeBranch root (Can.CaseBranch pattern branch) =
132 do destructors <- destructHelp (Opt.Root root) pattern []
133 obranch <- optimize cycle branch
134 pure (pattern, foldr Opt.Destruct obranch destructors)
135 in
136 do temp <- Names.generate
137 oexpr <- optimize cycle expr
138 case oexpr of
139 Opt.VarLocal root ->
140 Case.optimize temp root <$> traverse (optimizeBranch root) branches
142 _ ->
143 do obranches <- traverse (optimizeBranch temp) branches
144 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
146 Can.Accessor field ->
147 Names.registerField field (Opt.Accessor field)
149 Can.Access record field ->
150 do optRecord <- optimize cycle record
151 Names.registerField field (Opt.Access optRecord field)
153 Can.Update record updates ->
154 Names.registerFieldDict updates Opt.Update
155 <*> optimize cycle record
156 <*> traverse (optimize cycle) updates
158 Can.Record fields ->
159 Names.registerFieldDict fields Opt.Record
160 <*> traverse (optimize cycle) fields
162 Can.Unit ->
163 Names.registerKernel N.utils Opt.Unit
165 Can.Tuple a b maybeC ->
166 Names.registerKernel N.utils Opt.Tuple
167 <*> optimize cycle a
168 <*> optimize cycle b
169 <*> traverse (optimize cycle) maybeC
171 Can.Shader _ src _ ->
172 pure (Opt.Shader src)
176-- DEFINITION
179optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
180optimizeDef cycle def body =
181 case def of
182 Can.Def (A.At _ name) args expr ->
183 optimizeDefHelp cycle name args expr body
185 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
186 optimizeDefHelp cycle name (map fst typedArgs) expr body
189optimizeDefHelp :: Cycle -> N.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
190optimizeDefHelp cycle name args expr body =
191 do oexpr <- optimize cycle expr
192 case args of
193 [] ->
194 pure $ Opt.Let (Opt.Def name oexpr) body
196 _ ->
197 do (argNames, destructors) <- destructArgs args
198 let ofunc = Opt.Function argNames (foldr Opt.Destruct oexpr destructors)
199 pure $ Opt.Let (Opt.Def name ofunc) body
203-- DESTRUCTURING
206destructArgs :: [Can.Pattern] -> Names.Tracker ([N.Name], [Opt.Destructor])
207destructArgs args =
208 do (argNames, destructorLists) <- unzip <$> traverse destruct args
209 return (argNames, concat destructorLists)
212destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Destructor])
213destruct pattern@(A.At _ ptrn) =
214 case ptrn of
215 Can.PVar name ->
216 pure (name, [])
218 Can.PAlias subPattern name ->
219 do revDs <- destructHelp (Opt.Root name) subPattern []
220 pure (name, reverse revDs)
222 _ ->
223 do name <- Names.generate
224 revDs <- destructHelp (Opt.Root name) pattern []
225 pure (name, reverse revDs)
228destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
229destructHelp path (A.At region pattern) revDs =
230 case pattern of
231 Can.PAnything ->
232 pure revDs
234 Can.PVar name ->
235 pure (Opt.Destructor name path : revDs)
237 Can.PRecord fields ->
238 let
239 toDestruct name =
240 Opt.Destructor name (Opt.Field name path)
241 in
242 Names.registerFieldList fields (map toDestruct fields ++ revDs)
244 Can.PAlias subPattern name ->
245 destructHelp (Opt.Root name) subPattern $
246 Opt.Destructor name path : revDs
248 Can.PUnit ->
249 pure revDs
251 Can.PTuple a b Nothing ->
252 destructTwo path a b revDs
254 Can.PTuple a b (Just c) ->
255 case path of
256 Opt.Root _ ->
257 destructHelp (Opt.Index Index.first path) a =<<
258 destructHelp (Opt.Index Index.second path) b =<<
259 destructHelp (Opt.Index Index.third path) c revDs
261 _ ->
262 do name <- Names.generate
263 let newRoot = Opt.Root name
264 destructHelp (Opt.Index Index.first newRoot) a =<<
265 destructHelp (Opt.Index Index.second newRoot) b =<<
266 destructHelp (Opt.Index Index.third newRoot) c (Opt.Destructor name path : revDs)
268 Can.PList [] ->
269 pure revDs
271 Can.PList (hd:tl) ->
272 destructTwo path hd (A.At region (Can.PList tl)) revDs
274 Can.PCons hd tl ->
275 destructTwo path hd tl revDs
277 Can.PChr _ ->
278 pure revDs
280 Can.PStr _ ->
281 pure revDs
283 Can.PInt _ ->
284 pure revDs
286 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
287 case args of
288 [Can.PatternCtorArg i _ arg] ->
289 case opts of
290 Can.Normal -> destructHelp (Opt.Index i path) arg revDs
291 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
292 Can.Enum -> destructHelp (Opt.Index i path) arg revDs
294 _ ->
295 case path of
296 Opt.Root _ ->
297 foldM (destructCtorArg path) revDs args
299 _ ->
300 do name <- Names.generate
301 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
304destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
305destructTwo path a b revDs =
306 case path of
307 Opt.Root _ ->
308 destructHelp (Opt.Index Index.first path) a =<<
309 destructHelp (Opt.Index Index.second path) b revDs
311 _ ->
312 do name <- Names.generate
313 let newRoot = Opt.Root name
314 destructHelp (Opt.Index Index.first newRoot) a =<<
315 destructHelp (Opt.Index Index.second newRoot) b (Opt.Destructor name path : revDs)
318destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
319destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
320 destructHelp (Opt.Index index path) arg revDs
324-- TAIL CALL
327optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
328optimizePotentialTailCall cycle def =
329 case def of
330 Can.Def (A.At _ name) args expr ->
331 do (argNames, destructors) <- destructArgs args
332 toTailDef name argNames destructors <$>
333 optimizeTail cycle name argNames expr
335 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
336 do (argNames, destructors) <- destructArgs (map fst typedArgs)
337 toTailDef name argNames destructors <$>
338 optimizeTail cycle name argNames expr
341optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
342optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
343 case expression of
344 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
345 do oargs <- traverse (optimize cycle) args
346 if name == rootName
347 then
348 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
349 Index.LengthMatch pairs ->
350 pure $ Opt.TailCall name pairs
352 Index.LengthMismatch _ _ ->
353 do ofunc <- optimize cycle func
354 pure $ Opt.Call ofunc oargs
355 else
356 do ofunc <- optimize cycle func
357 pure $ Opt.Call ofunc oargs
359 Can.If branches finally ->
360 let
361 optimizeBranch (condition, branch) =
362 (,)
363 <$> optimize cycle condition
364 <*> optimizeTail cycle rootName argNames branch
365 in
366 Opt.If
367 <$> traverse optimizeBranch branches
368 <*> optimizeTail cycle rootName argNames finally
370 Can.Let def body ->
371 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
373 Can.LetRec defs body ->
374 case defs of
375 [def] ->
376 Opt.Let
377 <$> optimizePotentialTailCall cycle def
378 <*> optimizeTail cycle rootName argNames body
380 _ ->
381 do obody <- optimizeTail cycle rootName argNames body
382 foldM (\bod def -> optimizeDef cycle def bod) obody defs
384 Can.LetDestruct pattern expr body ->
385 do (dname, destructors) <- destruct pattern
386 oexpr <- optimize cycle expr
387 obody <- optimizeTail cycle dname argNames body
388 pure $
389 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
391 Can.Case expr branches ->
392 let
393 optimizeBranch root (Can.CaseBranch pattern branch) =
394 do destructors <- destructHelp (Opt.Root root) pattern []
395 obranch <- optimizeTail cycle rootName argNames branch
396 pure (pattern, foldr Opt.Destruct obranch destructors)
397 in
398 do temp <- Names.generate
399 oexpr <- optimize cycle expr
400 case oexpr of
401 Opt.VarLocal root ->
402 Case.optimize temp root <$> traverse (optimizeBranch root) branches
404 _ ->
405 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
407 _ ->
408 optimize cycle locExpr
412-- DETECT TAIL CALLS
415toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
416toTailDef name argNames destructors body =
417 if hasTailCall body then
418 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
419 else
420 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
423hasTailCall :: Opt.Expr -> Bool
424hasTailCall expression =
425 case expression of
426 Opt.TailCall _ _ ->
427 True
429 Opt.If branches finally ->
430 hasTailCall finally || any (hasTailCall . snd) branches
432 Opt.Let _ body ->
433 hasTailCall body
435 Opt.Destruct _ body ->
436 hasTailCall body
438 Opt.Case _ _ decider jumps ->
439 decidecHasTailCall decider || any (hasTailCall . snd) jumps
441 _ ->
442 False
445decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
446decidecHasTailCall decider =
447 case decider of
448 Opt.Leaf choice ->
449 case choice of
450 Opt.Inline expr ->
451 hasTailCall expr
453 Opt.Jump _ ->
454 False
456 Opt.Chain _ success failure ->
457 decidecHasTailCall success || decidecHasTailCall failure
459 Opt.FanOut _ tests fallback ->
460 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests