on eb97f2a help sign in
elk/src/Optimize.hs
507 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.General as Expr
8import qualified AST.Expression.Canonical as Can
9import qualified AST.Expression.Optimized as Opt
10import qualified AST.Module.Name as ModuleName
11import qualified AST.Pattern as P
12import qualified AST.Variable as Var
13import qualified Optimize.Case as Case
14import qualified Optimize.DecisionTree as DT
15import qualified Optimize.Environment as Env
16import qualified Optimize.Inline as Inline
17import qualified Reporting.Annotation as A
18import qualified Reporting.Region as R
22-- OPTIMIZE
25optimize :: DT.VariantDict -> ModuleName.Canonical -> [Can.Def] -> [Opt.Def]
26optimize variantDict home defs =
27 Env.run variantDict home (concat <$> mapM (optimizeDef True) defs)
31-- CONVERT DEFINITIONS
34optimizeDef :: Bool -> Can.Def -> Env.Optimizer [Opt.Def]
35optimizeDef isRoot (Can.Def _ pattern expression _) =
36 let
37 (args, canBody) =
38 Expr.collectLambdas expression
40 maybeGetHome =
41 if isRoot then
42 Just <$> Env.getHome
44 else
45 return Nothing
46 in
47 do home <- maybeGetHome
48 optimizeDefHelp home pattern args canBody
51optimizeDefHelp
52 :: Maybe ModuleName.Canonical
53 -> P.Canonical
54 -> [P.Canonical]
55 -> Can.Expr
56 -> Env.Optimizer [Opt.Def]
57optimizeDefHelp home pattern@(A.A _ ptrn) args rawBody =
58 let
59 facts =
60 Opt.Facts home
61 in
62 case (ptrn, args) of
63 (P.Var name, _ : _) ->
65 do htc <- Env.getTailCall
66 Env.setTailCall False
68 (args, body) <- optimizeFunction (Just name) args rawBody
70 isTailCall <- Env.getTailCall
71 Env.setTailCall htc
73 return $
74 if isTailCall then
75 [ Opt.TailDef facts name args body ]
77 else
78 [ Opt.Def facts name (Opt.Function args body) ]
80 (P.Var name, []) ->
82 do hasTailCall <- Env.getTailCall
83 body <- optimizeExpr Nothing rawBody
84 Env.setTailCall hasTailCall
85 return [ Opt.Def facts name body ]
87 (_, []) ->
89 do name <- Env.freshName
90 optBody <- optimizeExpr Nothing rawBody
91 let optDef = Opt.Def facts name optBody
92 return (optDef : patternToDefs home name pattern)
94 _ ->
95 error "there should never be a function where the name is not a P.Var"
99-- OPTIMIZE FUNCTION
102optimizeFunction
103 :: Maybe String
104 -> [P.Canonical]
105 -> Can.Expr
106 -> Env.Optimizer ([String], Opt.Expr)
107optimizeFunction maybeName patterns givenBody =
108 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
109 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
110 return (argNames, optBody)
113makeContext :: [String] -> String -> Context
114makeContext argNames name =
115 Just (name, length argNames, argNames)
118depatternArgs :: ([String], Can.Expr) -> P.Canonical -> Env.Optimizer ([String], Can.Expr)
119depatternArgs (names, rawExpr) ptrn =
120 do (name, expr) <- depattern ptrn rawExpr
121 return (name : names, expr)
124patternToDefs :: Maybe ModuleName.Canonical -> String -> P.Canonical -> [Opt.Def]
125patternToDefs home rootName pattern =
126 let
127 rootExpr =
128 case home of
129 Nothing ->
130 Opt.Var (Var.local rootName)
132 Just moduleName ->
133 Opt.Var (Var.topLevel moduleName rootName)
135 toDef (name, accessExpr) =
136 Opt.Def (Opt.Facts home) name accessExpr
137 in
138 map toDef (patternToSubstitutions rootExpr pattern)
142-- TURN A PATTERN INTO A BUNCH OF SUBSTITUTIONS
145patternToSubstitutions :: Opt.Expr -> P.Canonical -> [(String, Opt.Expr)]
146patternToSubstitutions expr (A.A _ pattern) =
147 case pattern of
148 P.Var name ->
149 [ (name, expr) ]
151 P.Literal _ ->
152 []
154 P.Record fields ->
155 map (\name -> (name, Opt.Access expr name)) fields
157 P.Anything ->
158 []
160 P.Alias alias realPattern ->
161 (alias, expr) : patternToSubstitutions expr realPattern
163 P.Data _ patterns ->
164 concat (zipWith (patternToSubstitutions . Opt.DataAccess expr) [0..] patterns)
168-- OPTIMIZE EXPRESSIONS
171type Context = Maybe (String, Int, [String])
174{-| This optimizer detects tail calls. The `Context` holds a function name,
175function arity, and list of argument names. If a tail call is possible, the
176`Context` will be filled in and we check any function we run into. If we go
177down a route that cannot possibly have tail calls, we remove the context.
178-}
179optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
180optimizeExpr context annExpr@(A.A region expression) =
181 let
182 keepLooking =
183 optimizeExpr context
184 -- this *might* find a tail call, use this when something is a potential
185 -- tail call
187 justConvert =
188 optimizeExpr Nothing
189 -- this will never find tail calls, use this when you are working with
190 -- expressions that cannot be turned into tail calls
191 in
192 case expression of
193 Expr.Literal lit ->
194 pure (Opt.Literal lit)
196 Expr.Var name ->
197 if name == Var.inCore ["Debug"] "crash" then
198 do home <- Env.getHome
199 pure (Opt.Crash home region Nothing)
201 else
202 pure (Opt.Var name)
204 Expr.ExplicitList elements ->
205 Opt.ExplicitList
206 <$> T.traverse justConvert elements
208 Expr.Binop op leftExpr rightExpr ->
209 optimizeBinop context region op leftExpr rightExpr
211 Expr.Lambda _ _ ->
212 let
213 (patterns, body) =
214 Expr.collectLambdas annExpr
215 in
216 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
218 Expr.App _ _ ->
219 let
220 (func:args) =
221 Expr.collectApps annExpr
222 in
223 case isTailCall context func (length args) of
224 Just (name, argNames) ->
225 do optArgs <- T.traverse justConvert args
226 Env.setTailCall True
227 return (Opt.TailCall name argNames optArgs)
229 Nothing ->
230 do hasTailCall <- Env.getTailCall
231 optFunc <- justConvert func
232 optArgs <- T.traverse justConvert args
233 Env.setTailCall hasTailCall
234 return (Opt.Call optFunc optArgs)
236 Expr.If branches finally ->
237 let
238 crawlBranch (cond,branch) =
239 (,) <$> justConvert cond <*> keepLooking branch
240 in
241 Opt.If
242 <$> T.traverse crawlBranch branches
243 <*> keepLooking finally
245 Expr.Let defs body ->
246 do optDefs <- concat <$> T.traverse (optimizeDef False) defs
247 Opt.Let optDefs <$> keepLooking body
249 Expr.Case expr branches ->
250 do optExpr <- optimizeExpr Nothing expr
251 variantDict <- Env.getVariantDict
253 name <- Env.freshName
254 optBranches <- T.traverse (optimizeBranch context region name) branches
255 let optCase = Case.optimize variantDict name optBranches
256 return $ Opt.Let [ Opt.Def Opt.dummyFacts name optExpr ] optCase
258 Expr.Data name exprs ->
259 Opt.Data name <$> T.traverse justConvert exprs
261 Expr.Access record field ->
262 Opt.Access
263 <$> justConvert record
264 <*> pure field
266 Expr.Update record fields ->
267 Opt.Update
268 <$> justConvert record
269 <*> T.traverse (mapSnd justConvert) fields
271 Expr.Record fields ->
272 Opt.Record
273 <$> T.traverse (mapSnd justConvert) fields
275 Expr.Cmd moduleName ->
276 pure (Opt.Cmd moduleName)
278 Expr.Sub moduleName ->
279 pure (Opt.Sub moduleName)
281 Expr.OutgoingPort name tipe ->
282 pure (Opt.OutgoingPort name tipe)
284 Expr.IncomingPort name tipe ->
285 pure (Opt.IncomingPort name tipe)
287 Expr.Program kind expr ->
288 Opt.Program kind <$> justConvert expr
290 Expr.SaveEnv _ _ ->
291 error "save_the_environment should never make it to optimization phase"
293 Expr.GLShader uid src gltipe ->
294 pure (Opt.GLShader uid src gltipe)
297mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
298mapSnd func (x, a) =
299 (,) x <$> func a
303-- OPTIMIZE CASE BRANCHES
306optimizeBranch
307 :: Context
308 -> R.Region
309 -> String
310 -> (P.Canonical, Can.Expr)
311 -> Env.Optimizer (P.Canonical, Opt.Expr)
312optimizeBranch context region exprName (rawPattern, canBranch) =
313 let
314 root =
315 Opt.Var (Var.Canonical Var.Local exprName)
316 in
317 do optBranch <- optimizeExpr context canBranch
318 (pattern, branch) <- tagCrashBranch region rawPattern optBranch
320 let substitutions =
321 Map.fromList (patternToSubstitutions root pattern)
323 (,) pattern <$> Inline.inline substitutions branch
326tagCrashBranch
327 :: R.Region
328 -> P.Canonical
329 -> Opt.Expr
330 -> Env.Optimizer (P.Canonical, Opt.Expr)
331tagCrashBranch region pattern@(A.A pr _) expr =
332 case expr of
333 Opt.Call (Opt.Crash home _ _) [arg] ->
334 do name <- Env.freshName
335 return
336 ( A.A pr (P.Alias name pattern)
337 , Opt.Call (Opt.Crash home region (Just (Opt.Var (Var.local name)))) [arg]
338 )
340 _ ->
341 return (pattern, expr)
345-- DETECT TAIL CALL
348isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
349isTailCall context (A.A _ function) arity =
350 case (context, function) of
351 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
353 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
354 Just (name, argNames)
356 else
357 Nothing
359 _ ->
360 Nothing
364-- DEPATTERN
365-- given a pattern and an expression, push the actual pattern matching into the
366-- expression as a case-expression if necessary. The pattern compiler will work
367-- it out from there
370depattern :: P.Canonical -> Can.Expr -> Env.Optimizer (String, Can.Expr)
371depattern canPattern@(A.A _ pattern) expr =
372 let
373 ann =
374 A.A (error "the annotation added in 'depattern' should not be observed!")
376 caseExpr e branches =
377 ann (Expr.Case e branches)
379 var name =
380 ann (Expr.Var (Var.local name))
382 caseify =
383 do name <- Env.freshName
384 return
385 ( name
386 , caseExpr (var name) [ (canPattern, expr) ]
387 )
388 in
389 case pattern of
390 P.Var name ->
391 return (name, expr)
393 P.Literal _ ->
394 caseify
396 P.Record _ ->
397 caseify
399 P.Anything ->
400 do name <- Env.freshName
401 return (name, expr)
403 P.Alias _ _ ->
404 caseify
406 P.Data _ _ ->
407 caseify
411-- OPTIMIZE BINOPS
414optimizeBinop
415 :: Context
416 -> R.Region
417 -> Var.Canonical
418 -> Can.Expr
419 -> Can.Expr
420 -> Env.Optimizer Opt.Expr
421optimizeBinop context region op leftExpr rightExpr =
422 let
423 ann = A.A region
424 binop = ann (Expr.Binop op leftExpr rightExpr)
425 in
426 if op == forwardApply then
428 optimizeExpr context (collect id left forwardApply binop)
430 else if op == backwardApply then
432 optimizeExpr context (collect id right backwardApply binop)
434 else if op == forwardCompose then
436 do var <- Env.freshName
437 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
438 let body = collect makeRoot left forwardCompose binop
439 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
441 else if op == backwardCompose then
443 do var <- Env.freshName
444 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
445 let body = collect makeRoot right backwardCompose binop
446 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
448 else if op == Var.Canonical Var.BuiltIn "::" then
450 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
452 else
454 Opt.Binop op
455 <$> optimizeExpr Nothing leftExpr
456 <*> optimizeExpr Nothing rightExpr
459-- left-associative ((x |> f) |> g)
460forwardApply :: Var.Canonical
461forwardApply =
462 Var.inCore ["Basics"] "|>"
465-- right-associative (g <| (f <| x))
466backwardApply :: Var.Canonical
467backwardApply =
468 Var.inCore ["Basics"] "<|"
471-- left-associative ((f >> g) >> h)
472forwardCompose :: Var.Canonical
473forwardCompose =
474 Var.inCore ["Basics"] ">>"
477-- right-associative (h << (g << f))
478backwardCompose :: Var.Canonical
479backwardCompose =
480 Var.inCore ["Basics"] "<<"
483type Assoc =
484 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
487left :: Assoc
488left =
489 flip (,)
492right :: Assoc
493right =
494 (,)
497collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
498collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
499 case expr of
500 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
501 let
502 (func, arg) = assoc leftExpr rightExpr
503 in
504 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
506 _ ->
507 makeRoot annExpr