git.iko.soy/elk
on 5fcbf86 help sign in
elk/src/Optimize.hs
570 lines 15 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize (optimize) where
5import Control.Arrow (second)
6import qualified Control.Monad as M
7import Data.Monoid ((<>))
8import qualified Data.Map as Map
9import qualified Data.Set as Set
10import qualified Data.Text as Text
11import Data.Text (Text)
13import qualified AST.Expression.Canonical as Can
14import qualified AST.Expression.Optimized as Opt
15import qualified AST.Module.Name as ModuleName
16import qualified AST.Pattern as P
17import qualified AST.Variable as Var
18import qualified Optimize.Case as Case
19import qualified Optimize.DecisionTree as DT
20import qualified Optimize.Environment as Env
21import qualified Optimize.Inline as Inline
22import qualified Optimize.Port as Port
23import qualified Reporting.Annotation as A
24import qualified Reporting.Region as R
28-- OPTIMIZE
31optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [(Text, Opt.Decl)]
32optimize variantDict home defs =
33 concat $ zipWith (optimizeDecl variantDict home) [1 .. length defs] defs
37-- OPTIMIZE DECLS
40optimizeDecl :: DT.VariantDict -> ModuleName.Canonical -> Int -> Can.Def -> [(Text, Opt.Decl)]
41optimizeDecl variantDict home index (Can.Def _ pattern@(A.A _ ptrn) expression _) =
42 let
43 (rawArgs, rawBody) =
44 Can.collectLambdas expression
45 in
46 case (ptrn, rawArgs) of
47 (P.Var name, _) ->
48 (:[]) $ Env.run variantDict home $
49 optimizeNamedDef name rawArgs rawBody
51 (_, []) ->
52 let
53 name = "_" <> Text.pack (show index)
54 var = Var.Global home name
55 deps = Set.singleton var
56 decl =
57 Env.run variantDict home $
58 ((,) name . Opt.Def) <$> optimizeExpr Nothing rawBody
59 destructors =
60 destruct (Opt.VarGlobal var) pattern
61 in
62 decl : map (second (Opt.Decl deps Set.empty Nothing)) destructors
64 _ ->
65 error "bug manifesting in Optimize.optimizeDecl, please report"
69-- OPTIMIZE DEFS
72optimizeLetDef :: Can.Def -> Env.Optimizer [(Text, Opt.Def)]
73optimizeLetDef (Can.Def _ pattern@(A.A _ ptrn) expression _) =
74 let
75 (rawArgs, rawBody) =
76 Can.collectLambdas expression
77 in
78 case (ptrn, rawArgs) of
79 (P.Var name, _) ->
80 do def <- optimizeNamedDef name rawArgs rawBody
81 return [def]
83 (_, []) ->
84 do name <- Env.freshName
85 body <- optimizeExpr Nothing rawBody
86 return $ (name, Opt.Def body) : destruct (Opt.VarLocal name) pattern
88 _ ->
89 error "bug manifesting in Optimize.optimizeLetDef, please report"
92optimizeNamedDef :: Text -> [P.Canonical] -> Can.Expr -> Env.Optimizer (Text, Opt.Def)
93optimizeNamedDef name rawArgs rawBody =
94 case rawArgs of
95 [] ->
96 do hasTailCall <- Env.getTailCall
97 body <- optimizeExpr Nothing rawBody
98 Env.setTailCall hasTailCall
99 return (name, Opt.Def body)
101 _ : _ ->
102 do htc <- Env.getTailCall
103 Env.setTailCall False
105 (args, body) <- optimizeFunction (Just name) rawArgs rawBody
107 isTailCall <- Env.getTailCall
108 Env.setTailCall htc
110 return $ (,) name $
111 if isTailCall then
112 Opt.TailDef args body
113 else
114 Opt.Def (Opt.Function args body)
117destruct :: Opt.Expr -> P.Canonical -> [(Text, Opt.Def)]
118destruct root pattern =
119 map (second Opt.Def) $
120 patternToSubstitutions root pattern
124-- OPTIMIZE FUNCTION
127optimizeFunction
128 :: Maybe Text
129 -> [P.Canonical]
130 -> Can.Expr
131 -> Env.Optimizer ([Text], Opt.Expr)
132optimizeFunction maybeName patterns givenBody =
133 Env.indirectly $
134 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
135 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
136 return (argNames, optBody)
139makeContext :: [Text] -> Text -> Context
140makeContext argNames name =
141 Just (name, length argNames, argNames)
144depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
145depatternArgs (names, rawExpr) ptrn =
146 do (name, expr) <- depattern ptrn rawExpr
147 return (name : names, expr)
151-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
154patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
155patternToSubstitutions expr (A.A _ pattern) =
156 case pattern of
157 P.Var name ->
158 [ (name, expr) ]
160 P.Literal _ ->
161 []
163 P.Record fields ->
164 map (\name -> (name, Opt.Access expr name)) fields
166 P.Anything ->
167 []
169 P.Alias alias realPattern ->
170 (alias, expr) : patternToSubstitutions expr realPattern
172 P.Ctor _ patterns ->
173 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
177-- OPTIMIZE EXPRESSIONS
180type Context = Maybe (Text, Int, [Text])
183{-| This optimizer detects tail calls. The `Context` holds a function name,
184function arity, and list of argument names. If a tail call is possible, the
185`Context` will be filled in and we check any function we run into. If we go
186down a route that cannot possibly have tail calls, we remove the context.
187-}
188optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
189optimizeExpr context annExpr@(A.A region expression) =
190 let
191 keepLooking =
192 optimizeExpr context
193 -- this *might* find a tail call, use this when something is a potential
194 -- tail call
196 justConvert =
197 optimizeExpr Nothing
198 -- this will never find tail calls, use this when you are working with
199 -- expressions that cannot be turned into tail calls
200 in
201 case expression of
202 Can.Literal lit ->
203 pure (Opt.Literal lit)
205 Can.Var var ->
206 optimizeVariable region var
208 Can.List elements ->
209 Opt.List
210 <$> traverse justConvert elements
212 Can.Binop op leftExpr rightExpr ->
213 optimizeBinop context region op leftExpr rightExpr
215 Can.Lambda _ _ ->
216 let
217 (patterns, body) =
218 Can.collectLambdas annExpr
219 in
220 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
222 Can.App _ _ ->
223 let
224 (func:args) =
225 Can.collectApps annExpr
226 in
227 case getTailCallInfo context func (length args) of
228 Just (name, argNames) ->
229 do optArgs <- traverse justConvert args
230 Env.setTailCall True
231 return (Opt.TailCall name argNames optArgs)
233 Nothing ->
234 do hasTailCall <- Env.getTailCall
235 optFunc <- justConvert func
236 optArgs <- traverse justConvert args
237 Env.setTailCall hasTailCall
238 return (Opt.Call optFunc optArgs)
240 Can.If branches finally ->
241 let
242 crawlBranch (cond,branch) =
243 (,) <$> justConvert cond <*> keepLooking branch
244 in
245 Opt.If
246 <$> traverse crawlBranch branches
247 <*> keepLooking finally
249 Can.Let defs body ->
250 do optDefs <- concat <$> traverse optimizeLetDef defs
251 Opt.Let optDefs <$> keepLooking body
253 Can.Case expr branches ->
254 do optExpr <- optimizeExpr Nothing expr
255 variantDict <- Env.getVariantDict
257 name <- Env.freshName
258 optBranches <- traverse (optimizeBranch context region name) branches
259 let optCase = Case.optimize variantDict name optBranches
260 return $ Opt.Let [ (name, Opt.Def optExpr) ] optCase
262 Can.Ctor (Var.Canonical _ name) exprs ->
263 Opt.Ctor name <$> traverse justConvert exprs
265 Can.Access record field ->
266 Opt.Access
267 <$> justConvert record
268 <*> pure field
270 Can.Update record fields ->
271 Opt.Update
272 <$> justConvert record
273 <*> traverse (traverse justConvert) fields
275 Can.Record fields ->
276 Opt.Record
277 <$> traverse (traverse justConvert) fields
279 Can.Cmd home manager ->
280 Opt.Cmd home <$> Env.registerEffects home manager
282 Can.Sub home manager ->
283 Opt.Sub home <$> Env.registerEffects home manager
285 Can.OutgoingPort name tipe ->
286 Opt.OutgoingPort name <$> Port.toEncoder tipe
288 Can.IncomingPort name tipe ->
289 Opt.IncomingPort name <$> Port.toDecoder tipe
291 Can.Program main expr ->
292 Opt.Program <$> optimizeMain main <*> justConvert expr
294 Can.SaveEnv _ _ ->
295 error "save_the_environment should never make it to optimization phase"
297 Can.GLShader _ src _ ->
298 pure (Opt.GLShader src)
302-- OPTIMIZE VARIABLE
305optimizeVariable :: R.Region -> Var.Canonical -> Env.Optimizer Opt.Expr
306optimizeVariable region (Var.Canonical home name) =
307 case home of
308 Var.BuiltIn ->
309 pure (Opt.VarLocal name)
311 Var.Local ->
312 pure (Opt.VarLocal name)
314 Var.TopLevel modul ->
315 do let var = Var.Global modul name
316 Env.register var
317 pure (Opt.VarGlobal var)
319 Var.Module modul ->
320 do let var = Var.Global modul name
321 Env.register var
322 if name == "crash" && modul == ModuleName.inCore "Debug"
323 then
324 do here <- Env.getHome
325 pure (Opt.Crash here region Nothing)
326 else
327 pure (Opt.VarGlobal var)
331-- OPTIMIZE MAIN
334optimizeMain :: Can.Main -> Env.Optimizer Opt.Main
335optimizeMain main =
336 case main of
337 Can.Static ->
338 return Opt.Static
340 Can.Dynamic tipe ->
341 Opt.Dynamic <$> Port.toFlagsDecoder tipe
345-- OPTIMIZE CASE BRANCHES
348optimizeBranch
349 :: Context
350 -> R.Region
351 -> Text
352 -> (P.Canonical, Can.Expr)
353 -> Env.Optimizer (P.Canonical, Opt.Expr)
354optimizeBranch context region exprName (rawPattern, canBranch) =
355 do optBranch <- optimizeExpr context canBranch
356 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
358 let substitutions = Map.fromList $
359 patternToSubstitutions (Opt.VarLocal exprName) pattern
361 (,) pattern <$> Inline.inline substitutions branch
364tagCrashBranch
365 :: R.Region
366 -> P.Canonical
367 -> Opt.Expr
368 -> Env.Optimizer (P.Canonical, Opt.Expr)
369tagCrashBranch region pattern@(A.A pr _) expr =
370 case expr of
371 Opt.Call (Opt.Crash home _ _) [arg] ->
372 do name <- Env.freshName
373 return
374 ( A.A pr (P.Alias name pattern)
375 , Opt.Call (Opt.Crash home region (Just (Opt.VarLocal name))) [arg]
376 )
378 _ ->
379 return (pattern, expr)
383-- DETECT TAIL CALL
386getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
387getTailCallInfo context (A.A _ function) arity =
388 case (context, function) of
389 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
391 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
392 Just (name, argNames)
394 else
395 Nothing
397 _ ->
398 Nothing
402-- DEPATTERN
403-- given a pattern and an expression, push the actual pattern matching into the
404-- expression as a case-expression if necessary. The pattern compiler will work
405-- it out from there
408depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
409depattern canPattern@(A.A _ pattern) expr =
410 let
411 ann =
412 A.A (error "the annotation added in 'depattern' should not be observed!")
414 caseExpr e branches =
415 ann (Can.Case e branches)
417 var name =
418 ann (Can.Var (Var.local name))
420 caseify =
421 do name <- Env.freshName
422 return
423 ( name
424 , caseExpr (var name) [ (canPattern, expr) ]
425 )
426 in
427 case pattern of
428 P.Var name ->
429 return (name, expr)
431 P.Literal _ ->
432 caseify
434 P.Record _ ->
435 caseify
437 P.Anything ->
438 do name <- Env.freshName
439 return (name, expr)
441 P.Alias _ _ ->
442 caseify
444 P.Ctor _ _ ->
445 caseify
449-- OPTIMIZE BINOPS
452optimizeBinop
453 :: Context
454 -> R.Region
455 -> Var.Canonical
456 -> Can.Expr
457 -> Can.Expr
458 -> Env.Optimizer Opt.Expr
459optimizeBinop context region op leftExpr rightExpr =
460 let
461 ann = A.A region
462 binop = ann (Can.Binop op leftExpr rightExpr)
463 in
464 if op == forwardApply then
466 optimizeExpr context (collect id left forwardApply binop)
468 else if op == backwardApply then
470 optimizeExpr context (collect id right backwardApply binop)
472 else if op == forwardCompose then
474 do var <- Env.freshName
475 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
476 let body = collect makeRoot left forwardCompose binop
477 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
479 else if op == backwardCompose then
481 do var <- Env.freshName
482 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
483 let body = collect makeRoot right backwardCompose binop
484 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
486 else if op == cons then
488 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
490 else
492 do optLeft <- optimizeExpr Nothing leftExpr
493 optRight <- optimizeExpr Nothing rightExpr
494 makeOptBinop op optLeft optRight
497-- left-associative ((x |> f) |> g)
498forwardApply :: Var.Canonical
499forwardApply =
500 Var.inCore "Basics" "|>"
503-- right-associative (g <| (f <| x))
504backwardApply :: Var.Canonical
505backwardApply =
506 Var.inCore "Basics" "<|"
509-- left-associative ((f >> g) >> h)
510forwardCompose :: Var.Canonical
511forwardCompose =
512 Var.inCore "Basics" ">>"
515-- right-associative (h << (g << f))
516backwardCompose :: Var.Canonical
517backwardCompose =
518 Var.inCore "Basics" "<<"
521cons :: Var.Canonical
522cons =
523 Var.Canonical Var.BuiltIn "::"
526type Assoc =
527 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
530left :: Assoc
531left =
532 flip (,)
535right :: Assoc
536right =
537 (,)
540collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
541collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
542 case expr of
543 Can.Binop op leftExpr rightExpr | op == desiredOp ->
544 let
545 (func, arg) = assoc leftExpr rightExpr
546 in
547 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
549 _ ->
550 makeRoot annExpr
553makeOptBinop :: Var.Canonical -> Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
554makeOptBinop (Var.Canonical home op) leftExpr rightExpr =
555 case home of
556 Var.Local ->
557 error "bug manifesting in Optimize.makeOptBinop"
559 Var.BuiltIn ->
560 error "bug manifesting in Optimize.makeOptBinop"
562 Var.Module name ->
563 do let var = Var.Global name op
564 Env.register var
565 return $ Opt.Binop var leftExpr rightExpr
567 Var.TopLevel name ->
568 do let var = Var.Global name op
569 Env.register var
570 return $ Opt.Binop var leftExpr rightExpr