git.iko.soy/elk
on f4a182c help sign in
elk/src/Optimize.hs
587 lines 15 kB history raw
1module Optimize (optimize) where
3import Control.Applicative
4import Control.Arrow (second)
5import qualified Control.Monad as M
6import qualified Control.Monad.State as State
7import qualified Data.Map as Map
8import Data.Map ((!))
9import qualified Data.Maybe as Maybe
10import qualified Data.Traversable as T
12import qualified AST.Expression.General as Expr
13import qualified AST.Expression.Canonical as Can
14import qualified AST.Expression.Optimized as Opt
15import qualified AST.Module as Module
16import qualified AST.Module.Name as ModuleName
17import qualified AST.Pattern as P
18import qualified AST.Variable as Var
19import Elm.Utils ((|>))
20import qualified Optimize.Cases as Cases
21import qualified Reporting.Annotation as A
22import qualified Reporting.Region as R
25-- OPTIMIZE - turn a canonical module into a list of optimized defs
27optimize :: Module.CanonicalModule -> Module.Optimized
28optimize module_ =
29 let
30 home =
31 Module.name module_
33 (defs, _) =
34 flattenLets [] (Module.program (Module.body module_))
35 in
36 State.evalState
37 (concat <$> mapM (optimizeDef (Just home)) defs)
38 (Env False 0)
41flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
42flattenLets defs annExpr@(A.A _ expr) =
43 case expr of
44 Expr.Let ds body ->
45 flattenLets (defs ++ ds) body
47 _ ->
48 (defs, annExpr)
51-- TRACKER TO THREAD FINDINGS THROUGH TRAVERSAL
53type Optimizer a =
54 State.State Env a
57data Env = Env
58 { _hasTailCall :: Bool
59 , _uid :: Int
60 }
63setTailCall :: Bool -> Optimizer ()
64setTailCall bool =
65 do uid <- State.gets _uid
66 State.put (Env bool uid)
69mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
70mapSnd func (x, a) =
71 (,) x <$> func a
74-- CONVERT DEFINITIONS
76optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Optimizer [Opt.Def]
77optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
78 let
79 (args, rawBody) =
80 Expr.collectLambdas expression
81 in
82 case (ptrn, args) of
83 (P.Var name, _ : _) ->
85 do (Env htc uid) <- State.get
86 State.put (Env False uid)
88 (args, body) <- optimizeFunction (Just name) args rawBody
90 (Env isTailCall uid') <- State.get
91 State.put (Env htc uid')
93 let facts = Opt.Facts home deps
95 return $
96 if isTailCall then
98 [ Opt.TailDef facts name args body ]
100 else
102 [ Opt.Def facts name (Opt.Function args body) ]
104 (P.Var name, []) ->
106 do hasTailCall <- State.gets _hasTailCall
107 body <- optimizeExpr Nothing rawBody
108 setTailCall hasTailCall
109 return [ Opt.Def (Opt.Facts home deps) name body ]
111 (_, []) ->
113 do name <- freshName
114 optBody <- optimizeExpr Nothing rawBody
115 let optDef = Opt.Def (Opt.Facts home deps) name optBody
116 return (optDef : patternToDefs home name pattern)
118 _ ->
119 error "there should never be a function where the name is not a P.Var"
122-- OPTIMIZE FUNCTION
124optimizeFunction
125 :: Maybe String
126 -> [P.CanonicalPattern]
127 -> Can.Expr
128 -> Optimizer ([String], Opt.Expr)
129optimizeFunction maybeName patterns givenBody =
130 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
131 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
132 return (argNames, optBody)
135makeContext :: [String] -> String -> Context
136makeContext argNames name =
137 Just (name, length argNames, argNames)
140depatternArgs
141 :: ([String], Can.Expr)
142 -> P.CanonicalPattern
143 -> Optimizer ([String], Can.Expr)
144depatternArgs (names, rawExpr) ptrn =
145 do (name, expr) <- depattern ptrn rawExpr
146 return (name : names, expr)
149patternToDefs
150 :: Maybe ModuleName.Canonical
151 -> String
152 -> P.CanonicalPattern
153 -> [Opt.Def]
154patternToDefs home rootName pattern =
155 let
156 (deps, rootExpr) =
157 case home of
158 Nothing ->
159 ( []
160 , Opt.Var (Var.local rootName)
161 )
163 Just moduleName ->
164 ( [ Var.TopLevelVar moduleName rootName ]
165 , Opt.Var (Var.topLevel moduleName rootName)
166 )
168 toDef (name, accessExpr) =
169 Opt.Def (Opt.Facts home deps) name accessExpr
170 in
171 map toDef (substitutions rootExpr pattern)
174substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
175substitutions expr (A.A _ pattern) =
176 case pattern of
177 P.Var name ->
178 [ (name, expr) ]
180 P.Literal _ ->
181 []
183 P.Record fields ->
184 map (\name -> (name, Opt.Access expr name)) fields
186 P.Anything ->
187 []
189 P.Alias alias realPattern ->
190 (alias, expr) : substitutions expr realPattern
192 P.Data _ patterns ->
193 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
196-- OPTIMIZE EXPRESSIONS
198type Context = Maybe (String, Int, [String])
201{-| This optimizer detects tail calls. The `Context` holds a function name,
202function arity, and list of argument names. If a tail call is possible, the
203`Context` will be filled in and we check any function we run into. If we go
204down a route that cannot possibly have tail calls, we remove the context.
205-}
206optimizeExpr :: Context -> Can.Expr -> Optimizer Opt.Expr
207optimizeExpr context annExpr@(A.A region expression) =
208 let
209 keepLooking =
210 optimizeExpr context
211 -- this *might* find a tail call, use this when something is a potential
212 -- tail call
214 justConvert =
215 optimizeExpr Nothing
216 -- this will never find tail calls, use this when you are working with
217 -- expressions that cannot be turned into tail calls
218 in
219 case expression of
220 Expr.Literal lit ->
221 pure (Opt.Literal lit)
223 Expr.Var name ->
224 pure (Opt.Var name)
226 Expr.Range lowExpr highExpr ->
227 Opt.Range
228 <$> justConvert lowExpr
229 <*> justConvert highExpr
231 Expr.ExplicitList elements ->
232 Opt.ExplicitList
233 <$> T.traverse justConvert elements
235 Expr.Binop op leftExpr rightExpr ->
236 optimizeBinop context region op leftExpr rightExpr
238 Expr.Lambda _ _ ->
239 let
240 (patterns, body) =
241 Expr.collectLambdas annExpr
242 in
243 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
245 Expr.App _ _ ->
246 let
247 (func:args) =
248 Expr.collectApps annExpr
249 in
250 case isTailCall context func (length args) of
251 Just (name, argNames) ->
252 do optArgs <- T.traverse justConvert args
253 setTailCall True
254 return (Opt.TailCall name argNames optArgs)
256 Nothing ->
257 do hasTailCall <- State.gets _hasTailCall
258 optFunc <- justConvert func
259 optArgs <- T.traverse justConvert args
260 setTailCall hasTailCall
261 return (Opt.Call optFunc optArgs)
264 Expr.If branches finally ->
265 let
266 crawlBranch (cond,branch) =
267 (,) <$> justConvert cond <*> keepLooking branch
268 in
269 Opt.If
270 <$> T.traverse crawlBranch branches
271 <*> keepLooking finally
273 Expr.Let defs body ->
274 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
275 Opt.Let optDefs <$> keepLooking body
277 Expr.Case expr branches ->
278 optimizeCase context region expr branches
280 Expr.Data name exprs ->
281 Opt.Data name <$> T.traverse justConvert exprs
283 Expr.Access record field ->
284 Opt.Access
285 <$> justConvert record
286 <*> pure field
288 Expr.Update record fields ->
289 Opt.Update
290 <$> justConvert record
291 <*> T.traverse (mapSnd justConvert) fields
293 Expr.Record fields ->
294 Opt.Record
295 <$> T.traverse (mapSnd justConvert) fields
297 Expr.GLShader uid src gltipe ->
298 pure (Opt.GLShader uid src gltipe)
300 Expr.Port impl ->
301 Opt.Port <$>
302 case impl of
303 Expr.In name tipe ->
304 pure (Expr.In name tipe)
306 Expr.Out name expr tipe ->
307 (\e -> Expr.Out name e tipe) <$> justConvert expr
309 Expr.Task name expr tipe ->
310 (\e -> Expr.Task name e tipe) <$> justConvert expr
313-- DETECT TAIL CALL
315isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
316isTailCall context (A.A _ function) arity =
317 case (context, function) of
318 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
320 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
321 Just (name, argNames)
323 else
324 Nothing
326 _ ->
327 Nothing
331-- DEPATTERN
332-- given a pattern and an expression, push the actual pattern matching into the
333-- expression as a case-expression if necessary. The pattern compiler will work
334-- it out from there
336depattern :: P.CanonicalPattern -> Can.Expr -> Optimizer (String, Can.Expr)
337depattern canPattern@(A.A _ pattern) expr =
338 let
339 ann =
340 A.A (error "the annotation added in 'depattern' should not be observed!")
342 caseExpr e branches =
343 ann (Expr.Case e branches)
345 var name =
346 ann (Expr.Var (Var.local name))
348 caseify =
349 do name <- freshName
350 return
351 ( name
352 , caseExpr (var name) [ (canPattern, expr) ]
353 )
354 in
355 case pattern of
356 P.Var name ->
357 return (name, expr)
359 P.Literal _ ->
360 caseify
362 P.Record _ ->
363 caseify
365 P.Anything ->
366 do name <- freshName
367 return (name, expr)
369 P.Alias _ _ ->
370 caseify
372 P.Data _ _ ->
373 caseify
376freshName :: Optimizer String
377freshName =
378 do (Env htc uid) <- State.get
379 State.put (Env htc (uid + 1))
380 return ("_p" ++ show uid)
383-- OPTIMIZE CASES
385optimizeCase
386 :: Context
387 -> R.Region
388 -> Can.Expr
389 -> [(P.CanonicalPattern, Can.Expr)]
390 -> Optimizer Opt.Expr
391optimizeCase context region expr rawBranches =
392 let
393 indexify index (pattern, branch) =
394 ( (pattern, index)
395 , (index, branch)
396 )
397 in
398 do optBranches <- T.traverse (mapSnd (optimizeExpr context)) rawBranches
399 crashBranches <-
400 case last optBranches of
401 (A.A ann P.Anything, Opt.Call (Opt.Crash _ _) [arg]) ->
402 do name <- freshName
403 let newPattern = A.A ann (P.Var name)
404 let newCall = Opt.Call (Opt.Crash region (Just name)) [arg]
405 return (init optBranches ++ [(newPattern, newCall)])
407 _ ->
408 return optBranches
410 let (patterns, indexedBranches) =
411 unzip (zipWith indexify [0..] crashBranches)
413 let decisionTree =
414 Cases.compile (error "need a VariantDict") patterns
416 optExpr <- optimizeExpr Nothing expr
418 return (inlinedCase optExpr decisionTree indexedBranches)
421inlinedCase
422 :: Opt.Expr
423 -> Cases.DecisionTree Cases.Jump
424 -> [(Int, Opt.Expr)]
425 -> Opt.Expr
426inlinedCase expr decisionTree jumpTargets =
427 let
428 substitutionDict =
429 gatherSubstitutions decisionTree
431 (jumps, sharedBranches) =
432 unzip (map (toJumps substitutionDict) jumpTargets)
433 in
434 Opt.Case
435 expr
436 (inlineJumps (Map.fromList jumps) decisionTree)
437 (Maybe.catMaybes sharedBranches)
440gatherSubstitutions
441 :: Cases.DecisionTree Cases.Jump
442 -> Map.Map Int [[(String, Cases.Path)]]
443gatherSubstitutions decisionTree =
444 case decisionTree of
445 Cases.Decision _ edges fallback ->
446 map snd edges
447 |> maybe id (:) fallback
448 |> map gatherSubstitutions
449 |> Map.unionsWith (++)
451 Cases.Match (Cases.Jump target substitutions) ->
452 Map.singleton target [substitutions]
455toJumps
456 :: Map.Map Int [[(String, Cases.Path)]]
457 -> (Int, Opt.Expr)
458 -> ( (Int, Opt.Jump), Maybe (Int, Opt.Branch) )
459toJumps substitutionDict (target, branch) =
460 case substitutionDict ! target of
461 [substitutions] ->
462 ( (target, Opt.Inline (Opt.Branch substitutions branch))
463 , Nothing
464 )
466 substitutions : _ ->
467 ( (target, Opt.Jump target)
468 , Just (target, Opt.Branch substitutions branch)
469 )
471 _ ->
472 error "it is not possible for there to be no substitutions"
475inlineJumps
476 :: Map.Map Int Opt.Jump
477 -> Cases.DecisionTree Cases.Jump
478 -> Cases.DecisionTree Opt.Jump
479inlineJumps jumpDict decisionTree =
480 let
481 go =
482 inlineJumps jumpDict
483 in
484 case decisionTree of
485 Cases.Decision test edges fallback ->
486 Cases.Decision test (map (second go) edges) (fmap go fallback)
488 Cases.Match (Cases.Jump target _) ->
489 Cases.Match (jumpDict ! target)
492-- OPTIMIZE BINOPS
494optimizeBinop
495 :: Context
496 -> R.Region
497 -> Var.Canonical
498 -> Can.Expr
499 -> Can.Expr
500 -> Optimizer Opt.Expr
501optimizeBinop context region op leftExpr rightExpr =
502 let
503 ann = A.A region
504 binop = ann (Expr.Binop op leftExpr rightExpr)
505 in
506 if op == forwardApply then
508 optimizeExpr context (collect id left forwardApply binop)
510 else if op == backwardApply then
512 optimizeExpr context (collect id right backwardApply binop)
514 else if op == forwardCompose then
516 do var <- freshName
517 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
518 let body = collect makeRoot left forwardCompose binop
519 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
521 else if op == backwardCompose then
523 do var <- freshName
524 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
525 let body = collect makeRoot right backwardCompose binop
526 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
528 else if op == Var.Canonical Var.BuiltIn "::" then
530 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
532 else
534 Opt.Binop op
535 <$> optimizeExpr Nothing leftExpr
536 <*> optimizeExpr Nothing rightExpr
539-- left-associative ((x |> f) |> g)
540forwardApply :: Var.Canonical
541forwardApply =
542 Var.inCore ["Basics"] "|>"
545-- right-associative (g <| (f <| x))
546backwardApply :: Var.Canonical
547backwardApply =
548 Var.inCore ["Basics"] "<|"
551-- left-associative ((f >> g) >> h)
552forwardCompose :: Var.Canonical
553forwardCompose =
554 Var.inCore ["Basics"] ">>"
557-- right-associative (h << (g << f))
558backwardCompose :: Var.Canonical
559backwardCompose =
560 Var.inCore ["Basics"] "<<"
563type Assoc =
564 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
567left :: Assoc
568left =
569 flip (,)
572right :: Assoc
573right =
574 (,)
577collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
578collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
579 case expr of
580 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
581 let
582 (func, arg) = assoc leftExpr rightExpr
583 in
584 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
586 _ ->
587 makeRoot annExpr