git.iko.soy/elk
on 32ff950 help sign in
elk/compiler/src/Optimize/DecisionTree.hs
615 lines 14 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2{-# LANGUAGE OverloadedStrings #-}
3module Optimize.DecisionTree
4 ( DecisionTree(..)
5 , compile
6 , Path(..)
7 , Test(..)
8 )
9 where
12{- To learn more about how this works, definitely read through:
14 "When Do Match-Compilation Heuristics Matter?"
16by Kevin Scott and Norman Ramsey. The rough idea is that we start with a simple
17list of patterns and expressions, and then turn that into a "decision tree"
18that requires as few tests as possible to make it to a leaf. Read the paper, it
19explains this extraordinarily well! We are currently using the same heuristics
20as SML/NJ to get nice trees.
21-}
23import Control.Arrow (second)
24import Control.Monad (liftM, liftM2, liftM4)
25import Data.Binary
26import qualified Data.List as List
27import qualified Data.Maybe as Maybe
28import qualified Data.Set as Set
29import Data.Text (Text)
31import qualified AST.Canonical as Can
32import qualified Data.Index as Index
33import qualified Elm.Name as N
34import qualified Reporting.Annotation as A
38-- COMPILE CASES
41{-| Users of this module will mainly interact with this function. It takes
42some normal branches and gives out a decision tree that has "labels" at all
43the leafs and a dictionary that maps these "labels" to the code that should
44run.
46If 2 or more leaves point to the same label, we need to do some tricks in JS to
47make that work nicely. When is JS getting goto?! ;) That is outside the scope
48of this module though.
49-}
50compile :: [(Can.Pattern, Int)] -> DecisionTree
51compile rawBranches =
52 let
53 format (pattern, index) =
54 Branch index [(Empty, pattern)]
55 in
56 toDecisionTree (map format rawBranches)
60-- DECISION TREES
63data DecisionTree
64 = Match Int
65 | Decision
66 { _path :: Path
67 , _edges :: [(Test, DecisionTree)]
68 , _default :: Maybe DecisionTree
69 }
70 deriving (Eq)
73data Test
74 = IsCtor N.Name Index.ZeroBased Int Can.CtorOpts
75 | IsCons
76 | IsNil
77 | IsTuple
78 | IsInt Int
79 | IsChr Text
80 | IsStr Text
81 deriving (Eq, Ord)
84data Path
85 = Index Index.ZeroBased Path
86 | Unbox Path
87 | Empty
88 deriving (Eq)
92-- ACTUALLY BUILD DECISION TREES
95data Branch =
96 Branch
97 { _goal :: Int
98 , _patterns :: [(Path, Can.Pattern)]
99 }
102toDecisionTree :: [Branch] -> DecisionTree
103toDecisionTree rawBranches =
104 let
105 branches =
106 map flattenPatterns rawBranches
107 in
108 case checkForMatch branches of
109 Just goal ->
110 Match goal
112 Nothing ->
113 let
114 path =
115 pickPath branches
117 (edges, fallback) =
118 gatherEdges branches path
120 decisionEdges =
121 map (second toDecisionTree) edges
122 in
123 case (decisionEdges, fallback) of
124 ([(_tag, decisionTree)], []) ->
125 decisionTree
127 (_, []) ->
128 Decision path decisionEdges Nothing
130 ([], _ : _) ->
131 toDecisionTree fallback
133 (_, _) ->
134 Decision path decisionEdges (Just (toDecisionTree fallback))
137isComplete :: [Test] -> Bool
138isComplete tests =
139 case head tests of
140 IsCtor _ _ numAlts _ ->
141 numAlts == length tests
143 IsCons ->
144 length tests == 2
146 IsNil ->
147 length tests == 2
149 IsTuple ->
150 True
152 IsChr _ ->
153 False
155 IsStr _ ->
156 False
158 IsInt _ ->
159 False
163-- FLATTEN PATTERNS
166{-| Flatten type aliases and use the VariantDict to figure out when a tag is
167the only variant so we can skip doing any tests on it.
168-}
169flattenPatterns :: Branch -> Branch
170flattenPatterns (Branch goal pathPatterns) =
171 Branch goal (foldr flatten [] pathPatterns)
174flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)] -> [(Path, Can.Pattern)]
175flatten pathPattern@(path, A.At region pattern) otherPathPatterns =
176 case pattern of
177 Can.PVar _ ->
178 pathPattern : otherPathPatterns
180 Can.PAnything ->
181 pathPattern : otherPathPatterns
183 Can.PCtor _ _ (Can.Union _ _ numAlts _) _ _ ctorArgs ->
184 if numAlts == 1 then
185 case map dearg ctorArgs of
186 [arg] ->
187 flatten (Unbox path, arg) otherPathPatterns
189 args ->
190 foldr flatten otherPathPatterns (subPositions path args)
191 else
192 pathPattern : otherPathPatterns
194 Can.PTuple a b maybeC ->
195 flatten (Index Index.first path, a) $
196 flatten (Index Index.second path, b) $
197 case maybeC of
198 Nothing ->
199 otherPathPatterns
201 Just c ->
202 flatten (Index Index.third path, c) otherPathPatterns
204 Can.PUnit ->
205 otherPathPatterns
207 Can.PAlias realPattern alias ->
208 flatten (path, realPattern) $
209 (path, A.At region (Can.PVar alias)) : otherPathPatterns
211 Can.PRecord _ ->
212 pathPattern : otherPathPatterns
214 Can.PList _ ->
215 pathPattern : otherPathPatterns
217 Can.PCons _ _ ->
218 pathPattern : otherPathPatterns
220 Can.PChr _ ->
221 pathPattern : otherPathPatterns
223 Can.PStr _ ->
224 pathPattern : otherPathPatterns
226 Can.PInt _ ->
227 pathPattern : otherPathPatterns
230subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
231subPositions path patterns =
232 Index.indexedMap (\index pattern -> (Index index path, pattern)) patterns
235dearg :: Can.PatternCtorArg -> Can.Pattern
236dearg (Can.PatternCtorArg _ _ pattern) =
237 pattern
241-- SUCCESSFULLY MATCH
244{-| If the first branch has no more "decision points" we can finally take that
245path. If that is the case we give the resulting label and a mapping from free
246variables to "how to get their value". So a pattern like (Just (x,_)) will give
247us something like ("x" => value.0.0)
248-}
249checkForMatch :: [Branch] -> Maybe Int
250checkForMatch branches =
251 case branches of
252 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
253 Just goal
255 _ ->
256 Nothing
260-- GATHER OUTGOING EDGES
263gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
264gatherEdges branches path =
265 let
266 relevantTests =
267 testsAtPath path branches
269 allEdges =
270 map (edgesFor path branches) relevantTests
272 fallbacks =
273 if isComplete relevantTests then
274 []
275 else
276 filter (isIrrelevantTo path) branches
277 in
278 ( allEdges, fallbacks )
282-- FIND RELEVANT TESTS
285testsAtPath :: Path -> [Branch] -> [Test]
286testsAtPath selectedPath branches =
287 let
288 allTests =
289 Maybe.mapMaybe (testAtPath selectedPath) branches
291 skipVisited test curr@(uniqueTests, visitedTests) =
292 if Set.member test visitedTests then
293 curr
294 else
295 ( test : uniqueTests
296 , Set.insert test visitedTests
297 )
298 in
299 fst (foldr skipVisited ([], Set.empty) allTests)
302testAtPath :: Path -> Branch -> Maybe Test
303testAtPath selectedPath (Branch _ pathPatterns) =
304 case List.lookup selectedPath pathPatterns of
305 Nothing ->
306 Nothing
308 Just (A.At _ pattern) ->
309 case pattern of
310 Can.PCtor _ _ (Can.Union _ _ numAlts opts) name index _ ->
311 Just (IsCtor name index numAlts opts)
313 Can.PList ps ->
314 Just (case ps of { [] -> IsNil ; _ -> IsCons })
316 Can.PCons _ _ ->
317 Just IsCons
319 Can.PTuple _ _ _ ->
320 Just IsTuple
322 Can.PUnit ->
323 Just IsTuple
325 Can.PVar _ ->
326 Nothing
328 Can.PAnything ->
329 Nothing
331 Can.PInt int ->
332 Just (IsInt int)
334 Can.PStr str ->
335 Just (IsStr str)
337 Can.PChr chr ->
338 Just (IsChr chr)
340 Can.PRecord _ ->
341 Nothing
343 Can.PAlias _ _ ->
344 error "aliases should never reach 'testAtPath' function"
348-- BUILD EDGES
351edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
352edgesFor path branches test =
353 ( test
354 , Maybe.mapMaybe (toRelevantBranch test path) branches
355 )
358toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
359toRelevantBranch test path branch@(Branch goal pathPatterns) =
360 case extract path pathPatterns of
361 Found start (A.At region pattern) end ->
362 case pattern of
363 Can.PCtor _ _ (Can.Union _ _ numAlts _) name _ ctorArgs ->
364 case test of
365 IsCtor testName _ _ _ | name == testName ->
366 Just $ Branch goal $
367 case map dearg ctorArgs of
368 [arg] | numAlts == 1 ->
369 start ++ [(Unbox path, arg)] ++ end
371 args ->
372 start ++ subPositions path args ++ end
374 _ ->
375 Nothing
377 Can.PList [] ->
378 case test of
379 IsNil ->
380 Just (Branch goal (start ++ end))
382 _ ->
383 Nothing
385 Can.PList (hd:tl) ->
386 case test of
387 IsCons ->
388 let tl' = A.At region (Can.PList tl) in
389 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
391 _ ->
392 Nothing
394 Can.PCons hd tl ->
395 case test of
396 IsCons ->
397 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
399 _ ->
400 Nothing
402 Can.PChr chr ->
403 case test of
404 IsChr testChr | chr == testChr ->
405 Just (Branch goal (start ++ end))
406 _ ->
407 Nothing
409 Can.PStr str ->
410 case test of
411 IsStr testStr | str == testStr ->
412 Just (Branch goal (start ++ end))
414 _ ->
415 Nothing
417 Can.PInt int ->
418 case test of
419 IsInt testInt | int == testInt ->
420 Just (Branch goal (start ++ end))
422 _ ->
423 Nothing
425 Can.PUnit ->
426 Just (Branch goal (start ++ end))
428 Can.PTuple a b maybeC ->
429 Just (Branch goal (start ++ subPositions path (a : b : Maybe.maybeToList maybeC) ++ end))
431 Can.PVar _ ->
432 Just branch
434 Can.PAnything ->
435 Just branch
437 Can.PRecord _ ->
438 Just branch
440 Can.PAlias _ _ ->
441 Just branch
443 NotFound ->
444 Just branch
447data Extract
448 = NotFound
449 | Found [(Path, Can.Pattern)] Can.Pattern [(Path, Can.Pattern)]
452extract :: Path -> [(Path, Can.Pattern)] -> Extract
453extract selectedPath pathPatterns =
454 case pathPatterns of
455 [] ->
456 NotFound
458 first@(path, pattern) : rest ->
459 if path == selectedPath then
460 Found [] pattern rest
462 else
463 case extract selectedPath rest of
464 NotFound ->
465 NotFound
467 Found start foundPattern end ->
468 Found (first : start) foundPattern end
472-- FIND IRRELEVANT BRANCHES
475isIrrelevantTo :: Path -> Branch -> Bool
476isIrrelevantTo selectedPath (Branch _ pathPatterns) =
477 case List.lookup selectedPath pathPatterns of
478 Nothing ->
479 True
481 Just pattern ->
482 not (needsTests pattern)
485needsTests :: Can.Pattern -> Bool
486needsTests (A.At _ pattern) =
487 case pattern of
488 Can.PVar _ -> False
489 Can.PAnything -> False
490 Can.PRecord _ -> False
491 Can.PCtor _ _ _ _ _ _ -> True
492 Can.PList _ -> True
493 Can.PCons _ _ -> True
494 Can.PUnit -> True
495 Can.PTuple _ _ _ -> True
496 Can.PChr _ -> True
497 Can.PStr _ -> True
498 Can.PInt _ -> True
499 Can.PAlias _ _ ->
500 error "aliases should never reach 'isIrrelevantTo' function"
505-- PICK A PATH
508pickPath :: [Branch] -> Path
509pickPath branches =
510 let
511 allPaths =
512 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
513 in
514 case bests (addWeights (smallDefaults branches) allPaths) of
515 [path] ->
516 path
518 tiedPaths ->
519 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
522isChoicePath :: (Path, Can.Pattern) -> Maybe Path
523isChoicePath (path, pattern) =
524 if needsTests pattern then
525 Just path
526 else
527 Nothing
530addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
531addWeights toWeight paths =
532 map (\path -> (path, toWeight path)) paths
535bests :: [(Path, Int)] -> [Path]
536bests allPaths =
537 case allPaths of
538 [] ->
539 error "Cannot choose the best of zero paths. This should never happen."
541 (headPath, headWeight) : weightedPaths ->
542 let
543 gatherMinimum acc@(minWeight, paths) (path, weight) =
544 if weight == minWeight then
545 (minWeight, path : paths)
547 else if weight < minWeight then
548 (weight, [path])
550 else
551 acc
552 in
553 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
557-- PATH PICKING HEURISTICS
560smallDefaults :: [Branch] -> Path -> Int
561smallDefaults branches path =
562 length (filter (isIrrelevantTo path) branches)
565smallBranchingFactor :: [Branch] -> Path -> Int
566smallBranchingFactor branches path =
567 let
568 (edges, fallback) =
569 gatherEdges branches path
570 in
571 length edges + (if null fallback then 0 else 1)
575-- BINARY
578instance Binary Test where
579 put test =
580 case test of
581 IsCtor a b c d -> putWord8 0 >> put a >> put b >> put c >> put d
582 IsCons -> putWord8 1
583 IsNil -> putWord8 2
584 IsTuple -> putWord8 3
585 IsChr a -> putWord8 4 >> put a
586 IsStr a -> putWord8 5 >> put a
587 IsInt a -> putWord8 6 >> put a
589 get =
590 do word <- getWord8
591 case word of
592 0 -> liftM4 IsCtor get get get get
593 1 -> pure IsCons
594 2 -> pure IsNil
595 3 -> pure IsTuple
596 4 -> liftM IsChr get
597 5 -> liftM IsStr get
598 6 -> liftM IsInt get
599 _ -> error "problem getting DecisionTree.Test binary"
602instance Binary Path where
603 put path =
604 case path of
605 Index a b -> putWord8 0 >> put a >> put b
606 Unbox a -> putWord8 1 >> put a
607 Empty -> putWord8 2
609 get =
610 do word <- getWord8
611 case word of
612 0 -> liftM2 Index get get
613 1 -> liftM Unbox get
614 2 -> pure Empty
615 _ -> error "problem getting DecisionTree.Path binary"