git.iko.soy/elk
on 2ec15df help sign in
elk/compiler/src/Optimize/Expression.hs
466 lines 14 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.registerDebug name home region
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 <- destructCase 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)
212destructCase :: N.Name -> Can.Pattern -> Names.Tracker [Opt.Destructor]
213destructCase rootName pattern =
214 reverse <$> destructHelp (Opt.Root rootName) pattern []
217destruct :: Can.Pattern -> Names.Tracker (N.Name, [Opt.Destructor])
218destruct pattern@(A.At _ ptrn) =
219 case ptrn of
220 Can.PVar name ->
221 pure (name, [])
223 Can.PAlias subPattern name ->
224 do revDs <- destructHelp (Opt.Root name) subPattern []
225 pure (name, reverse revDs)
227 _ ->
228 do name <- Names.generate
229 revDs <- destructHelp (Opt.Root name) pattern []
230 pure (name, reverse revDs)
233destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
234destructHelp path (A.At region pattern) revDs =
235 case pattern of
236 Can.PAnything ->
237 pure revDs
239 Can.PVar name ->
240 pure (Opt.Destructor name path : revDs)
242 Can.PRecord fields ->
243 let
244 toDestruct name =
245 Opt.Destructor name (Opt.Field name path)
246 in
247 Names.registerFieldList fields (map toDestruct fields ++ revDs)
249 Can.PAlias subPattern name ->
250 destructHelp (Opt.Root name) subPattern $
251 Opt.Destructor name path : revDs
253 Can.PUnit ->
254 pure revDs
256 Can.PTuple a b Nothing ->
257 destructTwo path a b revDs
259 Can.PTuple a b (Just c) ->
260 case path of
261 Opt.Root _ ->
262 destructHelp (Opt.Index Index.third path) c =<<
263 destructHelp (Opt.Index Index.second path) b =<<
264 destructHelp (Opt.Index Index.first path) a revDs
266 _ ->
267 do name <- Names.generate
268 let newRoot = Opt.Root name
269 destructHelp (Opt.Index Index.third newRoot) c =<<
270 destructHelp (Opt.Index Index.second newRoot) b =<<
271 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
273 Can.PList [] ->
274 pure revDs
276 Can.PList (hd:tl) ->
277 destructTwo path hd (A.At region (Can.PList tl)) revDs
279 Can.PCons hd tl ->
280 destructTwo path hd tl revDs
282 Can.PChr _ ->
283 pure revDs
285 Can.PStr _ ->
286 pure revDs
288 Can.PInt _ ->
289 pure revDs
291 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
292 case args of
293 [Can.PatternCtorArg _ _ arg] ->
294 case opts of
295 Can.Normal -> destructHelp (Opt.Index Index.first path) arg revDs
296 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
297 Can.Enum -> destructHelp (Opt.Index Index.first path) arg revDs
299 _ ->
300 case path of
301 Opt.Root _ ->
302 foldM (destructCtorArg path) revDs args
304 _ ->
305 do name <- Names.generate
306 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
309destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
310destructTwo path a b revDs =
311 case path of
312 Opt.Root _ ->
313 destructHelp (Opt.Index Index.second path) b =<<
314 destructHelp (Opt.Index Index.first path) a revDs
316 _ ->
317 do name <- Names.generate
318 let newRoot = Opt.Root name
319 destructHelp (Opt.Index Index.second newRoot) b =<<
320 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
323destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
324destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
325 destructHelp (Opt.Index index path) arg revDs
329-- TAIL CALL
332optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
333optimizePotentialTailCall cycle def =
334 case def of
335 Can.Def (A.At _ name) args expr ->
336 do (argNames, destructors) <- destructArgs args
337 toTailDef name argNames destructors <$>
338 optimizeTail cycle name argNames expr
340 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
341 do (argNames, destructors) <- destructArgs (map fst typedArgs)
342 toTailDef name argNames destructors <$>
343 optimizeTail cycle name argNames expr
346optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
347optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
348 case expression of
349 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
350 do oargs <- traverse (optimize cycle) args
351 if name == rootName
352 then
353 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
354 Index.LengthMatch pairs ->
355 pure $ Opt.TailCall name pairs
357 Index.LengthMismatch _ _ ->
358 do ofunc <- optimize cycle func
359 pure $ Opt.Call ofunc oargs
360 else
361 do ofunc <- optimize cycle func
362 pure $ Opt.Call ofunc oargs
364 Can.If branches finally ->
365 let
366 optimizeBranch (condition, branch) =
367 (,)
368 <$> optimize cycle condition
369 <*> optimizeTail cycle rootName argNames branch
370 in
371 Opt.If
372 <$> traverse optimizeBranch branches
373 <*> optimizeTail cycle rootName argNames finally
375 Can.Let def body ->
376 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
378 Can.LetRec defs body ->
379 case defs of
380 [def] ->
381 Opt.Let
382 <$> optimizePotentialTailCall cycle def
383 <*> optimizeTail cycle rootName argNames body
385 _ ->
386 do obody <- optimizeTail cycle rootName argNames body
387 foldM (\bod def -> optimizeDef cycle def bod) obody defs
389 Can.LetDestruct pattern expr body ->
390 do (dname, destructors) <- destruct pattern
391 oexpr <- optimize cycle expr
392 obody <- optimizeTail cycle dname argNames body
393 pure $
394 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
396 Can.Case expr branches ->
397 let
398 optimizeBranch root (Can.CaseBranch pattern branch) =
399 do destructors <- destructCase root pattern
400 obranch <- optimizeTail cycle rootName argNames branch
401 pure (pattern, foldr Opt.Destruct obranch destructors)
402 in
403 do temp <- Names.generate
404 oexpr <- optimize cycle expr
405 case oexpr of
406 Opt.VarLocal root ->
407 Case.optimize temp root <$> traverse (optimizeBranch root) branches
409 _ ->
410 do obranches <- traverse (optimizeBranch temp) branches
411 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
413 _ ->
414 optimize cycle locExpr
418-- DETECT TAIL CALLS
421toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
422toTailDef name argNames destructors body =
423 if hasTailCall body then
424 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
425 else
426 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
429hasTailCall :: Opt.Expr -> Bool
430hasTailCall expression =
431 case expression of
432 Opt.TailCall _ _ ->
433 True
435 Opt.If branches finally ->
436 hasTailCall finally || any (hasTailCall . snd) branches
438 Opt.Let _ body ->
439 hasTailCall body
441 Opt.Destruct _ body ->
442 hasTailCall body
444 Opt.Case _ _ decider jumps ->
445 decidecHasTailCall decider || any (hasTailCall . snd) jumps
447 _ ->
448 False
451decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
452decidecHasTailCall decider =
453 case decider of
454 Opt.Leaf choice ->
455 case choice of
456 Opt.Inline expr ->
457 hasTailCall expr
459 Opt.Jump _ ->
460 False
462 Opt.Chain _ success failure ->
463 decidecHasTailCall success || decidecHasTailCall failure
465 Opt.FanOut _ tests fallback ->
466 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests