git.iko.soy/elk
on 1238607 help sign in
elk/compiler/src/Optimize/DecisionTree.hs
637 lines 15 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, liftM5)
25import Data.Binary
26import qualified Data.List as List
27import qualified Data.Maybe as Maybe
28import qualified Data.Name as Name
29import qualified Data.Set as Set
30import qualified Data.Utf8 as Utf8
32import qualified AST.Canonical as Can
33import qualified Data.Index as Index
34import qualified Elm.ModuleName as ModuleName
35import qualified Reporting.Annotation as A
39-- COMPILE CASES
42{-| Users of this module will mainly interact with this function. It takes
43some normal branches and gives out a decision tree that has "labels" at all
44the leafs and a dictionary that maps these "labels" to the code that should
45run.
47If 2 or more leaves point to the same label, we need to do some tricks in JS to
48make that work nicely. When is JS getting goto?! ;) That is outside the scope
49of this module though.
50-}
51compile :: [(Can.Pattern, Int)] -> DecisionTree
52compile rawBranches =
53 let
54 format (pattern, index) =
55 Branch index [(Empty, pattern)]
56 in
57 toDecisionTree (map format rawBranches)
61-- DECISION TREES
64data DecisionTree
65 = Match Int
66 | Decision
67 { _path :: Path
68 , _edges :: [(Test, DecisionTree)]
69 , _default :: Maybe DecisionTree
70 }
71 deriving (Eq)
74data Test
75 = IsCtor ModuleName.Canonical Name.Name Index.ZeroBased Int Can.CtorOpts
76 | IsCons
77 | IsNil
78 | IsTuple
79 | IsInt Int
80 | IsChr Utf8.String
81 | IsStr Utf8.String
82 | IsBool Bool
83 deriving (Eq, Ord)
86data Path
87 = Index Index.ZeroBased Path
88 | Unbox Path
89 | Empty
90 deriving (Eq)
94-- ACTUALLY BUILD DECISION TREES
97data Branch =
98 Branch
99 { _goal :: Int
100 , _patterns :: [(Path, Can.Pattern)]
101 }
104toDecisionTree :: [Branch] -> DecisionTree
105toDecisionTree rawBranches =
106 let
107 branches =
108 map flattenPatterns rawBranches
109 in
110 case checkForMatch branches of
111 Just goal ->
112 Match goal
114 Nothing ->
115 let
116 path =
117 pickPath branches
119 (edges, fallback) =
120 gatherEdges branches path
122 decisionEdges =
123 map (second toDecisionTree) edges
124 in
125 case (decisionEdges, fallback) of
126 ([(_tag, decisionTree)], []) ->
127 decisionTree
129 (_, []) ->
130 Decision path decisionEdges Nothing
132 ([], _ : _) ->
133 toDecisionTree fallback
135 (_, _) ->
136 Decision path decisionEdges (Just (toDecisionTree fallback))
139isComplete :: [Test] -> Bool
140isComplete tests =
141 case head tests of
142 IsCtor _ _ _ numAlts _ ->
143 numAlts == length tests
145 IsCons ->
146 length tests == 2
148 IsNil ->
149 length tests == 2
151 IsTuple ->
152 True
154 IsChr _ ->
155 False
157 IsStr _ ->
158 False
160 IsInt _ ->
161 False
163 IsBool _ ->
164 length tests == 2
168-- FLATTEN PATTERNS
171{-| Flatten type aliases and use the VariantDict to figure out when a tag is
172the only variant so we can skip doing any tests on it.
173-}
174flattenPatterns :: Branch -> Branch
175flattenPatterns (Branch goal pathPatterns) =
176 Branch goal (foldr flatten [] pathPatterns)
179flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)] -> [(Path, Can.Pattern)]
180flatten pathPattern@(path, A.At region pattern) otherPathPatterns =
181 case pattern of
182 Can.PVar _ ->
183 pathPattern : otherPathPatterns
185 Can.PAnything ->
186 pathPattern : otherPathPatterns
188 Can.PCtor _ _ (Can.Union _ _ numAlts _) _ _ ctorArgs ->
189 if numAlts == 1 then
190 case map dearg ctorArgs of
191 [arg] ->
192 flatten (Unbox path, arg) otherPathPatterns
194 args ->
195 foldr flatten otherPathPatterns (subPositions path args)
196 else
197 pathPattern : otherPathPatterns
199 Can.PTuple a b maybeC ->
200 flatten (Index Index.first path, a) $
201 flatten (Index Index.second path, b) $
202 case maybeC of
203 Nothing ->
204 otherPathPatterns
206 Just c ->
207 flatten (Index Index.third path, c) otherPathPatterns
209 Can.PUnit ->
210 otherPathPatterns
212 Can.PAlias realPattern alias ->
213 flatten (path, realPattern) $
214 (path, A.At region (Can.PVar alias)) : otherPathPatterns
216 Can.PRecord _ ->
217 pathPattern : otherPathPatterns
219 Can.PList _ ->
220 pathPattern : otherPathPatterns
222 Can.PCons _ _ ->
223 pathPattern : otherPathPatterns
225 Can.PChr _ ->
226 pathPattern : otherPathPatterns
228 Can.PStr _ ->
229 pathPattern : otherPathPatterns
231 Can.PInt _ ->
232 pathPattern : otherPathPatterns
234 Can.PBool _ _ ->
235 pathPattern : otherPathPatterns
238subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
239subPositions path patterns =
240 Index.indexedMap (\index pattern -> (Index index path, pattern)) patterns
243dearg :: Can.PatternCtorArg -> Can.Pattern
244dearg (Can.PatternCtorArg _ _ pattern) =
245 pattern
249-- SUCCESSFULLY MATCH
252{-| If the first branch has no more "decision points" we can finally take that
253path. If that is the case we give the resulting label and a mapping from free
254variables to "how to get their value". So a pattern like (Just (x,_)) will give
255us something like ("x" => value.0.0)
256-}
257checkForMatch :: [Branch] -> Maybe Int
258checkForMatch branches =
259 case branches of
260 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
261 Just goal
263 _ ->
264 Nothing
268-- GATHER OUTGOING EDGES
271gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
272gatherEdges branches path =
273 let
274 relevantTests =
275 testsAtPath path branches
277 allEdges =
278 map (edgesFor path branches) relevantTests
280 fallbacks =
281 if isComplete relevantTests then
282 []
283 else
284 filter (isIrrelevantTo path) branches
285 in
286 ( allEdges, fallbacks )
290-- FIND RELEVANT TESTS
293testsAtPath :: Path -> [Branch] -> [Test]
294testsAtPath selectedPath branches =
295 let
296 allTests =
297 Maybe.mapMaybe (testAtPath selectedPath) branches
299 skipVisited test curr@(uniqueTests, visitedTests) =
300 if Set.member test visitedTests then
301 curr
302 else
303 ( test : uniqueTests
304 , Set.insert test visitedTests
305 )
306 in
307 fst (foldr skipVisited ([], Set.empty) allTests)
310testAtPath :: Path -> Branch -> Maybe Test
311testAtPath selectedPath (Branch _ pathPatterns) =
312 case List.lookup selectedPath pathPatterns of
313 Nothing ->
314 Nothing
316 Just (A.At _ pattern) ->
317 case pattern of
318 Can.PCtor home _ (Can.Union _ _ numAlts opts) name index _ ->
319 Just (IsCtor home name index numAlts opts)
321 Can.PList ps ->
322 Just (case ps of { [] -> IsNil ; _ -> IsCons })
324 Can.PCons _ _ ->
325 Just IsCons
327 Can.PTuple _ _ _ ->
328 Just IsTuple
330 Can.PUnit ->
331 Just IsTuple
333 Can.PVar _ ->
334 Nothing
336 Can.PAnything ->
337 Nothing
339 Can.PInt int ->
340 Just (IsInt int)
342 Can.PStr str ->
343 Just (IsStr str)
345 Can.PChr chr ->
346 Just (IsChr chr)
348 Can.PBool _ bool ->
349 Just (IsBool bool)
351 Can.PRecord _ ->
352 Nothing
354 Can.PAlias _ _ ->
355 error "aliases should never reach 'testAtPath' function"
359-- BUILD EDGES
362edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
363edgesFor path branches test =
364 ( test
365 , Maybe.mapMaybe (toRelevantBranch test path) branches
366 )
369toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
370toRelevantBranch test path branch@(Branch goal pathPatterns) =
371 case extract path pathPatterns of
372 Found start (A.At region pattern) end ->
373 case pattern of
374 Can.PCtor _ _ (Can.Union _ _ numAlts _) name _ ctorArgs ->
375 case test of
376 IsCtor _ testName _ _ _ | name == testName ->
377 Just $ Branch goal $
378 case map dearg ctorArgs of
379 [arg] | numAlts == 1 ->
380 start ++ [(Unbox path, arg)] ++ end
382 args ->
383 start ++ subPositions path args ++ end
385 _ ->
386 Nothing
388 Can.PList [] ->
389 case test of
390 IsNil ->
391 Just (Branch goal (start ++ end))
393 _ ->
394 Nothing
396 Can.PList (hd:tl) ->
397 case test of
398 IsCons ->
399 let tl' = A.At region (Can.PList tl) in
400 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
402 _ ->
403 Nothing
405 Can.PCons hd tl ->
406 case test of
407 IsCons ->
408 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
410 _ ->
411 Nothing
413 Can.PChr chr ->
414 case test of
415 IsChr testChr | chr == testChr ->
416 Just (Branch goal (start ++ end))
417 _ ->
418 Nothing
420 Can.PStr str ->
421 case test of
422 IsStr testStr | str == testStr ->
423 Just (Branch goal (start ++ end))
425 _ ->
426 Nothing
428 Can.PInt int ->
429 case test of
430 IsInt testInt | int == testInt ->
431 Just (Branch goal (start ++ end))
433 _ ->
434 Nothing
436 Can.PBool _ bool ->
437 case test of
438 IsBool testBool | bool == testBool ->
439 Just (Branch goal (start ++ end))
441 _ ->
442 Nothing
444 Can.PUnit ->
445 Just (Branch goal (start ++ end))
447 Can.PTuple a b maybeC ->
448 Just (Branch goal (start ++ subPositions path (a : b : Maybe.maybeToList maybeC) ++ end))
450 Can.PVar _ ->
451 Just branch
453 Can.PAnything ->
454 Just branch
456 Can.PRecord _ ->
457 Just branch
459 Can.PAlias _ _ ->
460 Just branch
462 NotFound ->
463 Just branch
466data Extract
467 = NotFound
468 | Found [(Path, Can.Pattern)] Can.Pattern [(Path, Can.Pattern)]
471extract :: Path -> [(Path, Can.Pattern)] -> Extract
472extract selectedPath pathPatterns =
473 case pathPatterns of
474 [] ->
475 NotFound
477 first@(path, pattern) : rest ->
478 if path == selectedPath then
479 Found [] pattern rest
481 else
482 case extract selectedPath rest of
483 NotFound ->
484 NotFound
486 Found start foundPattern end ->
487 Found (first : start) foundPattern end
491-- FIND IRRELEVANT BRANCHES
494isIrrelevantTo :: Path -> Branch -> Bool
495isIrrelevantTo selectedPath (Branch _ pathPatterns) =
496 case List.lookup selectedPath pathPatterns of
497 Nothing ->
498 True
500 Just pattern ->
501 not (needsTests pattern)
504needsTests :: Can.Pattern -> Bool
505needsTests (A.At _ pattern) =
506 case pattern of
507 Can.PVar _ -> False
508 Can.PAnything -> False
509 Can.PRecord _ -> False
510 Can.PCtor _ _ _ _ _ _ -> True
511 Can.PList _ -> True
512 Can.PCons _ _ -> True
513 Can.PUnit -> True
514 Can.PTuple _ _ _ -> True
515 Can.PChr _ -> True
516 Can.PStr _ -> True
517 Can.PInt _ -> True
518 Can.PBool _ _ -> True
519 Can.PAlias _ _ ->
520 error "aliases should never reach 'isIrrelevantTo' function"
525-- PICK A PATH
528pickPath :: [Branch] -> Path
529pickPath branches =
530 let
531 allPaths =
532 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
533 in
534 case bests (addWeights (smallDefaults branches) allPaths) of
535 [path] ->
536 path
538 tiedPaths ->
539 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
542isChoicePath :: (Path, Can.Pattern) -> Maybe Path
543isChoicePath (path, pattern) =
544 if needsTests pattern then
545 Just path
546 else
547 Nothing
550addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
551addWeights toWeight paths =
552 map (\path -> (path, toWeight path)) paths
555bests :: [(Path, Int)] -> [Path]
556bests allPaths =
557 case allPaths of
558 [] ->
559 error "Cannot choose the best of zero paths. This should never happen."
561 (headPath, headWeight) : weightedPaths ->
562 let
563 gatherMinimum acc@(minWeight, paths) (path, weight) =
564 if weight == minWeight then
565 (minWeight, path : paths)
567 else if weight < minWeight then
568 (weight, [path])
570 else
571 acc
572 in
573 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
577-- PATH PICKING HEURISTICS
580smallDefaults :: [Branch] -> Path -> Int
581smallDefaults branches path =
582 length (filter (isIrrelevantTo path) branches)
585smallBranchingFactor :: [Branch] -> Path -> Int
586smallBranchingFactor branches path =
587 let
588 (edges, fallback) =
589 gatherEdges branches path
590 in
591 length edges + (if null fallback then 0 else 1)
595-- BINARY
598instance Binary Test where
599 put test =
600 case test of
601 IsCtor a b c d e -> putWord8 0 >> put a >> put b >> put c >> put d >> put e
602 IsCons -> putWord8 1
603 IsNil -> putWord8 2
604 IsTuple -> putWord8 3
605 IsChr a -> putWord8 4 >> put a
606 IsStr a -> putWord8 5 >> put a
607 IsInt a -> putWord8 6 >> put a
608 IsBool a -> putWord8 7 >> put a
610 get =
611 do word <- getWord8
612 case word of
613 0 -> liftM5 IsCtor get get get get get
614 1 -> pure IsCons
615 2 -> pure IsNil
616 3 -> pure IsTuple
617 4 -> liftM IsChr get
618 5 -> liftM IsStr get
619 6 -> liftM IsInt get
620 7 -> liftM IsBool get
621 _ -> error "problem getting DecisionTree.Test binary"
624instance Binary Path where
625 put path =
626 case path of
627 Index a b -> putWord8 0 >> put a >> put b
628 Unbox a -> putWord8 1 >> put a
629 Empty -> putWord8 2
631 get =
632 do word <- getWord8
633 case word of
634 0 -> liftM2 Index get get
635 1 -> liftM Unbox get
636 2 -> pure Empty
637 _ -> error "problem getting DecisionTree.Path binary"