git.iko.soy/elk
on 1238607 help sign in
elk/compiler/src/Optimize/Expression.hs
490 lines 14 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.Expression
4 ( optimize
5 , destructArgs
6 , optimizePotentialTailCall
7 )
8 where
11import Prelude hiding (cycle)
12import Control.Monad (foldM)
13import qualified Data.Map as Map
14import qualified Data.Name as Name
15import qualified Data.Set as Set
17import qualified AST.Canonical as Can
18import qualified AST.Optimized as Opt
19import qualified AST.Utils.Shader as Shader
20import qualified Data.Index as Index
21import qualified Elm.ModuleName as ModuleName
22import qualified Optimize.Case as Case
23import qualified Optimize.Names as Names
24import qualified Reporting.Annotation as A
28-- OPTIMIZE
31type Cycle =
32 Set.Set Name.Name
35optimize :: Cycle -> Can.Expr -> Names.Tracker Opt.Expr
36optimize cycle (A.At region expression) =
37 case expression of
38 Can.VarLocal name ->
39 pure (Opt.VarLocal name)
41 Can.VarTopLevel home name ->
42 if Set.member name cycle then
43 pure (Opt.VarCycle home name)
44 else
45 Names.registerGlobal home name
47 Can.VarKernel home name ->
48 Names.registerKernel home (Opt.VarKernel home name)
50 Can.VarForeign home name _ ->
51 Names.registerGlobal home name
53 Can.VarCtor opts home name index _ ->
54 Names.registerCtor home name index opts
56 Can.VarDebug home name _ ->
57 Names.registerDebug name home region
59 Can.VarOperator _ home name _ ->
60 Names.registerGlobal home name
62 Can.Chr chr ->
63 Names.registerKernel Name.utils (Opt.Chr chr)
65 Can.Str str ->
66 pure (Opt.Str str)
68 Can.Int int ->
69 pure (Opt.Int int)
71 Can.Float float ->
72 pure (Opt.Float float)
74 Can.List entries ->
75 Names.registerKernel Name.list Opt.List
76 <*> traverse (optimize cycle) entries
78 Can.Negate expr ->
79 do func <- Names.registerGlobal ModuleName.basics Name.negate
80 arg <- optimize cycle expr
81 pure $ Opt.Call func [arg]
83 Can.Binop _ home name _ left right ->
84 do optFunc <- Names.registerGlobal home name
85 optLeft <- optimize cycle left
86 optRight <- optimize cycle right
87 return (Opt.Call optFunc [optLeft, optRight])
89 Can.Lambda args body ->
90 do (argNames, destructors) <- destructArgs args
91 obody <- optimize cycle body
92 pure $ Opt.Function argNames (foldr Opt.Destruct obody destructors)
94 Can.Call func args ->
95 Opt.Call
96 <$> optimize cycle func
97 <*> traverse (optimize cycle) args
99 Can.If branches finally ->
100 let
101 optimizeBranch (condition, branch) =
102 (,)
103 <$> optimize cycle condition
104 <*> optimize cycle branch
105 in
106 Opt.If
107 <$> traverse optimizeBranch branches
108 <*> optimize cycle finally
110 Can.Let def body ->
111 optimizeDef cycle def =<< optimize cycle body
113 Can.LetRec defs body ->
114 case defs of
115 [def] ->
116 Opt.Let
117 <$> optimizePotentialTailCallDef cycle def
118 <*> optimize cycle body
120 _ ->
121 do obody <- optimize cycle body
122 foldM (\bod def -> optimizeDef cycle def bod) obody defs
124 Can.LetDestruct pattern expr body ->
125 do (name, destructs) <- destruct pattern
126 oexpr <- optimize cycle expr
127 obody <- optimize cycle body
128 pure $
129 Opt.Let (Opt.Def name oexpr) (foldr Opt.Destruct obody destructs)
131 Can.Case expr branches ->
132 let
133 optimizeBranch root (Can.CaseBranch pattern branch) =
134 do destructors <- destructCase root pattern
135 obranch <- optimize cycle branch
136 pure (pattern, foldr Opt.Destruct obranch destructors)
137 in
138 do temp <- Names.generate
139 oexpr <- optimize cycle expr
140 case oexpr of
141 Opt.VarLocal root ->
142 Case.optimize temp root <$> traverse (optimizeBranch root) branches
144 _ ->
145 do obranches <- traverse (optimizeBranch temp) branches
146 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
148 Can.Accessor field ->
149 Names.registerField field (Opt.Accessor field)
151 Can.Access record (A.At _ field) ->
152 do optRecord <- optimize cycle record
153 Names.registerField field (Opt.Access optRecord field)
155 Can.Update _ record updates ->
156 Names.registerFieldDict updates Opt.Update
157 <*> optimize cycle record
158 <*> traverse (optimizeUpdate cycle) updates
160 Can.Record fields ->
161 Names.registerFieldDict fields Opt.Record
162 <*> traverse (optimize cycle) fields
164 Can.Unit ->
165 Names.registerKernel Name.utils Opt.Unit
167 Can.Tuple a b maybeC ->
168 Names.registerKernel Name.utils Opt.Tuple
169 <*> optimize cycle a
170 <*> optimize cycle b
171 <*> traverse (optimize cycle) maybeC
173 Can.Shader src (Shader.Types attributes uniforms _varyings) ->
174 pure (Opt.Shader src (Map.keysSet attributes) (Map.keysSet uniforms))
178-- UPDATE
181optimizeUpdate :: Cycle -> Can.FieldUpdate -> Names.Tracker Opt.Expr
182optimizeUpdate cycle (Can.FieldUpdate _ expr) =
183 optimize cycle expr
187-- DEFINITION
190optimizeDef :: Cycle -> Can.Def -> Opt.Expr -> Names.Tracker Opt.Expr
191optimizeDef cycle def body =
192 case def of
193 Can.Def (A.At _ name) args expr ->
194 optimizeDefHelp cycle name args expr body
196 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
197 optimizeDefHelp cycle name (map fst typedArgs) expr body
200optimizeDefHelp :: Cycle -> Name.Name -> [Can.Pattern] -> Can.Expr -> Opt.Expr -> Names.Tracker Opt.Expr
201optimizeDefHelp cycle name args expr body =
202 do oexpr <- optimize cycle expr
203 case args of
204 [] ->
205 pure $ Opt.Let (Opt.Def name oexpr) body
207 _ ->
208 do (argNames, destructors) <- destructArgs args
209 let ofunc = Opt.Function argNames (foldr Opt.Destruct oexpr destructors)
210 pure $ Opt.Let (Opt.Def name ofunc) body
214-- DESTRUCTURING
217destructArgs :: [Can.Pattern] -> Names.Tracker ([Name.Name], [Opt.Destructor])
218destructArgs args =
219 do (argNames, destructorLists) <- unzip <$> traverse destruct args
220 return (argNames, concat destructorLists)
223destructCase :: Name.Name -> Can.Pattern -> Names.Tracker [Opt.Destructor]
224destructCase rootName pattern =
225 reverse <$> destructHelp (Opt.Root rootName) pattern []
228destruct :: Can.Pattern -> Names.Tracker (Name.Name, [Opt.Destructor])
229destruct pattern@(A.At _ ptrn) =
230 case ptrn of
231 Can.PVar name ->
232 pure (name, [])
234 Can.PAlias subPattern name ->
235 do revDs <- destructHelp (Opt.Root name) subPattern []
236 pure (name, reverse revDs)
238 _ ->
239 do name <- Names.generate
240 revDs <- destructHelp (Opt.Root name) pattern []
241 pure (name, reverse revDs)
244destructHelp :: Opt.Path -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
245destructHelp path (A.At region pattern) revDs =
246 case pattern of
247 Can.PAnything ->
248 pure revDs
250 Can.PVar name ->
251 pure (Opt.Destructor name path : revDs)
253 Can.PRecord fields ->
254 let
255 toDestruct name =
256 Opt.Destructor name (Opt.Field name path)
257 in
258 Names.registerFieldList fields (map toDestruct fields ++ revDs)
260 Can.PAlias subPattern name ->
261 destructHelp (Opt.Root name) subPattern $
262 Opt.Destructor name path : revDs
264 Can.PUnit ->
265 pure revDs
267 Can.PTuple a b Nothing ->
268 destructTwo path a b revDs
270 Can.PTuple a b (Just c) ->
271 case path of
272 Opt.Root _ ->
273 destructHelp (Opt.Index Index.third path) c =<<
274 destructHelp (Opt.Index Index.second path) b =<<
275 destructHelp (Opt.Index Index.first path) a revDs
277 _ ->
278 do name <- Names.generate
279 let newRoot = Opt.Root name
280 destructHelp (Opt.Index Index.third newRoot) c =<<
281 destructHelp (Opt.Index Index.second newRoot) b =<<
282 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
284 Can.PList [] ->
285 pure revDs
287 Can.PList (hd:tl) ->
288 destructTwo path hd (A.At region (Can.PList tl)) revDs
290 Can.PCons hd tl ->
291 destructTwo path hd tl revDs
293 Can.PChr _ ->
294 pure revDs
296 Can.PStr _ ->
297 pure revDs
299 Can.PInt _ ->
300 pure revDs
302 Can.PBool _ _ ->
303 pure revDs
305 Can.PCtor _ _ (Can.Union _ _ _ opts) _ _ args ->
306 case args of
307 [Can.PatternCtorArg _ _ arg] ->
308 case opts of
309 Can.Normal -> destructHelp (Opt.Index Index.first path) arg revDs
310 Can.Unbox -> destructHelp (Opt.Unbox path) arg revDs
311 Can.Enum -> destructHelp (Opt.Index Index.first path) arg revDs
313 _ ->
314 case path of
315 Opt.Root _ ->
316 foldM (destructCtorArg path) revDs args
318 _ ->
319 do name <- Names.generate
320 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
323destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
324destructTwo path a b revDs =
325 case path of
326 Opt.Root _ ->
327 destructHelp (Opt.Index Index.second path) b =<<
328 destructHelp (Opt.Index Index.first path) a revDs
330 _ ->
331 do name <- Names.generate
332 let newRoot = Opt.Root name
333 destructHelp (Opt.Index Index.second newRoot) b =<<
334 destructHelp (Opt.Index Index.first newRoot) a (Opt.Destructor name path : revDs)
337destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
338destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
339 destructHelp (Opt.Index index path) arg revDs
343-- TAIL CALL
346optimizePotentialTailCallDef :: Cycle -> Can.Def -> Names.Tracker Opt.Def
347optimizePotentialTailCallDef cycle def =
348 case def of
349 Can.Def (A.At _ name) args expr ->
350 optimizePotentialTailCall cycle name args expr
352 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
353 optimizePotentialTailCall cycle name (map fst typedArgs) expr
356optimizePotentialTailCall :: Cycle -> Name.Name -> [Can.Pattern] -> Can.Expr -> Names.Tracker Opt.Def
357optimizePotentialTailCall cycle name args expr =
358 do (argNames, destructors) <- destructArgs args
359 toTailDef name argNames destructors <$>
360 optimizeTail cycle name argNames expr
363optimizeTail :: Cycle -> Name.Name -> [Name.Name] -> Can.Expr -> Names.Tracker Opt.Expr
364optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
365 case expression of
366 Can.Call func args ->
367 do oargs <- traverse (optimize cycle) args
369 let isMatchingName =
370 case A.toValue func of
371 Can.VarLocal name -> rootName == name
372 Can.VarTopLevel _ name -> rootName == name
373 _ -> False
375 if isMatchingName
376 then
377 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
378 Index.LengthMatch pairs ->
379 pure $ Opt.TailCall rootName pairs
381 Index.LengthMismatch _ _ ->
382 do ofunc <- optimize cycle func
383 pure $ Opt.Call ofunc oargs
384 else
385 do ofunc <- optimize cycle func
386 pure $ Opt.Call ofunc oargs
388 Can.If branches finally ->
389 let
390 optimizeBranch (condition, branch) =
391 (,)
392 <$> optimize cycle condition
393 <*> optimizeTail cycle rootName argNames branch
394 in
395 Opt.If
396 <$> traverse optimizeBranch branches
397 <*> optimizeTail cycle rootName argNames finally
399 Can.Let def body ->
400 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
402 Can.LetRec defs body ->
403 case defs of
404 [def] ->
405 Opt.Let
406 <$> optimizePotentialTailCallDef cycle def
407 <*> optimizeTail cycle rootName argNames body
409 _ ->
410 do obody <- optimizeTail cycle rootName argNames body
411 foldM (\bod def -> optimizeDef cycle def bod) obody defs
413 Can.LetDestruct pattern expr body ->
414 do (dname, destructors) <- destruct pattern
415 oexpr <- optimize cycle expr
416 obody <- optimizeTail cycle rootName argNames body
417 pure $
418 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
420 Can.Case expr branches ->
421 let
422 optimizeBranch root (Can.CaseBranch pattern branch) =
423 do destructors <- destructCase root pattern
424 obranch <- optimizeTail cycle rootName argNames branch
425 pure (pattern, foldr Opt.Destruct obranch destructors)
426 in
427 do temp <- Names.generate
428 oexpr <- optimize cycle expr
429 case oexpr of
430 Opt.VarLocal root ->
431 Case.optimize temp root <$> traverse (optimizeBranch root) branches
433 _ ->
434 do obranches <- traverse (optimizeBranch temp) branches
435 return $ Opt.Let (Opt.Def temp oexpr) (Case.optimize temp temp obranches)
437 _ ->
438 optimize cycle locExpr
442-- DETECT TAIL CALLS
445toTailDef :: Name.Name -> [Name.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
446toTailDef name argNames destructors body =
447 if hasTailCall body then
448 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
449 else
450 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
453hasTailCall :: Opt.Expr -> Bool
454hasTailCall expression =
455 case expression of
456 Opt.TailCall _ _ ->
457 True
459 Opt.If branches finally ->
460 hasTailCall finally || any (hasTailCall . snd) branches
462 Opt.Let _ body ->
463 hasTailCall body
465 Opt.Destruct _ body ->
466 hasTailCall body
468 Opt.Case _ _ decider jumps ->
469 decidecHasTailCall decider || any (hasTailCall . snd) jumps
471 _ ->
472 False
475decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
476decidecHasTailCall decider =
477 case decider of
478 Opt.Leaf choice ->
479 case choice of
480 Opt.Inline expr ->
481 hasTailCall expr
483 Opt.Jump _ ->
484 False
486 Opt.Chain _ success failure ->
487 decidecHasTailCall success || decidecHasTailCall failure
489 Opt.FanOut _ tests fallback ->
490 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests