git.iko.soy/elk
on bca019d help sign in
elk/src/Optimize.hs
573 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.VDom ->
338 return Opt.VDom
340 Can.NoFlags ->
341 return Opt.NoFlags
343 Can.Flags tipe ->
344 Opt.Flags <$> Port.toDecoder tipe
348-- OPTIMIZE CASE BRANCHES
351optimizeBranch
352 :: Context
353 -> R.Region
354 -> Text
355 -> (P.Canonical, Can.Expr)
356 -> Env.Optimizer (P.Canonical, Opt.Expr)
357optimizeBranch context region exprName (rawPattern, canBranch) =
358 do optBranch <- optimizeExpr context canBranch
359 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
361 let substitutions = Map.fromList $
362 patternToSubstitutions (Opt.VarLocal exprName) pattern
364 (,) pattern <$> Inline.inline substitutions branch
367tagCrashBranch
368 :: R.Region
369 -> P.Canonical
370 -> Opt.Expr
371 -> Env.Optimizer (P.Canonical, Opt.Expr)
372tagCrashBranch region pattern@(A.A pr _) expr =
373 case expr of
374 Opt.Call (Opt.Crash home _ _) [arg] ->
375 do name <- Env.freshName
376 return
377 ( A.A pr (P.Alias name pattern)
378 , Opt.Call (Opt.Crash home region (Just (Opt.VarLocal name))) [arg]
379 )
381 _ ->
382 return (pattern, expr)
386-- DETECT TAIL CALL
389getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
390getTailCallInfo context (A.A _ function) arity =
391 case (context, function) of
392 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
394 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
395 Just (name, argNames)
397 else
398 Nothing
400 _ ->
401 Nothing
405-- DEPATTERN
406-- given a pattern and an expression, push the actual pattern matching into the
407-- expression as a case-expression if necessary. The pattern compiler will work
408-- it out from there
411depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
412depattern canPattern@(A.A _ pattern) expr =
413 let
414 ann =
415 A.A (error "the annotation added in 'depattern' should not be observed!")
417 caseExpr e branches =
418 ann (Can.Case e branches)
420 var name =
421 ann (Can.Var (Var.local name))
423 caseify =
424 do name <- Env.freshName
425 return
426 ( name
427 , caseExpr (var name) [ (canPattern, expr) ]
428 )
429 in
430 case pattern of
431 P.Var name ->
432 return (name, expr)
434 P.Literal _ ->
435 caseify
437 P.Record _ ->
438 caseify
440 P.Anything ->
441 do name <- Env.freshName
442 return (name, expr)
444 P.Alias _ _ ->
445 caseify
447 P.Ctor _ _ ->
448 caseify
452-- OPTIMIZE BINOPS
455optimizeBinop
456 :: Context
457 -> R.Region
458 -> Var.Canonical
459 -> Can.Expr
460 -> Can.Expr
461 -> Env.Optimizer Opt.Expr
462optimizeBinop context region op leftExpr rightExpr =
463 let
464 ann = A.A region
465 binop = ann (Can.Binop op leftExpr rightExpr)
466 in
467 if op == forwardApply then
469 optimizeExpr context (collect id left forwardApply binop)
471 else if op == backwardApply then
473 optimizeExpr context (collect id right backwardApply binop)
475 else if op == forwardCompose then
477 do var <- Env.freshName
478 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
479 let body = collect makeRoot left forwardCompose binop
480 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
482 else if op == backwardCompose then
484 do var <- Env.freshName
485 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
486 let body = collect makeRoot right backwardCompose binop
487 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
489 else if op == cons then
491 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
493 else
495 do optLeft <- optimizeExpr Nothing leftExpr
496 optRight <- optimizeExpr Nothing rightExpr
497 makeOptBinop op optLeft optRight
500-- left-associative ((x |> f) |> g)
501forwardApply :: Var.Canonical
502forwardApply =
503 Var.inCore "Basics" "|>"
506-- right-associative (g <| (f <| x))
507backwardApply :: Var.Canonical
508backwardApply =
509 Var.inCore "Basics" "<|"
512-- left-associative ((f >> g) >> h)
513forwardCompose :: Var.Canonical
514forwardCompose =
515 Var.inCore "Basics" ">>"
518-- right-associative (h << (g << f))
519backwardCompose :: Var.Canonical
520backwardCompose =
521 Var.inCore "Basics" "<<"
524cons :: Var.Canonical
525cons =
526 Var.Canonical Var.BuiltIn "::"
529type Assoc =
530 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
533left :: Assoc
534left =
535 flip (,)
538right :: Assoc
539right =
540 (,)
543collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
544collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
545 case expr of
546 Can.Binop op leftExpr rightExpr | op == desiredOp ->
547 let
548 (func, arg) = assoc leftExpr rightExpr
549 in
550 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
552 _ ->
553 makeRoot annExpr
556makeOptBinop :: Var.Canonical -> Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
557makeOptBinop (Var.Canonical home op) leftExpr rightExpr =
558 case home of
559 Var.Local ->
560 error "bug manifesting in Optimize.makeOptBinop"
562 Var.BuiltIn ->
563 error "bug manifesting in Optimize.makeOptBinop"
565 Var.Module name ->
566 do let var = Var.Global name op
567 Env.register var
568 return $ Opt.Binop var leftExpr rightExpr
570 Var.TopLevel name ->
571 do let var = Var.Global name op
572 Env.register var
573 return $ Opt.Binop var leftExpr rightExpr