git.iko.soy/elk
on 97552fd help sign in
elk/src/Optimize.hs
451 lines 12 kB history raw
1module Optimize (optimize) where
3import qualified Control.Monad as M
4import qualified Data.Traversable as T
6import qualified AST.Expression.General as Expr
7import qualified AST.Expression.Canonical as Can
8import qualified AST.Expression.Optimized as Opt
9import qualified AST.Module as Module
10import qualified AST.Module.Name as ModuleName
11import qualified AST.Pattern as P
12import qualified AST.Variable as Var
13import qualified Optimize.Environment as Env
14import qualified Optimize.Patterns as Patterns
15import qualified Optimize.Patterns.DecisionTree as DT
16import qualified Reporting.Annotation as A
17import qualified Reporting.Region as R
20-- OPTIMIZE
22optimize :: DT.VariantDict -> Module.CanonicalModule -> Module.Optimized
23optimize variantDict module_@(Module.Module home _ _ _ _ body) =
24 let
25 (defs, _) =
26 flattenLets [] (Module.program body)
28 optDefs =
29 Env.run variantDict (concat <$> mapM (optimizeDef (Just home)) defs)
30 in
31 module_ { Module.body = body { Module.program = optDefs } }
34flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
35flattenLets defs annExpr@(A.A _ expr) =
36 case expr of
37 Expr.Let ds body ->
38 flattenLets (defs ++ ds) body
40 _ ->
41 (defs, annExpr)
44-- CONVERT DEFINITIONS
46optimizeDef :: Maybe ModuleName.Canonical -> Can.Def -> Env.Optimizer [Opt.Def]
47optimizeDef home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
48 let
49 (args, rawBody) =
50 Expr.collectLambdas expression
51 in
52 case (ptrn, args) of
53 (P.Var name, _ : _) ->
55 do htc <- Env.getTailCall
56 Env.setTailCall False
58 (args, body) <- optimizeFunction (Just name) args rawBody
60 isTailCall <- Env.getTailCall
61 Env.setTailCall htc
63 let facts = Opt.Facts home deps
65 return $
66 if isTailCall then
68 [ Opt.TailDef facts name args body ]
70 else
72 [ Opt.Def facts name (Opt.Function args body) ]
74 (P.Var name, []) ->
76 do hasTailCall <- Env.getTailCall
77 body <- optimizeExpr Nothing rawBody
78 Env.setTailCall hasTailCall
79 return [ Opt.Def (Opt.Facts home deps) name body ]
81 (_, []) ->
83 do name <- Env.freshName
84 optBody <- optimizeExpr Nothing rawBody
85 let optDef = Opt.Def (Opt.Facts home deps) name optBody
86 return (optDef : patternToDefs home name pattern)
88 _ ->
89 error "there should never be a function where the name is not a P.Var"
92-- OPTIMIZE FUNCTION
94optimizeFunction
95 :: Maybe String
96 -> [P.CanonicalPattern]
97 -> Can.Expr
98 -> Env.Optimizer ([String], Opt.Expr)
99optimizeFunction maybeName patterns givenBody =
100 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
101 optBody <- optimizeExpr (makeContext argNames =<< maybeName) body
102 return (argNames, optBody)
105makeContext :: [String] -> String -> Context
106makeContext argNames name =
107 Just (name, length argNames, argNames)
110depatternArgs
111 :: ([String], Can.Expr)
112 -> P.CanonicalPattern
113 -> Env.Optimizer ([String], Can.Expr)
114depatternArgs (names, rawExpr) ptrn =
115 do (name, expr) <- depattern ptrn rawExpr
116 return (name : names, expr)
119patternToDefs
120 :: Maybe ModuleName.Canonical
121 -> String
122 -> P.CanonicalPattern
123 -> [Opt.Def]
124patternToDefs home rootName pattern =
125 let
126 (deps, rootExpr) =
127 case home of
128 Nothing ->
129 ( []
130 , Opt.Var (Var.local rootName)
131 )
133 Just moduleName ->
134 ( [ Var.TopLevelVar moduleName rootName ]
135 , Opt.Var (Var.topLevel moduleName rootName)
136 )
138 toDef (name, accessExpr) =
139 Opt.Def (Opt.Facts home deps) name accessExpr
140 in
141 map toDef (substitutions rootExpr pattern)
144substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
145substitutions 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) : substitutions expr realPattern
162 P.Data _ patterns ->
163 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
166-- OPTIMIZE EXPRESSIONS
168type Context = Maybe (String, Int, [String])
171{-| This optimizer detects tail calls. The `Context` holds a function name,
172function arity, and list of argument names. If a tail call is possible, the
173`Context` will be filled in and we check any function we run into. If we go
174down a route that cannot possibly have tail calls, we remove the context.
175-}
176optimizeExpr :: Context -> Can.Expr -> Env.Optimizer Opt.Expr
177optimizeExpr context annExpr@(A.A region expression) =
178 let
179 keepLooking =
180 optimizeExpr context
181 -- this *might* find a tail call, use this when something is a potential
182 -- tail call
184 justConvert =
185 optimizeExpr Nothing
186 -- this will never find tail calls, use this when you are working with
187 -- expressions that cannot be turned into tail calls
188 in
189 case expression of
190 Expr.Literal lit ->
191 pure (Opt.Literal lit)
193 Expr.Var name ->
194 if name == Var.inCore ["Debug"] "crash" then
195 pure (Opt.Crash region Nothing)
197 else
198 pure (Opt.Var name)
200 Expr.Range lowExpr highExpr ->
201 Opt.Range
202 <$> justConvert lowExpr
203 <*> justConvert highExpr
205 Expr.ExplicitList elements ->
206 Opt.ExplicitList
207 <$> T.traverse justConvert elements
209 Expr.Binop op leftExpr rightExpr ->
210 optimizeBinop context region op leftExpr rightExpr
212 Expr.Lambda _ _ ->
213 let
214 (patterns, body) =
215 Expr.collectLambdas annExpr
216 in
217 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
219 Expr.App _ _ ->
220 let
221 (func:args) =
222 Expr.collectApps annExpr
223 in
224 case isTailCall context func (length args) of
225 Just (name, argNames) ->
226 do optArgs <- T.traverse justConvert args
227 Env.setTailCall True
228 return (Opt.TailCall name argNames optArgs)
230 Nothing ->
231 do hasTailCall <- Env.getTailCall
232 optFunc <- justConvert func
233 optArgs <- T.traverse justConvert args
234 Env.setTailCall hasTailCall
235 return (Opt.Call optFunc optArgs)
237 Expr.If branches finally ->
238 let
239 crawlBranch (cond,branch) =
240 (,) <$> justConvert cond <*> keepLooking branch
241 in
242 Opt.If
243 <$> T.traverse crawlBranch branches
244 <*> keepLooking finally
246 Expr.Let defs body ->
247 do optDefs <- concat <$> T.traverse (optimizeDef Nothing) defs
248 Opt.Let optDefs <$> keepLooking body
250 Expr.Case expr branches ->
251 do optExpr <- optimizeExpr Nothing expr
252 optBranches <- T.traverse (mapSnd (optimizeExpr context)) branches
253 variantDict <- Env.getVariantDict
254 Patterns.optimize variantDict region optExpr optBranches
256 Expr.Data name exprs ->
257 Opt.Data name <$> T.traverse justConvert exprs
259 Expr.Access record field ->
260 Opt.Access
261 <$> justConvert record
262 <*> pure field
264 Expr.Update record fields ->
265 Opt.Update
266 <$> justConvert record
267 <*> T.traverse (mapSnd justConvert) fields
269 Expr.Record fields ->
270 Opt.Record
271 <$> T.traverse (mapSnd justConvert) fields
273 Expr.GLShader uid src gltipe ->
274 pure (Opt.GLShader uid src gltipe)
276 Expr.Port impl ->
277 Opt.Port <$>
278 case impl of
279 Expr.In name tipe ->
280 pure (Expr.In name tipe)
282 Expr.Out name expr tipe ->
283 (\e -> Expr.Out name e tipe) <$> justConvert expr
285 Expr.Task name expr tipe ->
286 (\e -> Expr.Task name e tipe) <$> justConvert expr
289mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
290mapSnd func (x, a) =
291 (,) x <$> func a
294-- DETECT TAIL CALL
296isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
297isTailCall context (A.A _ function) arity =
298 case (context, function) of
299 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
301 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
302 Just (name, argNames)
304 else
305 Nothing
307 _ ->
308 Nothing
311-- DEPATTERN
312-- given a pattern and an expression, push the actual pattern matching into the
313-- expression as a case-expression if necessary. The pattern compiler will work
314-- it out from there
316depattern :: P.CanonicalPattern -> Can.Expr -> Env.Optimizer (String, Can.Expr)
317depattern canPattern@(A.A _ pattern) expr =
318 let
319 ann =
320 A.A (error "the annotation added in 'depattern' should not be observed!")
322 caseExpr e branches =
323 ann (Expr.Case e branches)
325 var name =
326 ann (Expr.Var (Var.local name))
328 caseify =
329 do name <- Env.freshName
330 return
331 ( name
332 , caseExpr (var name) [ (canPattern, expr) ]
333 )
334 in
335 case pattern of
336 P.Var name ->
337 return (name, expr)
339 P.Literal _ ->
340 caseify
342 P.Record _ ->
343 caseify
345 P.Anything ->
346 do name <- Env.freshName
347 return (name, expr)
349 P.Alias _ _ ->
350 caseify
352 P.Data _ _ ->
353 caseify
356-- OPTIMIZE BINOPS
358optimizeBinop
359 :: Context
360 -> R.Region
361 -> Var.Canonical
362 -> Can.Expr
363 -> Can.Expr
364 -> Env.Optimizer Opt.Expr
365optimizeBinop context region op leftExpr rightExpr =
366 let
367 ann = A.A region
368 binop = ann (Expr.Binop op leftExpr rightExpr)
369 in
370 if op == forwardApply then
372 optimizeExpr context (collect id left forwardApply binop)
374 else if op == backwardApply then
376 optimizeExpr context (collect id right backwardApply binop)
378 else if op == forwardCompose then
380 do var <- Env.freshName
381 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
382 let body = collect makeRoot left forwardCompose binop
383 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
385 else if op == backwardCompose then
387 do var <- Env.freshName
388 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
389 let body = collect makeRoot right backwardCompose binop
390 optimizeExpr context (ann (Expr.Lambda (ann (P.Var var)) body))
392 else if op == Var.Canonical Var.BuiltIn "::" then
394 optimizeExpr context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
396 else
398 Opt.Binop op
399 <$> optimizeExpr Nothing leftExpr
400 <*> optimizeExpr Nothing rightExpr
403-- left-associative ((x |> f) |> g)
404forwardApply :: Var.Canonical
405forwardApply =
406 Var.inCore ["Basics"] "|>"
409-- right-associative (g <| (f <| x))
410backwardApply :: Var.Canonical
411backwardApply =
412 Var.inCore ["Basics"] "<|"
415-- left-associative ((f >> g) >> h)
416forwardCompose :: Var.Canonical
417forwardCompose =
418 Var.inCore ["Basics"] ">>"
421-- right-associative (h << (g << f))
422backwardCompose :: Var.Canonical
423backwardCompose =
424 Var.inCore ["Basics"] "<<"
427type Assoc =
428 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
431left :: Assoc
432left =
433 flip (,)
436right :: Assoc
437right =
438 (,)
441collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
442collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
443 case expr of
444 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
445 let
446 (func, arg) = assoc leftExpr rightExpr
447 in
448 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
450 _ ->
451 makeRoot annExpr