on ff0b05b help sign in
elk/src/Optimize.hs
506 lines 13 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.Canonical as Can
8import qualified AST.Expression.Optimized as Opt
9import qualified AST.Module.Name as ModuleName
10import qualified AST.Pattern as P
11import qualified AST.Variable as Var
12import qualified Optimize.Case as Case
13import qualified Optimize.DecisionTree as DT
14import qualified Optimize.Environment as Env
15import qualified Optimize.Inline as Inline
16import qualified Reporting.Annotation as A
17import qualified Reporting.Region as R
21-- OPTIMIZE
24optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [Opt.Def]
25optimize variantDict home defs =
26 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
30-- CONVERT DEFINITIONS
33optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
34optimizeDef isRoot (Can.Def _ pattern expression _) =
35 let
36 (args, canBody) =
37 Can.collectLambdas expression
39 maybeGetHome =
40 if isRoot then
41 Just <$> Env.getHome
43 else
44 return Nothing
45 in
46 do home <- maybeGetHome
47 optimizeDefHelp home pattern args canBody
50optimizeDefHelp
51 :: Maybe ModuleName.Canonical
52 -> P.Canonical
53 -> [P.Canonical]
54 -> Can.Expr
55 -> Env.Optimizer [Opt.Def]
56optimizeDefHelp home pattern@(A.A _ ptrn) args rawBody =
57 let
58 facts =
59 Opt.Facts home
60 in
61 case (ptrn, args) of
62 (P.Var name, _ : _) ->
64 do htc <- Env.getTailCall
65 Env.setTailCall False
67 (args, body) <- optimizeFunction (Just name) args rawBody
69 isTailCall <- Env.getTailCall
70 Env.setTailCall htc
72 return $
73 if isTailCall then
74 [ Opt.TailDef facts name args body ]
76 else
77 [ Opt.Def facts name (Opt.Function args body) ]
79 (P.Var name, []) ->
81 do hasTailCall <- Env.getTailCall
82 body <- optimizeExpr Nothing rawBody
83 Env.setTailCall hasTailCall
84 return [ Opt.Def facts name body ]
86 (_, []) ->
88 do name <- Env.freshName
89 optBody <- optimizeExpr Nothing rawBody
90 let optDef = Opt.Def facts name optBody
91 return (optDef : patternToDefs home name pattern)
93 _ ->
94 error "there should never be a function where the name is not a P.Var"
98-- OPTIMIZE FUNCTION
101optimizeFunction
102 :: Maybe String
103 -> [P.Canonical]
104 -> Can.Expr
105 -> Env.Optimizer ([String], Opt.Expr)
106optimizeFunction maybeName patterns givenBody =
107 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
108 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
109 return (argNames, optBody)
112makeContext :: [String] -> String -> Context
113makeContext argNames name =
114 Just (name, length argNames, argNames)
117depatternArgs :: ([String], Can.Expr) -> P.Canonical -> Env.Optimizer ([String], Can.Expr)
118depatternArgs (names, rawExpr) ptrn =
119 do (name, expr) <- depattern ptrn rawExpr
120 return (name : names, expr)
123patternToDefs :: Maybe ModuleName.Canonical -> String -> P.Canonical -> [Opt.Def]
124patternToDefs home rootName pattern =
125 let
126 rootExpr =
127 case home of
128 Nothing ->
129 Opt.Var (Var.local rootName)
131 Just moduleName ->
132 Opt.Var (Var.topLevel moduleName rootName)
134 toDef (name, accessExpr) =
135 Opt.Def (Opt.Facts home) name accessExpr
136 in
137 map toDef (patternToSubstitutions rootExpr pattern)
141-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
144patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(String, Opt.Expr)]
145patternToSubstitutions expr (A.A _ pattern) =
146 case pattern of
147 P.Var name ->
148 [ (name, expr) ]
150 P.Literal _ ->
151 []
153 P.Record fields ->
154 map (\name -> (name, Opt.Access expr name)) fields
156 P.Anything ->
157 []
159 P.Alias alias realPattern ->
160 (alias, expr) : patternToSubstitutions expr realPattern
162 P.Data _ patterns ->
163 concat (zipWith (patternToSubstitutions . Opt.CtorAccess expr) [0..] patterns)
167-- OPTIMIZE EXPRESSIONS
170type Context = Maybe (String, Int, [String])
173{-| This optimizer detects tail calls. The `Context` holds a function name,
174function arity, and list of argument names. If a tail call is possible, the
175`Context` will be filled in and we check any function we run into. If we go
176down a route that cannot possibly have tail calls, we remove the context.
177-}
178optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
179optimizeExpr context annExpr@(A.A region expression) =
180 let
181 keepLooking =
182 optimizeExpr context
183 -- this *might* find a tail call, use this when something is a potential
184 -- tail call
186 justConvert =
187 optimizeExpr Nothing
188 -- this will never find tail calls, use this when you are working with
189 -- expressions that cannot be turned into tail calls
190 in
191 case expression of
192 Can.Literal lit ->
193 pure (Opt.Literal lit)
195 Can.Var name ->
196 if name == Var.inCore ["Debug"] "crash" then
197 do home <- Env.getHome
198 pure (Opt.Crash home region Nothing)
200 else
201 pure (Opt.Var name)
203 Can.List elements ->
204 Opt.List
205 <$> T.traverse justConvert elements
207 Can.Binop op leftExpr rightExpr ->
208 optimizeBinop context region op leftExpr rightExpr
210 Can.Lambda _ _ ->
211 let
212 (patterns, body) =
213 Can.collectLambdas annExpr
214 in
215 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
217 Can.App _ _ ->
218 let
219 (func:args) =
220 Can.collectApps annExpr
221 in
222 case isTailCall context func (length args) of
223 Just (name, argNames) ->
224 do optArgs <- T.traverse justConvert args
225 Env.setTailCall True
226 return (Opt.TailCall name argNames optArgs)
228 Nothing ->
229 do hasTailCall <- Env.getTailCall
230 optFunc <- justConvert func
231 optArgs <- T.traverse justConvert args
232 Env.setTailCall hasTailCall
233 return (Opt.Call optFunc optArgs)
235 Can.If branches finally ->
236 let
237 crawlBranch (cond,branch) =
238 (,) <$> justConvert cond <*> keepLooking branch
239 in
240 Opt.If
241 <$> T.traverse crawlBranch branches
242 <*> keepLooking finally
244 Can.Let defs body ->
245 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
246 Opt.Let optDefs <$> keepLooking body
248 Can.Case expr branches ->
249 do optExpr <- optimizeExpr Nothing expr
250 variantDict <- Env.getVariantDict
252 name <- Env.freshName
253 optBranches <- T.traverse (optimizeBranch context region name) branches
254 let optCase = Case.optimize variantDict name optBranches
255 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
257 Can.Ctor name exprs ->
258 Opt.Ctor name <$> T.traverse justConvert exprs
260 Can.Access record field ->
261 Opt.Access
262 <$> justConvert record
263 <*> pure field
265 Can.Update record fields ->
266 Opt.Update
267 <$> justConvert record
268 <*> T.traverse (mapSnd justConvert) fields
270 Can.Record fields ->
271 Opt.Record
272 <$> T.traverse (mapSnd justConvert) fields
274 Can.Cmd moduleName ->
275 pure (Opt.Cmd moduleName)
277 Can.Sub moduleName ->
278 pure (Opt.Sub moduleName)
280 Can.OutgoingPort name tipe ->
281 pure (Opt.OutgoingPort name tipe)
283 Can.IncomingPort name tipe ->
284 pure (Opt.IncomingPort name tipe)
286 Can.Program kind expr ->
287 Opt.Program kind <$> justConvert expr
289 Can.SaveEnv _ _ ->
290 error "save_the_environment should never make it to optimization phase"
292 Can.GLShader uid src gltipe ->
293 pure (Opt.GLShader uid src gltipe)
296mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
297mapSnd func (x, a) =
298 (,) x <$> func a
302-- OPTIMIZE CASE BRANCHES
305optimizeBranch
306 :: Context
307 -> R.Region
308 -> String
309 -> (P.Canonical, Can.Expr)
310 -> Env.Optimizer (P.Canonical, Opt.Expr)
311optimizeBranch context region exprName (rawPattern, canBranch) =
312 let
313 root =
314 Opt.Var (Var.Canonical Var.Local exprName)
315 in
316 do optBranch <- optimizeExpr context canBranch
317 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
319 let substitutions =
320 Map.fromList (patternToSubstitutions root pattern)
322 (,) pattern <$> Inline.inline substitutions branch
325tagCrashBranch
326 :: R.Region
327 -> P.Canonical
328 -> Opt.Expr
329 -> Env.Optimizer (P.Canonical, Opt.Expr)
330tagCrashBranch region pattern@(A.A pr _) expr =
331 case expr of
332 Opt.Call (Opt.Crash home _ _) [arg] ->
333 do name <- Env.freshName
334 return
335 ( A.A pr (P.Alias name pattern)
336 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
337 )
339 _ ->
340 return (pattern, expr)
344-- DETECT TAIL CALL
347isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
348isTailCall context (A.A _ function) arity =
349 case (context, function) of
350 ( Just (ctxName, ctxArity, argNames), Can.Var (Var.Canonical home name) ) ->
352 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
353 Just (name, argNames)
355 else
356 Nothing
358 _ ->
359 Nothing
363-- DEPATTERN
364-- given a pattern and an expression, push the actual pattern matching into the
365-- expression as a case-expression if necessary. The pattern compiler will work
366-- it out from there
369depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (String, Can.Expr)
370depattern canPattern@(A.A _ pattern) expr =
371 let
372 ann =
373 A.A (error "the annotation added in 'depattern' should not be observed!")
375 caseExpr e branches =
376 ann (Can.Case e branches)
378 var name =
379 ann (Can.Var (Var.local name))
381 caseify =
382 do name <- Env.freshName
383 return
384 ( name
385 , caseExpr (var name) [ (canPattern, expr) ]
386 )
387 in
388 case pattern of
389 P.Var name ->
390 return (name, expr)
392 P.Literal _ ->
393 caseify
395 P.Record _ ->
396 caseify
398 P.Anything ->
399 do name <- Env.freshName
400 return (name, expr)
402 P.Alias _ _ ->
403 caseify
405 P.Data _ _ ->
406 caseify
410-- OPTIMIZE BINOPS
413optimizeBinop
414 :: Context
415 -> R.Region
416 -> Var.Canonical
417 -> Can.Expr
418 -> Can.Expr
419 -> Env.Optimizer Opt.Expr
420optimizeBinop context region op leftExpr rightExpr =
421 let
422 ann = A.A region
423 binop = ann (Can.Binop op leftExpr rightExpr)
424 in
425 if op == forwardApply then
427 optimizeExpr context (collect id left forwardApply binop)
429 else if op == backwardApply then
431 optimizeExpr context (collect id right backwardApply binop)
433 else if op == forwardCompose then
435 do var <- Env.freshName
436 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
437 let body = collect makeRoot left forwardCompose binop
438 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
440 else if op == backwardCompose then
442 do var <- Env.freshName
443 let makeRoot func = ann (Can.App func (ann (Can.Var (Var.local var))))
444 let body = collect makeRoot right backwardCompose binop
445 optimizeExpr context (ann (Can.Lambda (ann (P.Var var)) body))
447 else if op == Var.Canonical Var.BuiltIn "::" then
449 optimizeExpr context (ann (Can.Ctor "::" [ leftExpr, rightExpr ]))
451 else
453 Opt.Binop op
454 <$> optimizeExpr Nothing leftExpr
455 <*> optimizeExpr Nothing rightExpr
458-- left-associative ((x |> f) |> g)
459forwardApply :: Var.Canonical
460forwardApply =
461 Var.inCore ["Basics"] "|>"
464-- right-associative (g <| (f <| x))
465backwardApply :: Var.Canonical
466backwardApply =
467 Var.inCore ["Basics"] "<|"
470-- left-associative ((f >> g) >> h)
471forwardCompose :: Var.Canonical
472forwardCompose =
473 Var.inCore ["Basics"] ">>"
476-- right-associative (h << (g << f))
477backwardCompose :: Var.Canonical
478backwardCompose =
479 Var.inCore ["Basics"] "<<"
482type Assoc =
483 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
486left :: Assoc
487left =
488 flip (,)
491right :: Assoc
492right =
493 (,)
496collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
497collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
498 case expr of
499 Can.Binop op leftExpr rightExpr | op == desiredOp ->
500 let
501 (func, arg) = assoc leftExpr rightExpr
502 in
503 A.A ann (Can.App func (collect makeRoot assoc desiredOp arg))
505 _ ->
506 makeRoot annExpr