on 65e6ff1 help sign in
elk/src/Optimize.hs
514 lines 13 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize (optimize) where
5import qualified Control.Monad as M
6import qualified Data.Map as Map
7import qualified Data.Traversable as T
8import Data.Text (Text)
10import qualified AST.Expression.Canonical as Can
11import qualified AST.Expression.Optimized as Opt
12import qualified AST.Module.Name as ModuleName
13import qualified AST.Pattern as P
14import qualified AST.Variable as Var
15import qualified Optimize.Case as Case
16import qualified Optimize.DecisionTree as DT
17import qualified Optimize.Environment as Env
18import qualified Optimize.Inline as Inline
19import qualified Reporting.Annotation as A
20import qualified Reporting.Region as R
24-- OPTIMIZE
27optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [Opt.Def]
28optimize variantDict home defs =
29 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
33-- CONVERT DEFINITIONS
36optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
37optimizeDef isRoot (Can.Def _ pattern expression _) =
38 let
39 (args, canBody) =
40 Can.collectLambdas expression
42 maybeGetHome =
43 if isRoot then
44 Just <$> Env.getHome
46 else
47 return Nothing
48 in
49 do home <- maybeGetHome
50 optimizeDefHelp home pattern args canBody
53optimizeDefHelp
54 :: Maybe ModuleName.Canonical
55 -> P.Canonical
56 -> [P.Canonical]
57 -> Can.Expr
58 -> Env.Optimizer [Opt.Def]
59optimizeDefHelp home pattern@(A.A _ ptrn) rawArgs rawBody =
60 let
61 facts =
62 Opt.Facts home
63 in
64 case (ptrn, rawArgs) of
65 (P.Var name, _ : _) ->
67 do htc <- Env.getTailCall
68 Env.setTailCall False
70 (args, body) <- optimizeFunction (Just name) rawArgs rawBody
72 isTailCall <- Env.getTailCall
73 Env.setTailCall htc
75 return $
76 if isTailCall then
77 [ Opt.TailDef facts name args body ]
79 else
80 [ Opt.Def facts name (Opt.Function args body) ]
82 (P.Var name, []) ->
84 do hasTailCall <- Env.getTailCall
85 body <- optimizeExpr Nothing rawBody
86 Env.setTailCall hasTailCall
87 return [ Opt.Def facts name body ]
89 (_, []) ->
91 do name <- Env.freshName
92 optBody <- optimizeExpr Nothing rawBody
93 let optDef = Opt.Def facts name optBody
94 return (optDef : patternToDefs home name pattern)
96 _ ->
97 error "there should never be a function where the name is not a P.Var"
101-- OPTIMIZE FUNCTION
104optimizeFunction
105 :: Maybe Text
106 -> [P.Canonical]
107 -> Can.Expr
108 -> Env.Optimizer ([Text], Opt.Expr)
109optimizeFunction maybeName patterns givenBody =
110 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
111 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
112 return (argNames, optBody)
115makeContext :: [Text] -> Text -> Context
116makeContext argNames name =
117 Just (name, length argNames, argNames)
120depatternArgs :: ([Text], Can.Expr) -> P.Canonical -> Env.Optimizer ([Text], Can.Expr)
121depatternArgs (names, rawExpr) ptrn =
122 do (name, expr) <- depattern ptrn rawExpr
123 return (name : names, expr)
126patternToDefs :: Maybe ModuleName.Canonical -> Text -> P.Canonical -> [Opt.Def]
127patternToDefs home rootName pattern =
128 let
129 rootExpr =
130 case home of
131 Nothing ->
132 Opt.Var (Var.local rootName)
134 Just moduleName ->
135 Opt.Var (Var.topLevel moduleName rootName)
137 toDef (name, accessExpr) =
138 Opt.Def (Opt.Facts home) name accessExpr
139 in
140 map toDef (patternToSubstitutions rootExpr pattern)
144-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
147patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(Text, Opt.Expr)]
148patternToSubstitutions expr (A.A _ pattern) =
149 case pattern of
150 P.Var name ->
151 [ (name, expr) ]
153 P.Literal _ ->
154 []
156 P.Record fields ->
157 map (\name -> (name, Opt.Access expr name)) fields
159 P.Anything ->
160 []
162 P.Alias alias realPattern ->
163 (alias, expr) : patternToSubstitutions expr realPattern
165 P.Ctor _ patterns ->
166 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
170-- OPTIMIZE EXPRESSIONS
173type Context = Maybe (Text, Int, [Text])
176{-| This optimizer detects tail calls. The `Context` holds a function name,
177function arity, and list of argument names. If a tail call is possible, the
178`Context` will be filled in and we check any function we run into. If we go
179down a route that cannot possibly have tail calls, we remove the context.
180-}
181optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
182optimizeExpr context annExpr@(A.A region expression) =
183 let
184 keepLooking =
185 optimizeExpr context
186 -- this *might* find a tail call, use this when something is a potential
187 -- tail call
189 justConvert =
190 optimizeExpr Nothing
191 -- this will never find tail calls, use this when you are working with
192 -- expressions that cannot be turned into tail calls
193 in
194 case expression of
195 Can.Literal lit ->
196 pure (Opt.Literal lit)
198 Can.Var name ->
199 if name == Var.inCore "Debug" "crash" then
200 do home <- Env.getHome
201 pure (Opt.Crash home region Nothing)
203 else
204 pure (Opt.Var name)
206 Can.List elements ->
207 Opt.List
208 <$> T.traverse justConvert elements
210 Can.Binop op leftExpr rightExpr ->
211 optimizeBinop context region op leftExpr rightExpr
213 Can.Lambda _ _ ->
214 let
215 (patterns, body) =
216 Can.collectLambdas annExpr
217 in
218 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
220 Can.App _ _ ->
221 let
222 (func:args) =
223 Can.collectApps annExpr
224 in
225 case getTailCallInfo context func (length args) of
226 Just (name, argNames) ->
227 do optArgs <- T.traverse justConvert args
228 Env.setTailCall True
229 return (Opt.TailCall name argNames optArgs)
231 Nothing ->
232 do hasTailCall <- Env.getTailCall
233 optFunc <- justConvert func
234 optArgs <- T.traverse justConvert args
235 Env.setTailCall hasTailCall
236 return (Opt.Call optFunc optArgs)
238 Can.If branches finally ->
239 let
240 crawlBranch (cond,branch) =
241 (,) <$> justConvert cond <*> keepLooking branch
242 in
243 Opt.If
244 <$> T.traverse crawlBranch branches
245 <*> keepLooking finally
247 Can.Let defs body ->
248 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
249 Opt.Let optDefs <$> keepLooking body
251 Can.Case expr branches ->
252 do optExpr <- optimizeExpr Nothing expr
253 variantDict <- Env.getVariantDict
255 name <- Env.freshName
256 optBranches <- T.traverse (optimizeBranch context region name) branches
257 let optCase = Case.optimize variantDict name optBranches
258 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
260 Can.Ctor (Var.Canonical _ name) exprs ->
261 Opt.Ctor name <$> T.traverse justConvert exprs
263 Can.Access record field ->
264 Opt.Access
265 <$> justConvert record
266 <*> pure field
268 Can.Update record fields ->
269 Opt.Update
270 <$> justConvert record
271 <*> T.traverse (mapSnd justConvert) fields
273 Can.Record fields ->
274 Opt.Record
275 <$> T.traverse (mapSnd justConvert) fields
277 Can.Cmd moduleName ->
278 pure (Opt.Cmd moduleName)
280 Can.Sub moduleName ->
281 pure (Opt.Sub moduleName)
283 Can.OutgoingPort name tipe ->
284 pure (Opt.OutgoingPort name tipe)
286 Can.IncomingPort name tipe ->
287 pure (Opt.IncomingPort name tipe)
289 Can.Program kind expr ->
290 Opt.Program kind <$> justConvert expr
292 Can.SaveEnv _ _ ->
293 error "save_the_environment should never make it to optimization phase"
295 Can.GLShader uid src gltipe ->
296 pure (Opt.GLShader uid src gltipe)
299mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
300mapSnd func (x, a) =
301 (,) x <$> func a
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 "::"
490type Assoc =
491 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
494left :: Assoc
495left =
496 flip (,)
499right :: Assoc
500right =
501 (,)
504collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
505collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
506 case expr of
507 Can.Binop op leftExpr rightExpr | op == desiredOp ->
508 let
509 (func, arg) = assoc leftExpr rightExpr
510 in
511 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
513 _ ->
514 makeRoot annExpr