on dc356df help sign in
elk/src/Optimize/Expression.hs
452 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 _ _ _ _ _ args ->
287 case path of
288 Opt.Root _ ->
289 foldM (destructCtorArg path) revDs args
291 _ ->
292 do name <- Names.generate
293 foldM (destructCtorArg (Opt.Root name)) (Opt.Destructor name path : revDs) args
296destructTwo :: Opt.Path -> Can.Pattern -> Can.Pattern -> [Opt.Destructor] -> Names.Tracker [Opt.Destructor]
297destructTwo path a b revDs =
298 case path of
299 Opt.Root _ ->
300 destructHelp (Opt.Index Index.first path) a =<<
301 destructHelp (Opt.Index Index.second path) b revDs
303 _ ->
304 do name <- Names.generate
305 let newRoot = Opt.Root name
306 destructHelp (Opt.Index Index.first newRoot) a =<<
307 destructHelp (Opt.Index Index.second newRoot) b (Opt.Destructor name path : revDs)
310destructCtorArg :: Opt.Path -> [Opt.Destructor] -> Can.PatternCtorArg -> Names.Tracker [Opt.Destructor]
311destructCtorArg path revDs (Can.PatternCtorArg index _ arg) =
312 destructHelp (Opt.Index index path) arg revDs
316-- TAIL CALL
319optimizePotentialTailCall :: Cycle -> Can.Def -> Names.Tracker Opt.Def
320optimizePotentialTailCall cycle def =
321 case def of
322 Can.Def (A.At _ name) args expr ->
323 do (argNames, destructors) <- destructArgs args
324 toTailDef name argNames destructors <$>
325 optimizeTail cycle name argNames expr
327 Can.TypedDef (A.At _ name) _ typedArgs expr _ ->
328 do (argNames, destructors) <- destructArgs (map fst typedArgs)
329 toTailDef name argNames destructors <$>
330 optimizeTail cycle name argNames expr
333optimizeTail :: Cycle -> N.Name -> [N.Name] -> Can.Expr -> Names.Tracker Opt.Expr
334optimizeTail cycle rootName argNames locExpr@(A.At _ expression) =
335 case expression of
336 Can.Call func@(A.At _ (Can.VarTopLevel _ name)) args ->
337 do oargs <- traverse (optimize cycle) args
338 if name == rootName
339 then
340 case Index.indexedZipWith (\_ a b -> (a,b)) argNames oargs of
341 Index.LengthMatch pairs ->
342 pure $ Opt.TailCall name pairs
344 Index.LengthMismatch _ _ ->
345 do ofunc <- optimize cycle func
346 pure $ Opt.Call ofunc oargs
347 else
348 do ofunc <- optimize cycle func
349 pure $ Opt.Call ofunc oargs
351 Can.If branches finally ->
352 let
353 optimizeBranch (condition, branch) =
354 (,)
355 <$> optimize cycle condition
356 <*> optimizeTail cycle rootName argNames branch
357 in
358 Opt.If
359 <$> traverse optimizeBranch branches
360 <*> optimizeTail cycle rootName argNames finally
362 Can.Let def body ->
363 optimizeDef cycle def =<< optimizeTail cycle rootName argNames body
365 Can.LetRec defs body ->
366 case defs of
367 [def] ->
368 Opt.Let
369 <$> optimizePotentialTailCall cycle def
370 <*> optimizeTail cycle rootName argNames body
372 _ ->
373 do obody <- optimizeTail cycle rootName argNames body
374 foldM (\bod def -> optimizeDef cycle def bod) obody defs
376 Can.LetDestruct pattern expr body ->
377 do (dname, destructors) <- destruct pattern
378 oexpr <- optimize cycle expr
379 obody <- optimizeTail cycle dname argNames body
380 pure $
381 Opt.Let (Opt.Def dname oexpr) (foldr Opt.Destruct obody destructors)
383 Can.Case expr branches ->
384 let
385 optimizeBranch root (Can.CaseBranch pattern branch) =
386 do destructors <- destructHelp (Opt.Root root) pattern []
387 obranch <- optimizeTail cycle rootName argNames branch
388 pure (pattern, foldr Opt.Destruct obranch destructors)
389 in
390 do temp <- Names.generate
391 oexpr <- optimize cycle expr
392 case oexpr of
393 Opt.VarLocal root ->
394 Case.optimize temp root <$> traverse (optimizeBranch root) branches
396 _ ->
397 Case.optimize temp temp <$> traverse (optimizeBranch temp) branches
399 _ ->
400 optimize cycle locExpr
404-- DETECT TAIL CALLS
407toTailDef :: N.Name -> [N.Name] -> [Opt.Destructor] -> Opt.Expr -> Opt.Def
408toTailDef name argNames destructors body =
409 if hasTailCall body then
410 Opt.TailDef name argNames (foldr Opt.Destruct body destructors)
411 else
412 Opt.Def name (Opt.Function argNames (foldr Opt.Destruct body destructors))
415hasTailCall :: Opt.Expr -> Bool
416hasTailCall expression =
417 case expression of
418 Opt.TailCall _ _ ->
419 True
421 Opt.If branches finally ->
422 hasTailCall finally || any (hasTailCall . snd) branches
424 Opt.Let _ body ->
425 hasTailCall body
427 Opt.Destruct _ body ->
428 hasTailCall body
430 Opt.Case _ _ decider jumps ->
431 decidecHasTailCall decider || any (hasTailCall . snd) jumps
433 _ ->
434 False
437decidecHasTailCall :: Opt.Decider Opt.Choice -> Bool
438decidecHasTailCall decider =
439 case decider of
440 Opt.Leaf choice ->
441 case choice of
442 Opt.Inline expr ->
443 hasTailCall expr
445 Opt.Jump _ ->
446 False
448 Opt.Chain _ success failure ->
449 decidecHasTailCall success || decidecHasTailCall failure
451 Opt.FanOut _ tests fallback ->
452 decidecHasTailCall fallback || any (decidecHasTailCall . snd) tests