git.iko.soy/elk
on 0cd50f9 help sign in
elk/src/Optimize.hs
487 lines 12 kB history raw
1module Optimize (optimize) where
3import Control.Applicative
4import qualified Control.Monad as M
5import qualified Control.Monad.State as State
6import qualified Data.Map as Map
7import qualified Data.Traversable as T
9import qualified AST.Expression.General as Expr
10import qualified AST.Expression.Canonical as Can
11import qualified AST.Expression.Optimized as Opt
12import qualified AST.Module as Module
13import qualified AST.Module.Name as ModuleName
14import qualified AST.Pattern as P
15import qualified AST.Variable as Var
16import qualified Optimize.Cases as Cases
17import qualified Reporting.Annotation as A
18import qualified Reporting.Region as R
21-- OPTIMIZE FOR TAIL CALLS
23optimize :: Module.CanonicalModule -> Module.Optimized
24optimize module_ =
25 let
26 home =
27 Module.name module_
29 (defs, _) =
30 flattenLets [] (Module.program (Module.body module_))
31 in
32 State.evalState
33 (concat <$> mapM (definition (Just home)) defs)
34 (Env False 0)
37flattenLets :: [Can.Def] -> Can.Expr -> ([Can.Def], Can.Expr)
38flattenLets defs annExpr@(A.A _ expr) =
39 case expr of
40 Expr.Let ds body ->
41 flattenLets (defs ++ ds) body
43 _ ->
44 (defs, annExpr)
47-- TRACKER TO THREAD FINDINGS THROUGH TRAVERSAL
49type Optimizer a =
50 State.State Env a
53data Env = Env
54 { _hasTailCall :: Bool
55 , _uid :: Int
56 }
59setTailCall :: Bool -> Optimizer ()
60setTailCall bool =
61 do uid <- State.gets _uid
62 State.put (Env bool uid)
65mapSnd :: M.Functor box => (a -> box b) -> (x, a) -> box (x, b)
66mapSnd func (x, a) =
67 (,) x <$> func a
70-- CONVERT DEFINITIONS
72definition :: Maybe ModuleName.Canonical -> Can.Def -> Optimizer [Opt.Def]
73definition home (Can.Definition (Can.Facts deps) pattern@(A.A _ ptrn) expression _) =
74 let
75 (args, rawBody) =
76 Expr.collectLambdas expression
77 in
78 case (ptrn, args) of
79 (P.Var name, _ : _) ->
81 do (Env htc uid) <- State.get
82 State.put (Env False uid)
84 (args, body) <- optimizeFunction (Just name) args rawBody
86 (Env isTailCall uid') <- State.get
87 State.put (Env htc uid')
89 let facts = Opt.Facts home deps
91 return $
92 if isTailCall then
94 [ Opt.TailDef facts name args body ]
96 else
98 [ Opt.Def facts name (Opt.Function args body) ]
100 (P.Var name, []) ->
102 do hasTailCall <- State.gets _hasTailCall
103 body <- findTailCalls Nothing rawBody
104 setTailCall hasTailCall
105 return [ Opt.Def (Opt.Facts home deps) name body ]
107 (_, []) ->
109 do name <- freshName
110 optBody <- findTailCalls Nothing rawBody
111 let optDef = Opt.Def (Opt.Facts home deps) name optBody
112 return (optDef : patternToDefs home name pattern)
114 _ ->
115 error "there should never be a function where the name is not a P.Var"
118-- OPTIMIZE FUNCTION
120optimizeFunction
121 :: Maybe String
122 -> [P.CanonicalPattern]
123 -> Can.Expr
124 -> Optimizer ([String], Opt.Expr)
125optimizeFunction maybeName patterns givenBody =
126 do (argNames, body) <- M.foldM depatternArgs ([], givenBody) (reverse patterns)
127 optBody <- findTailCalls (makeContext argNames =<< maybeName) body
128 return (argNames, optBody)
131makeContext :: [String] -> String -> Context
132makeContext argNames name =
133 Just (name, length argNames, argNames)
136depatternArgs
137 :: ([String], Can.Expr)
138 -> P.CanonicalPattern
139 -> Optimizer ([String], Can.Expr)
140depatternArgs (names, rawExpr) ptrn =
141 do (name, expr) <- depattern ptrn rawExpr
142 return (name : names, expr)
145patternToDefs
146 :: Maybe ModuleName.Canonical
147 -> String
148 -> P.CanonicalPattern
149 -> [Opt.Def]
150patternToDefs home rootName pattern =
151 let
152 (deps, rootExpr) =
153 case home of
154 Nothing ->
155 ( []
156 , Opt.Var (Var.local rootName)
157 )
159 Just moduleName ->
160 ( [ Var.TopLevelVar moduleName rootName ]
161 , Opt.Var (Var.topLevel moduleName rootName)
162 )
164 toDef (name, accessExpr) =
165 Opt.Def (Opt.Facts home deps) name accessExpr
166 in
167 map toDef (substitutions rootExpr pattern)
170substitutions :: Opt.Expr -> P.CanonicalPattern -> [(String, Opt.Expr)]
171substitutions expr (A.A _ pattern) =
172 case pattern of
173 P.Var name ->
174 [ (name, expr) ]
176 P.Literal _ ->
177 []
179 P.Record fields ->
180 map (\name -> (name, Opt.Access expr name)) fields
182 P.Anything ->
183 []
185 P.Alias alias realPattern ->
186 (alias, expr) : substitutions expr realPattern
188 P.Data _ patterns ->
189 concat (zipWith (substitutions . Opt.DataAccess expr) [0..] patterns)
192-- CONVERT EXPRESSIONS
194type Context = Maybe (String, Int, [String])
197findTailCalls :: Context -> Can.Expr -> Optimizer Opt.Expr
198findTailCalls context annExpr@(A.A region expression) =
199 let
200 keepLooking =
201 findTailCalls context
202 -- this *might* find a tail call, use this when something is a potential
203 -- tail call
205 justConvert =
206 findTailCalls Nothing
207 -- this will never find tail calls, use this when you are working with
208 -- expressions that cannot be turned into tail calls
209 in
210 case expression of
211 Expr.Literal lit ->
212 pure (Opt.Literal lit)
214 Expr.Var name ->
215 pure (Opt.Var name)
217 Expr.Range lowExpr highExpr ->
218 Opt.Range
219 <$> justConvert lowExpr
220 <*> justConvert highExpr
222 Expr.ExplicitList elements ->
223 Opt.ExplicitList
224 <$> T.traverse justConvert elements
226 Expr.Binop op leftExpr rightExpr ->
227 optimizeBinop context region op leftExpr rightExpr
229 Expr.Lambda _ _ ->
230 let
231 (patterns, body) =
232 Expr.collectLambdas annExpr
233 in
234 uncurry Opt.Function <$> optimizeFunction Nothing patterns body
236 Expr.App _ _ ->
237 let
238 (func:args) =
239 Expr.collectApps annExpr
240 in
241 case isTailCall context func (length args) of
242 Just (name, argNames) ->
243 do optArgs <- T.traverse justConvert args
244 setTailCall True
245 return (Opt.TailCall name argNames optArgs)
247 Nothing ->
248 do hasTailCall <- State.gets _hasTailCall
249 optFunc <- justConvert func
250 optArgs <- T.traverse justConvert args
251 setTailCall hasTailCall
252 return (Opt.Call optFunc optArgs)
255 Expr.If branches finally ->
256 let
257 crawlBranch (cond,branch) =
258 (,) <$> justConvert cond <*> keepLooking branch
259 in
260 Opt.If
261 <$> T.traverse crawlBranch branches
262 <*> keepLooking finally
264 Expr.Let defs body ->
265 do optDefs <- concat <$> T.traverse (definition Nothing) defs
266 Opt.Let optDefs <$> keepLooking body
268 Expr.Case expr rawBranches ->
269 let
270 indexify index (pattern, branch) =
271 ( (pattern, index)
272 , (index, branch)
273 )
275 (patterns, indexedBranches) =
276 unzip (zipWith indexify [0..] rawBranches)
278 decisionTree =
279 Cases.compile (error "need a VariantDict") patterns
281 branches =
282 Map.fromList indexedBranches
283 in
284 Opt.Case
285 <$> justConvert expr
286 <*> pure decisionTree
287 <*> T.traverse keepLooking branches
289 Expr.Data name exprs ->
290 Opt.Data name <$> T.traverse justConvert exprs
292 Expr.Access record field ->
293 Opt.Access
294 <$> justConvert record
295 <*> pure field
297 Expr.Update record fields ->
298 Opt.Update
299 <$> justConvert record
300 <*> T.traverse (mapSnd justConvert) fields
302 Expr.Record fields ->
303 Opt.Record
304 <$> T.traverse (mapSnd justConvert) fields
306 Expr.GLShader uid src gltipe ->
307 pure (Opt.GLShader uid src gltipe)
309 Expr.Port impl ->
310 Opt.Port <$>
311 case impl of
312 Expr.In name tipe ->
313 pure (Expr.In name tipe)
315 Expr.Out name expr tipe ->
316 (\e -> Expr.Out name e tipe) <$> justConvert expr
318 Expr.Task name expr tipe ->
319 (\e -> Expr.Task name e tipe) <$> justConvert expr
322-- DETECT TAIL CALL
324isTailCall :: Context -> Can.Expr -> Int -> Maybe (String, [String])
325isTailCall context (A.A _ function) arity =
326 case (context, function) of
327 ( Just (ctxName, ctxArity, argNames), Expr.Var (Var.Canonical home name) ) ->
329 if name == ctxName && arity == ctxArity && Var.isLocalHome home then
330 Just (name, argNames)
332 else
333 Nothing
335 _ ->
336 Nothing
340-- DEPATTERN
341-- given a pattern and an expression, push the actual pattern matching into the
342-- expression as a case-expression if necessary. The pattern compiler will work
343-- it out from there
345depattern :: P.CanonicalPattern -> Can.Expr -> Optimizer (String, Can.Expr)
346depattern canPattern@(A.A _ pattern) expr =
347 let
348 ann =
349 A.A (error "the annotation added in 'depattern' should not be observed!")
351 caseExpr e branches =
352 ann (Expr.Case e branches)
354 var name =
355 ann (Expr.Var (Var.local name))
357 caseify =
358 do name <- freshName
359 return
360 ( name
361 , caseExpr (var name) [ (canPattern, expr) ]
362 )
363 in
364 case pattern of
365 P.Var name ->
366 return (name, expr)
368 P.Literal _ ->
369 caseify
371 P.Record _ ->
372 caseify
374 P.Anything ->
375 do name <- freshName
376 return (name, expr)
378 P.Alias _ _ ->
379 caseify
381 P.Data _ _ ->
382 caseify
385freshName :: Optimizer String
386freshName =
387 do (Env htc uid) <- State.get
388 State.put (Env htc (uid + 1))
389 return ("_p" ++ show uid)
392-- OPTIMIZE BINOPS
394optimizeBinop
395 :: Context
396 -> R.Region
397 -> Var.Canonical
398 -> Can.Expr
399 -> Can.Expr
400 -> Optimizer Opt.Expr
401optimizeBinop context region op leftExpr rightExpr =
402 let
403 ann = A.A region
404 binop = ann (Expr.Binop op leftExpr rightExpr)
405 in
406 if op == forwardApply then
408 findTailCalls context (collect id left forwardApply binop)
410 else if op == backwardApply then
412 findTailCalls context (collect id right backwardApply binop)
414 else if op == forwardCompose then
416 do var <- freshName
417 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
418 let body = collect makeRoot left forwardCompose binop
419 findTailCalls context (ann (Expr.Lambda (ann (P.Var var)) body))
421 else if op == backwardCompose then
423 do var <- freshName
424 let makeRoot func = ann (Expr.App func (ann (Expr.Var (Var.local var))))
425 let body = collect makeRoot right backwardCompose binop
426 findTailCalls context (ann (Expr.Lambda (ann (P.Var var)) body))
428 else if op == Var.Canonical Var.BuiltIn "::" then
430 findTailCalls context (ann (Expr.Data "::" [ leftExpr, rightExpr ]))
432 else
434 Opt.Binop op
435 <$> findTailCalls Nothing leftExpr
436 <*> findTailCalls Nothing rightExpr
439-- left-associative ((x |> f) |> g)
440forwardApply :: Var.Canonical
441forwardApply =
442 Var.inCore ["Basics"] "|>"
445-- right-associative (g <| (f <| x))
446backwardApply :: Var.Canonical
447backwardApply =
448 Var.inCore ["Basics"] "<|"
451-- left-associative ((f >> g) >> h)
452forwardCompose :: Var.Canonical
453forwardCompose =
454 Var.inCore ["Basics"] ">>"
457-- right-associative (h << (g << f))
458backwardCompose :: Var.Canonical
459backwardCompose =
460 Var.inCore ["Basics"] "<<"
463type Assoc =
464 Can.Expr -> Can.Expr -> (Can.Expr, Can.Expr)
467left :: Assoc
468left =
469 flip (,)
472right :: Assoc
473right =
474 (,)
477collect :: (Can.Expr -> Can.Expr) -> Assoc -> Var.Canonical -> Can.Expr -> Can.Expr
478collect makeRoot assoc desiredOp annExpr@(A.A ann expr) =
479 case expr of
480 Expr.Binop op leftExpr rightExpr | op == desiredOp ->
481 let
482 (func, arg) = assoc leftExpr rightExpr
483 in
484 A.A ann (Expr.App func (collect makeRoot assoc desiredOp arg))
486 _ ->
487 makeRoot annExpr