git.iko.soy/elk
on 3734eff help sign in
elk/src/Optimize.hs
572 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 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 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
134 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
135 return (argNames, optBody)
138makeContext :: [Text] -> Text -> Context
139makeContext argNames name =
140 Just (name, length argNames, argNames)
143depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
144depatternArgs (names, rawExpr) ptrn =
145 do (name, expr) <- depattern ptrn rawExpr
146 return (name : names, expr)
150-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
153patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
154patternToSubstitutions expr (A.A _ pattern) =
155 case pattern of
156 P.Var name ->
157 [ (name, expr) ]
159 P.Literal _ ->
160 []
162 P.Record fields ->
163 map (\name -> (name, Opt.Access expr name)) fields
165 P.Anything ->
166 []
168 P.Alias alias realPattern ->
169 (alias, expr) : patternToSubstitutions expr realPattern
171 P.Ctor _ patterns ->
172 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
176-- OPTIMIZE EXPRESSIONS
179type Context = Maybe (Text, Int, [Text])
182{-| This optimizer detects tail calls. The `Context` holds a function name,
183function arity, and list of argument names. If a tail call is possible, the
184`Context` will be filled in and we check any function we run into. If we go
185down a route that cannot possibly have tail calls, we remove the context.
186-}
187optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
188optimizeExpr context annExpr@(A.A region expression) =
189 let
190 keepLooking =
191 optimizeExpr context
192 -- this *might* find a tail call, use this when something is a potential
193 -- tail call
195 justConvert =
196 optimizeExpr Nothing
197 -- this will never find tail calls, use this when you are working with
198 -- expressions that cannot be turned into tail calls
199 in
200 case expression of
201 Can.Literal lit ->
202 pure (Opt.Literal lit)
204 Can.Var var ->
205 optimizeVariable region var
207 Can.List elements ->
208 Opt.List
209 <$> traverse justConvert elements
211 Can.Binop op leftExpr rightExpr ->
212 optimizeBinop context region op leftExpr rightExpr
214 Can.Lambda _ _ ->
215 let
216 (patterns, body) =
217 Can.collectLambdas annExpr
218 in
219 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
221 Can.App _ _ ->
222 let
223 (func:args) =
224 Can.collectApps annExpr
225 in
226 case getTailCallInfo context func (length args) of
227 Just (name, argNames) ->
228 do optArgs <- traverse justConvert args
229 Env.setTailCall True
230 return (Opt.TailCall name argNames optArgs)
232 Nothing ->
233 do hasTailCall <- Env.getTailCall
234 optFunc <- justConvert func
235 optArgs <- traverse justConvert args
236 Env.setTailCall hasTailCall
237 return (Opt.Call optFunc optArgs)
239 Can.If branches finally ->
240 let
241 crawlBranch (cond,branch) =
242 (,) <$> justConvert cond <*> keepLooking branch
243 in
244 Opt.If
245 <$> traverse crawlBranch branches
246 <*> keepLooking finally
248 Can.Let defs body ->
249 do optDefs <- concat <$> traverse optimizeLetDef defs
250 Opt.Let optDefs <$> keepLooking body
252 Can.Case expr branches ->
253 do optExpr <- optimizeExpr Nothing expr
254 variantDict <- Env.getVariantDict
256 name <- Env.freshName
257 optBranches <- traverse (optimizeBranch context region name) branches
258 let optCase = Case.optimize variantDict name optBranches
259 return $ Opt.Let [ (name, Opt.Def optExpr) ] optCase
261 Can.Ctor (Var.Canonical _ name) exprs ->
262 Opt.Ctor name <$> traverse justConvert exprs
264 Can.Access record field ->
265 Opt.Access
266 <$> justConvert record
267 <*> pure field
269 Can.Update record fields ->
270 Opt.Update
271 <$> justConvert record
272 <*> traverse (traverse justConvert) fields
274 Can.Record fields ->
275 Opt.Record
276 <$> traverse (traverse justConvert) fields
278 Can.Cmd home manager ->
279 Opt.Cmd home <$> Env.registerEffects home manager
281 Can.Sub home manager ->
282 Opt.Sub home <$> Env.registerEffects home manager
284 Can.OutgoingPort name tipe ->
285 Opt.OutgoingPort name <$> Port.toEncoder tipe
287 Can.IncomingPort name tipe ->
288 Opt.IncomingPort name <$> Port.toDecoder tipe
290 Can.Program main expr ->
291 Opt.Program <$> optimizeMain main <*> justConvert expr
293 Can.SaveEnv _ _ ->
294 error "save_the_environment should never make it to optimization phase"
296 Can.GLShader _ src _ ->
297 pure (Opt.GLShader src)
301-- OPTIMIZE VARIABLE
304optimizeVariable :: R.Region -> Var.Canonical -> Env.Optimizer Opt.Expr
305optimizeVariable region (Var.Canonical home name) =
306 case home of
307 Var.BuiltIn ->
308 pure (Opt.VarLocal name)
310 Var.Local ->
311 pure (Opt.VarLocal name)
313 Var.TopLevel modul ->
314 do let var = Var.Global modul name
315 Env.register var
316 pure (Opt.VarGlobal var)
318 Var.Module modul ->
319 do let var = Var.Global modul name
320 Env.register var
321 if name == "crash" && modul == ModuleName.inCore "Debug"
322 then
323 do here <- Env.getHome
324 pure (Opt.Crash here region Nothing)
325 else
326 pure (Opt.VarGlobal var)
330-- OPTIMIZE MAIN
333optimizeMain :: Can.Main -> Env.Optimizer Opt.Main
334optimizeMain main =
335 case main of
336 Can.VDom ->
337 return Opt.VDom
339 Can.NoFlags ->
340 return Opt.NoFlags
342 Can.Flags tipe ->
343 Opt.Flags <$> Port.toDecoder tipe
347-- OPTIMIZE CASE BRANCHES
350optimizeBranch
351 :: Context
352 -> R.Region
353 -> Text
354 -> (P.Canonical, Can.Expr)
355 -> Env.Optimizer (P.Canonical, Opt.Expr)
356optimizeBranch context region exprName (rawPattern, canBranch) =
357 do optBranch <- optimizeExpr context canBranch
358 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
360 let substitutions = Map.fromList $
361 patternToSubstitutions (Opt.VarLocal exprName) pattern
363 (,) pattern <$> Inline.inline substitutions branch
366tagCrashBranch
367 :: R.Region
368 -> P.Canonical
369 -> Opt.Expr
370 -> Env.Optimizer (P.Canonical, Opt.Expr)
371tagCrashBranch region pattern@(A.A pr _) expr =
372 case expr of
373 Opt.Call (Opt.Crash home _ _) [arg] ->
374 do name <- Env.freshName
375 return
376 ( A.A pr (P.Alias name pattern)
377 , Opt.Call (Opt.Crash home region (Just (Opt.VarLocal name))) [arg]
378 )
380 _ ->
381 return (pattern, expr)
385-- DETECT TAIL CALL
388getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
389getTailCallInfo context (A.A _ function) arity =
390 case (context, function) of
391 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
393 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
394 Just (name, argNames)
396 else
397 Nothing
399 _ ->
400 Nothing
404-- DEPATTERN
405-- given a pattern and an expression, push the actual pattern matching into the
406-- expression as a case-expression if necessary. The pattern compiler will work
407-- it out from there
410depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
411depattern canPattern@(A.A _ pattern) expr =
412 let
413 ann =
414 A.A (error "the annotation added in 'depattern' should not be observed!")
416 caseExpr e branches =
417 ann (Can.Case e branches)
419 var name =
420 ann (Can.Var (Var.local name))
422 caseify =
423 do name <- Env.freshName
424 return
425 ( name
426 , caseExpr (var name) [ (canPattern, expr) ]
427 )
428 in
429 case pattern of
430 P.Var name ->
431 return (name, expr)
433 P.Literal _ ->
434 caseify
436 P.Record _ ->
437 caseify
439 P.Anything ->
440 do name <- Env.freshName
441 return (name, expr)
443 P.Alias _ _ ->
444 caseify
446 P.Ctor _ _ ->
447 caseify
451-- OPTIMIZE BINOPS
454optimizeBinop
455 :: Context
456 -> R.Region
457 -> Var.Canonical
458 -> Can.Expr
459 -> Can.Expr
460 -> Env.Optimizer Opt.Expr
461optimizeBinop context region op leftExpr rightExpr =
462 let
463 ann = A.A region
464 binop = ann (Can.Binop op leftExpr rightExpr)
465 in
466 if op == forwardApply then
468 optimizeExpr context (collect id left forwardApply binop)
470 else if op == backwardApply then
472 optimizeExpr context (collect id right backwardApply binop)
474 else if op == forwardCompose then
476 do var <- Env.freshName
477 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
478 let body = collect makeRoot left forwardCompose binop
479 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
481 else if op == backwardCompose then
483 do var <- Env.freshName
484 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
485 let body = collect makeRoot right backwardCompose binop
486 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
488 else if op == cons then
490 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
492 else
494 do optLeft <- optimizeExpr Nothing leftExpr
495 optRight <- optimizeExpr Nothing rightExpr
496 makeOptBinop op optLeft optRight
499-- left-associative ((x |> f) |> g)
500forwardApply :: Var.Canonical
501forwardApply =
502 Var.inCore "Basics" "|>"
505-- right-associative (g <| (f <| x))
506backwardApply :: Var.Canonical
507backwardApply =
508 Var.inCore "Basics" "<|"
511-- left-associative ((f >> g) >> h)
512forwardCompose :: Var.Canonical
513forwardCompose =
514 Var.inCore "Basics" ">>"
517-- right-associative (h << (g << f))
518backwardCompose :: Var.Canonical
519backwardCompose =
520 Var.inCore "Basics" "<<"
523cons :: Var.Canonical
524cons =
525 Var.Canonical Var.BuiltIn "::"
528type Assoc =
529 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
532left :: Assoc
533left =
534 flip (,)
537right :: Assoc
538right =
539 (,)
542collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
543collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
544 case expr of
545 Can.Binop op leftExpr rightExpr | op == desiredOp ->
546 let
547 (func, arg) = assoc leftExpr rightExpr
548 in
549 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
551 _ ->
552 makeRoot annExpr
555makeOptBinop :: Var.Canonical -> Opt.Expr -> Opt.Expr -> Env.Optimizer Opt.Expr
556makeOptBinop (Var.Canonical home op) leftExpr rightExpr =
557 case home of
558 Var.Local ->
559 error "bug manifesting in Optimize.makeOptBinop"
561 Var.BuiltIn ->
562 error "bug manifesting in Optimize.makeOptBinop"
564 Var.Module name ->
565 do let var = Var.Global name op
566 Env.register var
567 return $ Opt.Binop var leftExpr rightExpr
569 Var.TopLevel name ->
570 do let var = Var.Global name op
571 Env.register var
572 return $ Opt.Binop var leftExpr rightExpr