git.iko.soy/elk
on f232da4 help sign in
elk/src/Optimize.hs
519 lines 14 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 Reporting.Annotation as A
23import qualified Reporting.Region as R
27-- OPTIMIZE
30optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [(Text, Opt.Decl)]
31optimize variantDict home defs =
32 concat $ zipWith (optimizeDecl variantDict home) [1 .. length defs] defs
36-- OPTIMIZE DECLS
39optimizeDecl :: DT.VariantDict -> ModuleName.Canonical -> Int -> Can.Def -> [(Text, Opt.Decl)]
40optimizeDecl variantDict home index (Can.Def _ pattern@(A.A _ ptrn) expression _) =
41 let
42 (rawArgs, rawBody) =
43 Can.collectLambdas expression
44 in
45 case (ptrn, rawArgs) of
46 (P.Var name, _) ->
47 (:[]) $ Env.run variantDict home $
48 optimizeNamedDef name rawArgs rawBody
50 (_, []) ->
51 let
52 name = "_" <> Text.pack (show index)
53 deps = Set.singleton (Var.Global home name)
54 decl =
55 Env.run variantDict home $
56 ((,) name . Opt.Def) <$> optimizeExpr Nothing rawBody
57 destructors =
58 destruct (Var.topLevel home name) pattern
59 in
60 decl : map (second (Opt.Decl deps Nothing)) destructors
62 _ ->
63 error "bug manifesting in Optimize.optimizeDecl, please report"
67-- OPTIMIZE DEFS
70optimizeLetDef :: Can.Def -> Env.Optimizer [(Text, Opt.Def)]
71optimizeLetDef (Can.Def _ pattern@(A.A _ ptrn) expression _) =
72 let
73 (rawArgs, rawBody) =
74 Can.collectLambdas expression
75 in
76 case (ptrn, rawArgs) of
77 (P.Var name, _) ->
78 do def <- optimizeNamedDef name rawArgs rawBody
79 return [def]
81 (_, []) ->
82 do name <- Env.freshName
83 body <- optimizeExpr Nothing rawBody
84 return $ (name, Opt.Def body) : destruct (Var.local name) pattern
86 _ ->
87 error "bug manifesting in Optimize.optimizeLetDef, please report"
90optimizeNamedDef :: Text -> [P.Canonical] -> Can.Expr -> Env.Optimizer (Text, Opt.Def)
91optimizeNamedDef name rawArgs rawBody =
92 case rawArgs of
93 [] ->
94 do hasTailCall <- Env.getTailCall
95 body <- optimizeExpr Nothing rawBody
96 Env.setTailCall hasTailCall
97 return (name, Opt.Def body)
99 _ : _ ->
100 do htc <- Env.getTailCall
101 Env.setTailCall False
103 (args, body) <- optimizeFunction (Just name) rawArgs rawBody
105 isTailCall <- Env.getTailCall
106 Env.setTailCall htc
108 return $ (,) name $
109 if isTailCall then
110 Opt.TailDef args body
111 else
112 Opt.Def (Opt.Function args body)
115destruct :: Var.Canonical -> P.Canonical -> [(Text, Opt.Def)]
116destruct var pattern =
117 map (second Opt.Def) $
118 patternToSubstitutions (Opt.Var var) pattern
122-- OPTIMIZE FUNCTION
125optimizeFunction
126 :: Maybe Text
127 -> [P.Canonical]
128 -> Can.Expr
129 -> Env.Optimizer ([Text], Opt.Expr)
130optimizeFunction maybeName patterns givenBody =
131 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
132 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
133 return (argNames, optBody)
136makeContext :: [Text] -> Text -> Context
137makeContext argNames name =
138 Just (name, length argNames, argNames)
141depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
142depatternArgs (names, rawExpr) ptrn =
143 do (name, expr) <- depattern ptrn rawExpr
144 return (name : names, expr)
148-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
151patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
152patternToSubstitutions expr (A.A _ pattern) =
153 case pattern of
154 P.Var name ->
155 [ (name, expr) ]
157 P.Literal _ ->
158 []
160 P.Record fields ->
161 map (\name -> (name, Opt.Access expr name)) fields
163 P.Anything ->
164 []
166 P.Alias alias realPattern ->
167 (alias, expr) : patternToSubstitutions expr realPattern
169 P.Ctor _ patterns ->
170 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
174-- OPTIMIZE EXPRESSIONS
177type Context = Maybe (Text, Int, [Text])
180{-| This optimizer detects tail calls. The `Context` holds a function name,
181function arity, and list of argument names. If a tail call is possible, the
182`Context` will be filled in and we check any function we run into. If we go
183down a route that cannot possibly have tail calls, we remove the context.
184-}
185optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
186optimizeExpr context annExpr@(A.A region expression) =
187 let
188 keepLooking =
189 optimizeExpr context
190 -- this *might* find a tail call, use this when something is a potential
191 -- tail call
193 justConvert =
194 optimizeExpr Nothing
195 -- this will never find tail calls, use this when you are working with
196 -- expressions that cannot be turned into tail calls
197 in
198 case expression of
199 Can.Literal lit ->
200 pure (Opt.Literal lit)
202 Can.Var name ->
203 do Env.register name
204 if name == debugCrash
205 then
206 do home <- Env.getHome
207 pure (Opt.Crash home region Nothing)
208 else
209 pure (Opt.Var name)
211 Can.List elements ->
212 Opt.List
213 <$> traverse justConvert elements
215 Can.Binop op leftExpr rightExpr ->
216 optimizeBinop context region op leftExpr rightExpr
218 Can.Lambda _ _ ->
219 let
220 (patterns, body) =
221 Can.collectLambdas annExpr
222 in
223 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
225 Can.App _ _ ->
226 let
227 (func:args) =
228 Can.collectApps annExpr
229 in
230 case getTailCallInfo context func (length args) of
231 Just (name, argNames) ->
232 do optArgs <- traverse justConvert args
233 Env.setTailCall True
234 return (Opt.TailCall name argNames optArgs)
236 Nothing ->
237 do hasTailCall <- Env.getTailCall
238 optFunc <- justConvert func
239 optArgs <- traverse justConvert args
240 Env.setTailCall hasTailCall
241 return (Opt.Call optFunc optArgs)
243 Can.If branches finally ->
244 let
245 crawlBranch (cond,branch) =
246 (,) <$> justConvert cond <*> keepLooking branch
247 in
248 Opt.If
249 <$> traverse crawlBranch branches
250 <*> keepLooking finally
252 Can.Let defs body ->
253 do optDefs <- concat <$> traverse optimizeLetDef defs
254 Opt.Let optDefs <$> keepLooking body
256 Can.Case expr branches ->
257 do optExpr <- optimizeExpr Nothing expr
258 variantDict <- Env.getVariantDict
260 name <- Env.freshName
261 optBranches <- traverse (optimizeBranch context region name) branches
262 let optCase = Case.optimize variantDict name optBranches
263 return $ Opt.Let [ (name, Opt.Def optExpr) ] optCase
265 Can.Ctor (Var.Canonical _ name) exprs ->
266 Opt.Ctor name <$> traverse justConvert exprs
268 Can.Access record field ->
269 Opt.Access
270 <$> justConvert record
271 <*> pure field
273 Can.Update record fields ->
274 Opt.Update
275 <$> justConvert record
276 <*> traverse (traverse justConvert) fields
278 Can.Record fields ->
279 Opt.Record
280 <$> traverse (traverse justConvert) fields
282 Can.Cmd moduleName manager ->
283 Opt.Cmd moduleName <$> Env.registerEffects manager
285 Can.Sub moduleName manager ->
286 Opt.Sub moduleName <$> Env.registerEffects manager
288 Can.OutgoingPort name tipe ->
289 pure (Opt.OutgoingPort name tipe)
291 Can.IncomingPort name tipe ->
292 pure (Opt.IncomingPort name tipe)
294 Can.Program kind expr ->
295 Opt.Program kind <$> justConvert expr
297 Can.SaveEnv _ _ ->
298 error "save_the_environment should never make it to optimization phase"
300 Can.GLShader uid src gltipe ->
301 pure (Opt.GLShader uid src gltipe)
305-- OPTIMIZE CASE BRANCHES
308optimizeBranch
309 :: Context
310 -> R.Region
311 -> Text
312 -> (P.Canonical, Can.Expr)
313 -> Env.Optimizer (P.Canonical, Opt.Expr)
314optimizeBranch context region exprName (rawPattern, canBranch) =
315 let
316 root =
317 Opt.Var (Var.Canonical Var.Local exprName)
318 in
319 do optBranch <- optimizeExpr context canBranch
320 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
322 let substitutions =
323 Map.fromList (patternToSubstitutions root pattern)
325 (,) pattern <$> Inline.inline substitutions branch
328tagCrashBranch
329 :: R.Region
330 -> P.Canonical
331 -> Opt.Expr
332 -> Env.Optimizer (P.Canonical, Opt.Expr)
333tagCrashBranch region pattern@(A.A pr _) expr =
334 case expr of
335 Opt.Call (Opt.Crash home _ _) [arg] ->
336 do name <- Env.freshName
337 return
338 ( A.A pr (P.Alias name pattern)
339 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
340 )
342 _ ->
343 return (pattern, expr)
347-- DETECT TAIL CALL
350getTailCallInfo :: Context -> Can.Expr -> Int -> Maybe (Text, [Text])
351getTailCallInfo context (A.A _ function) arity =
352 case (context, function) of
353 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
355 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
356 Just (name, argNames)
358 else
359 Nothing
361 _ ->
362 Nothing
366-- DEPATTERN
367-- given a pattern and an expression, push the actual pattern matching into the
368-- expression as a case-expression if necessary. The pattern compiler will work
369-- it out from there
372depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (Text, Can.Expr)
373depattern canPattern@(A.A _ pattern) expr =
374 let
375 ann =
376 A.A (error "the annotation added in 'depattern' should not be observed!")
378 caseExpr e branches =
379 ann (Can.Case e branches)
381 var name =
382 ann (Can.Var (Var.local name))
384 caseify =
385 do name <- Env.freshName
386 return
387 ( name
388 , caseExpr (var name) [ (canPattern, expr) ]
389 )
390 in
391 case pattern of
392 P.Var name ->
393 return (name, expr)
395 P.Literal _ ->
396 caseify
398 P.Record _ ->
399 caseify
401 P.Anything ->
402 do name <- Env.freshName
403 return (name, expr)
405 P.Alias _ _ ->
406 caseify
408 P.Ctor _ _ ->
409 caseify
413-- OPTIMIZE BINOPS
416optimizeBinop
417 :: Context
418 -> R.Region
419 -> Var.Canonical
420 -> Can.Expr
421 -> Can.Expr
422 -> Env.Optimizer Opt.Expr
423optimizeBinop context region op leftExpr rightExpr =
424 let
425 ann = A.A region
426 binop = ann (Can.Binop op leftExpr rightExpr)
427 in
428 if op == forwardApply then
430 optimizeExpr context (collect id left forwardApply binop)
432 else if op == backwardApply then
434 optimizeExpr context (collect id right backwardApply binop)
436 else if op == forwardCompose then
438 do var <- Env.freshName
439 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
440 let body = collect makeRoot left forwardCompose binop
441 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
443 else if op == backwardCompose then
445 do var <- Env.freshName
446 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
447 let body = collect makeRoot right backwardCompose binop
448 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
450 else if op == cons then
452 optimizeExpr context (ann (Can.Ctor cons [ leftExpr, rightExpr ]))
454 else
456 Opt.Binop op
457 <$> optimizeExpr Nothing leftExpr
458 <*> optimizeExpr Nothing rightExpr
461-- left-associative ((x |> f) |> g)
462forwardApply :: Var.Canonical
463forwardApply =
464 Var.inCore "Basics" "|>"
467-- right-associative (g <| (f <| x))
468backwardApply :: Var.Canonical
469backwardApply =
470 Var.inCore "Basics" "<|"
473-- left-associative ((f >> g) >> h)
474forwardCompose :: Var.Canonical
475forwardCompose =
476 Var.inCore "Basics" ">>"
479-- right-associative (h << (g << f))
480backwardCompose :: Var.Canonical
481backwardCompose =
482 Var.inCore "Basics" "<<"
485cons :: Var.Canonical
486cons =
487 Var.Canonical Var.BuiltIn "::"
490debugCrash :: Var.Canonical
491debugCrash =
492 Var.inCore "Debug" "crash"
495type Assoc =
496 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
499left :: Assoc
500left =
501 flip (,)
504right :: Assoc
505right =
506 (,)
509collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
510collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
511 case expr of
512 Can.Binop op leftExpr rightExpr | op == desiredOp ->
513 let
514 (func, arg) = assoc leftExpr rightExpr
515 in
516 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
518 _ ->
519 makeRoot annExpr