git.iko.soy/elk
on 25a0d74 help sign in
elk/src/Optimize.hs
534 lines 14 kB history raw
1module Optimize (optimize) where
3import qualified Control.Monad as M
4import qualified Data.Map as Map
5import qualified Data.Traversable as T
7import qualified AST.Expression.General as Expr
8import qualified AST.Expression.Canonical as Can
9import qualified AST.Expression.Optimized as Opt
10import qualified AST.Module as Module
11import qualified AST.Module.Name as ModuleName
12import qualified AST.Pattern as P
13import qualified AST.Variable as Var
14import qualified Optimize.Case as Case
15import qualified Optimize.DecisionTree as DT
16import qualified Optimize.Environment as Env
17import qualified Optimize.Inline as Inline
18import qualified Reporting.Annotation as A
19import qualified Reporting.Region as R
23-- OPTIMIZE
26optimize :: DT.VariantDict -> Module.CanonicalModule -> Module.Optimized
27optimize variantDict module_@(Module.Module home _ _ _ _ body) =
28 let
29 (defs, _) =
30 flattenLets [] (Module.program body)
32 optDefs =
33 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
34 in
35 module_ { Module.body = body { Module.program = optDefs } }
38flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
39flattenLets defs annExpr@(A.A _ expr) =
40 case expr of
41 Expr.Let ds body ->
42 flattenLets (defs ++ ds) body
44 _ ->
45 (defs, annExpr)
49-- CONVERT DEFINITIONS
52optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
53optimizeDef isRoot (Can.Definition (Can.Facts deps) pattern expression _) =
54 let
55 (args, canBody) =
56 Expr.collectLambdas expression
57 in
58 do home <-
59 if isRoot then
60 Just <$> Env.getHome
61 else
62 return Nothing
64 optimizeDefHelp home deps pattern args canBody
67optimizeDefHelp
68 :: Maybe ModuleName.Canonical
69 -> [Var.TopLevel]
70 -> P.CanonicalPattern
71 -> [P.CanonicalPattern]
72 -> Can.Expr
73 -> Env.Optimizer [Opt.Def]
74optimizeDefHelp home deps pattern@(A.A _ ptrn) args rawBody =
75 let
76 facts =
77 Opt.Facts home deps
78 in
79 case (ptrn, args) of
80 (P.Var name, _ : _) ->
82 do htc <- Env.getTailCall
83 Env.setTailCall False
85 (args, body) <- optimizeFunction (Just name) args rawBody
87 isTailCall <- Env.getTailCall
88 Env.setTailCall htc
90 return $
91 if isTailCall then
92 [ Opt.TailDef facts name args body ]
94 else
95 [ Opt.Def facts name (Opt.Function args body) ]
97 (P.Var name, []) ->
99 do hasTailCall <- Env.getTailCall
100 body <- optimizeExpr Nothing rawBody
101 Env.setTailCall hasTailCall
102 return [ Opt.Def facts name body ]
104 (_, []) ->
106 do name <- Env.freshName
107 optBody <- optimizeExpr Nothing rawBody
108 let optDef = Opt.Def facts name optBody
109 return (optDef : patternToDefs home name pattern)
111 _ ->
112 error "there should never be a function where the name is not a P.Var"
116-- OPTIMIZE FUNCTION
119optimizeFunction
120 :: Maybe String
121 -> [P.CanonicalPattern]
122 -> Can.Expr
123 -> Env.Optimizer ([String], Opt.Expr)
124optimizeFunction maybeName patterns givenBody =
125 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
126 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
127 return (argNames, optBody)
130makeContext :: [String] -> String -> Context
131makeContext argNames name =
132 Just (name, length argNames, argNames)
135depatternArgs
136 :: ([String], Can.Expr)
137 -> P.CanonicalPattern
138 -> Env.Optimizer ([String], Can.Expr)
139depatternArgs (names, rawExpr) ptrn =
140 do (name, expr) <- depattern ptrn rawExpr
141 return (name : names, expr)
144patternToDefs
145 :: Maybe ModuleName.Canonical
146 -> String
147 -> P.CanonicalPattern
148 -> [Opt.Def]
149patternToDefs home rootName pattern =
150 let
151 (deps, rootExpr) =
152 case home of
153 Nothing ->
154 ( []
155 , Opt.Var (Var.local rootName)
156 )
158 Just moduleName ->
159 ( [ Var.TopLevelVar moduleName rootName ]
160 , Opt.Var (Var.topLevel moduleName rootName)
161 )
163 toDef (name, accessExpr) =
164 Opt.Def (Opt.Facts home deps) name accessExpr
165 in
166 map toDef (patternToSubstitutions rootExpr pattern)
170-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
173patternToSubstitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
174patternToSubstitutions expr (A.A _ pattern) =
175 case pattern of
176 P.Var name ->
177 [ (name, expr) ]
179 P.Literal _ ->
180 []
182 P.Record fields ->
183 map (\name -> (name, Opt.Access expr name)) fields
185 P.Anything ->
186 []
188 P.Alias alias realPattern ->
189 (alias, expr) : patternToSubstitutions expr realPattern
191 P.Data _ patterns ->
192 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
196-- OPTIMIZE EXPRESSIONS
199type Context = Maybe (String, Int, [String])
202{-| This optimizer detects tail calls. The `Context` holds a function name,
203function arity, and list of argument names. If a tail call is possible, the
204`Context` will be filled in and we check any function we run into. If we go
205down a route that cannot possibly have tail calls, we remove the context.
206-}
207optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
208optimizeExpr context annExpr@(A.A region expression) =
209 let
210 keepLooking =
211 optimizeExpr context
212 -- this *might* find a tail call, use this when something is a potential
213 -- tail call
215 justConvert =
216 optimizeExpr Nothing
217 -- this will never find tail calls, use this when you are working with
218 -- expressions that cannot be turned into tail calls
219 in
220 case expression of
221 Expr.Literal lit ->
222 pure (Opt.Literal lit)
224 Expr.Var name ->
225 if name == Var.inCore ["Debug"] "crash" then
226 do home <- Env.getHome
227 pure (Opt.Crash home region Nothing)
229 else
230 pure (Opt.Var name)
232 Expr.Range lowExpr highExpr ->
233 Opt.Range
234 <$> justConvert lowExpr
235 <*> justConvert highExpr
237 Expr.ExplicitList elements ->
238 Opt.ExplicitList
239 <$> T.traverse justConvert elements
241 Expr.Binop op leftExpr rightExpr ->
242 optimizeBinop context region op leftExpr rightExpr
244 Expr.Lambda _ _ ->
245 let
246 (patterns, body) =
247 Expr.collectLambdas annExpr
248 in
249 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
251 Expr.App _ _ ->
252 let
253 (func:args) =
254 Expr.collectApps annExpr
255 in
256 case isTailCall context func (length args) of
257 Just (name, argNames) ->
258 do optArgs <- T.traverse justConvert args
259 Env.setTailCall True
260 return (Opt.TailCall name argNames optArgs)
262 Nothing ->
263 do hasTailCall <- Env.getTailCall
264 optFunc <- justConvert func
265 optArgs <- T.traverse justConvert args
266 Env.setTailCall hasTailCall
267 return (Opt.Call optFunc optArgs)
269 Expr.If branches finally ->
270 let
271 crawlBranch (cond,branch) =
272 (,) <$> justConvert cond <*> keepLooking branch
273 in
274 Opt.If
275 <$> T.traverse crawlBranch branches
276 <*> keepLooking finally
278 Expr.Let defs body ->
279 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
280 Opt.Let optDefs <$> keepLooking body
282 Expr.Case expr branches ->
283 do optExpr <- optimizeExpr Nothing expr
284 variantDict <- Env.getVariantDict
286 name <- Env.freshName
287 optBranches <- T.traverse (optimizeBranch context region name) branches
288 let optCase = Case.optimize variantDict name optBranches
289 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
291 Expr.Data name exprs ->
292 Opt.Data name <$> T.traverse justConvert exprs
294 Expr.Access record field ->
295 Opt.Access
296 <$> justConvert record
297 <*> pure field
299 Expr.Update record fields ->
300 Opt.Update
301 <$> justConvert record
302 <*> T.traverse (mapSnd justConvert) fields
304 Expr.Record fields ->
305 Opt.Record
306 <$> T.traverse (mapSnd justConvert) fields
308 Expr.GLShader uid src gltipe ->
309 pure (Opt.GLShader uid src gltipe)
311 Expr.Port impl ->
312 Opt.Port <$>
313 case impl of
314 Expr.In name tipe ->
315 pure (Expr.In name tipe)
317 Expr.Out name expr tipe ->
318 (\e -> Expr.Out name e tipe) <$> justConvert expr
320 Expr.Task name expr tipe ->
321 (\e -> Expr.Task name e tipe) <$> justConvert expr
324mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
325mapSnd func (x, a) =
326 (,) x <$> func a
330-- OPTIMIZE CASE BRANCHES
333optimizeBranch
334 :: Context
335 -> R.Region
336 -> String
337 -> (P.CanonicalPattern, Can.Expr)
338 -> Env.Optimizer (P.CanonicalPattern, Opt.Expr)
339optimizeBranch context region exprName (rawPattern, canBranch) =
340 let
341 root =
342 Opt.Var (Var.Canonical Var.Local exprName)
343 in
344 do optBranch <- optimizeExpr context canBranch
345 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
347 let substitutions =
348 Map.fromList (patternToSubstitutions root pattern)
350 (,) pattern <$> Inline.inline substitutions branch
353tagCrashBranch
354 :: R.Region
355 -> P.CanonicalPattern
356 -> Opt.Expr
357 -> Env.Optimizer (P.CanonicalPattern, Opt.Expr)
358tagCrashBranch region pattern@(A.A pr _) expr =
359 case expr of
360 Opt.Call (Opt.Crash home _ _) [arg] ->
361 do name <- Env.freshName
362 return
363 ( A.A pr (P.Alias name pattern)
364 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
365 )
367 _ ->
368 return (pattern, expr)
372-- DETECT TAIL CALL
375isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
376isTailCall context (A.A _ function) arity =
377 case (context, function) of
378 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
380 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
381 Just (name, argNames)
383 else
384 Nothing
386 _ ->
387 Nothing
391-- DEPATTERN
392-- given a pattern and an expression, push the actual pattern matching into the
393-- expression as a case-expression if necessary. The pattern compiler will work
394-- it out from there
397depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
398depattern canPattern@(A.A _ pattern) expr =
399 let
400 ann =
401 A.A (error "the annotation added in 'depattern' should not be observed!")
403 caseExpr e branches =
404 ann (Expr.Case e branches)
406 var name =
407 ann (Expr.Var (Var.local name))
409 caseify =
410 do name <- Env.freshName
411 return
412 ( name
413 , caseExpr (var name) [ (canPattern, expr) ]
414 )
415 in
416 case pattern of
417 P.Var name ->
418 return (name, expr)
420 P.Literal _ ->
421 caseify
423 P.Record _ ->
424 caseify
426 P.Anything ->
427 do name <- Env.freshName
428 return (name, expr)
430 P.Alias _ _ ->
431 caseify
433 P.Data _ _ ->
434 caseify
438-- OPTIMIZE BINOPS
441optimizeBinop
442 :: Context
443 -> R.Region
444 -> Var.Canonical
445 -> Can.Expr
446 -> Can.Expr
447 -> Env.Optimizer Opt.Expr
448optimizeBinop context region op leftExpr rightExpr =
449 let
450 ann = A.A region
451 binop = ann (Expr.Binop op leftExpr rightExpr)
452 in
453 if op == forwardApply then
455 optimizeExpr context (collect id left forwardApply binop)
457 else if op == backwardApply then
459 optimizeExpr context (collect id right backwardApply binop)
461 else if op == forwardCompose then
463 do var <- Env.freshName
464 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
465 let body = collect makeRoot left forwardCompose binop
466 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
468 else if op == backwardCompose then
470 do var <- Env.freshName
471 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
472 let body = collect makeRoot right backwardCompose binop
473 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
475 else if op == Var.Canonical Var.BuiltIn "::" then
477 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
479 else
481 Opt.Binop op
482 <$> optimizeExpr Nothing leftExpr
483 <*> optimizeExpr Nothing rightExpr
486-- left-associative ((x |> f) |> g)
487forwardApply :: Var.Canonical
488forwardApply =
489 Var.inCore ["Basics"] "|>"
492-- right-associative (g <| (f <| x))
493backwardApply :: Var.Canonical
494backwardApply =
495 Var.inCore ["Basics"] "<|"
498-- left-associative ((f >> g) >> h)
499forwardCompose :: Var.Canonical
500forwardCompose =
501 Var.inCore ["Basics"] ">>"
504-- right-associative (h << (g << f))
505backwardCompose :: Var.Canonical
506backwardCompose =
507 Var.inCore ["Basics"] "<<"
510type Assoc =
511 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
514left :: Assoc
515left =
516 flip (,)
519right :: Assoc
520right =
521 (,)
524collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
525collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
526 case expr of
527 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
528 let
529 (func, arg) = assoc leftExpr rightExpr
530 in
531 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
533 _ ->
534 makeRoot annExpr