git.iko.soy/elk
on 233cfc8 help sign in
elk/compiler/src/Optimize/DecisionTree.hs
616 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, liftM5)
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 AST.Module.Name as ModuleName
33import qualified Data.Index as Index
34import qualified Elm.Name as N
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 N.Name Index.ZeroBased Int Can.CtorOpts
76 | IsCons
77 | IsNil
78 | IsTuple
79 | IsInt Int
80 | IsChr Text
81 | IsStr Text
82 deriving (Eq, Ord)
85data Path
86 = Index Index.ZeroBased Path
87 | Unbox Path
88 | Empty
89 deriving (Eq)
93-- ACTUALLY BUILD DECISION TREES
96data Branch =
97 Branch
98 { _goal :: Int
99 , _patterns :: [(Path, Can.Pattern)]
100 }
103toDecisionTree :: [Branch] -> DecisionTree
104toDecisionTree rawBranches =
105 let
106 branches =
107 map flattenPatterns rawBranches
108 in
109 case checkForMatch branches of
110 Just goal ->
111 Match goal
113 Nothing ->
114 let
115 path =
116 pickPath branches
118 (edges, fallback) =
119 gatherEdges branches path
121 decisionEdges =
122 map (second toDecisionTree) edges
123 in
124 case (decisionEdges, fallback) of
125 ([(_tag, decisionTree)], []) ->
126 decisionTree
128 (_, []) ->
129 Decision path decisionEdges Nothing
131 ([], _ : _) ->
132 toDecisionTree fallback
134 (_, _) ->
135 Decision path decisionEdges (Just (toDecisionTree fallback))
138isComplete :: [Test] -> Bool
139isComplete tests =
140 case head tests of
141 IsCtor _ _ _ numAlts _ ->
142 numAlts == length tests
144 IsCons ->
145 length tests == 2
147 IsNil ->
148 length tests == 2
150 IsTuple ->
151 True
153 IsChr _ ->
154 False
156 IsStr _ ->
157 False
159 IsInt _ ->
160 False
164-- FLATTEN PATTERNS
167{-| Flatten type aliases and use the VariantDict to figure out when a tag is
168the only variant so we can skip doing any tests on it.
169-}
170flattenPatterns :: Branch -> Branch
171flattenPatterns (Branch goal pathPatterns) =
172 Branch goal (foldr flatten [] pathPatterns)
175flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)] -> [(Path, Can.Pattern)]
176flatten pathPattern@(path, A.At region pattern) otherPathPatterns =
177 case pattern of
178 Can.PVar _ ->
179 pathPattern : otherPathPatterns
181 Can.PAnything ->
182 pathPattern : otherPathPatterns
184 Can.PCtor _ _ (Can.Union _ _ numAlts _) _ _ ctorArgs ->
185 if numAlts == 1 then
186 case map dearg ctorArgs of
187 [arg] ->
188 flatten (Unbox path, arg) otherPathPatterns
190 args ->
191 foldr flatten otherPathPatterns (subPositions path args)
192 else
193 pathPattern : otherPathPatterns
195 Can.PTuple a b maybeC ->
196 flatten (Index Index.first path, a) $
197 flatten (Index Index.second path, b) $
198 case maybeC of
199 Nothing ->
200 otherPathPatterns
202 Just c ->
203 flatten (Index Index.third path, c) otherPathPatterns
205 Can.PUnit ->
206 otherPathPatterns
208 Can.PAlias realPattern alias ->
209 flatten (path, realPattern) $
210 (path, A.At region (Can.PVar alias)) : otherPathPatterns
212 Can.PRecord _ ->
213 pathPattern : otherPathPatterns
215 Can.PList _ ->
216 pathPattern : otherPathPatterns
218 Can.PCons _ _ ->
219 pathPattern : otherPathPatterns
221 Can.PChr _ ->
222 pathPattern : otherPathPatterns
224 Can.PStr _ ->
225 pathPattern : otherPathPatterns
227 Can.PInt _ ->
228 pathPattern : otherPathPatterns
231subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
232subPositions path patterns =
233 Index.indexedMap (\index pattern -> (Index index path, pattern)) patterns
236dearg :: Can.PatternCtorArg -> Can.Pattern
237dearg (Can.PatternCtorArg _ _ pattern) =
238 pattern
242-- SUCCESSFULLY MATCH
245{-| If the first branch has no more "decision points" we can finally take that
246path. If that is the case we give the resulting label and a mapping from free
247variables to "how to get their value". So a pattern like (Just (x,_)) will give
248us something like ("x" => value.0.0)
249-}
250checkForMatch :: [Branch] -> Maybe Int
251checkForMatch branches =
252 case branches of
253 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
254 Just goal
256 _ ->
257 Nothing
261-- GATHER OUTGOING EDGES
264gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
265gatherEdges branches path =
266 let
267 relevantTests =
268 testsAtPath path branches
270 allEdges =
271 map (edgesFor path branches) relevantTests
273 fallbacks =
274 if isComplete relevantTests then
275 []
276 else
277 filter (isIrrelevantTo path) branches
278 in
279 ( allEdges, fallbacks )
283-- FIND RELEVANT TESTS
286testsAtPath :: Path -> [Branch] -> [Test]
287testsAtPath selectedPath branches =
288 let
289 allTests =
290 Maybe.mapMaybe (testAtPath selectedPath) branches
292 skipVisited test curr@(uniqueTests, visitedTests) =
293 if Set.member test visitedTests then
294 curr
295 else
296 ( test : uniqueTests
297 , Set.insert test visitedTests
298 )
299 in
300 fst (foldr skipVisited ([], Set.empty) allTests)
303testAtPath :: Path -> Branch -> Maybe Test
304testAtPath selectedPath (Branch _ pathPatterns) =
305 case List.lookup selectedPath pathPatterns of
306 Nothing ->
307 Nothing
309 Just (A.At _ pattern) ->
310 case pattern of
311 Can.PCtor home _ (Can.Union _ _ numAlts opts) name index _ ->
312 Just (IsCtor home name index numAlts opts)
314 Can.PList ps ->
315 Just (case ps of { [] -> IsNil ; _ -> IsCons })
317 Can.PCons _ _ ->
318 Just IsCons
320 Can.PTuple _ _ _ ->
321 Just IsTuple
323 Can.PUnit ->
324 Just IsTuple
326 Can.PVar _ ->
327 Nothing
329 Can.PAnything ->
330 Nothing
332 Can.PInt int ->
333 Just (IsInt int)
335 Can.PStr str ->
336 Just (IsStr str)
338 Can.PChr chr ->
339 Just (IsChr chr)
341 Can.PRecord _ ->
342 Nothing
344 Can.PAlias _ _ ->
345 error "aliases should never reach 'testAtPath' function"
349-- BUILD EDGES
352edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
353edgesFor path branches test =
354 ( test
355 , Maybe.mapMaybe (toRelevantBranch test path) branches
356 )
359toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
360toRelevantBranch test path branch@(Branch goal pathPatterns) =
361 case extract path pathPatterns of
362 Found start (A.At region pattern) end ->
363 case pattern of
364 Can.PCtor _ _ (Can.Union _ _ numAlts _) name _ ctorArgs ->
365 case test of
366 IsCtor _ testName _ _ _ | name == testName ->
367 Just $ Branch goal $
368 case map dearg ctorArgs of
369 [arg] | numAlts == 1 ->
370 start ++ [(Unbox path, arg)] ++ end
372 args ->
373 start ++ subPositions path args ++ end
375 _ ->
376 Nothing
378 Can.PList [] ->
379 case test of
380 IsNil ->
381 Just (Branch goal (start ++ end))
383 _ ->
384 Nothing
386 Can.PList (hd:tl) ->
387 case test of
388 IsCons ->
389 let tl' = A.At region (Can.PList tl) in
390 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
392 _ ->
393 Nothing
395 Can.PCons hd tl ->
396 case test of
397 IsCons ->
398 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
400 _ ->
401 Nothing
403 Can.PChr chr ->
404 case test of
405 IsChr testChr | chr == testChr ->
406 Just (Branch goal (start ++ end))
407 _ ->
408 Nothing
410 Can.PStr str ->
411 case test of
412 IsStr testStr | str == testStr ->
413 Just (Branch goal (start ++ end))
415 _ ->
416 Nothing
418 Can.PInt int ->
419 case test of
420 IsInt testInt | int == testInt ->
421 Just (Branch goal (start ++ end))
423 _ ->
424 Nothing
426 Can.PUnit ->
427 Just (Branch goal (start ++ end))
429 Can.PTuple a b maybeC ->
430 Just (Branch goal (start ++ subPositions path (a : b : Maybe.maybeToList maybeC) ++ end))
432 Can.PVar _ ->
433 Just branch
435 Can.PAnything ->
436 Just branch
438 Can.PRecord _ ->
439 Just branch
441 Can.PAlias _ _ ->
442 Just branch
444 NotFound ->
445 Just branch
448data Extract
449 = NotFound
450 | Found [(Path, Can.Pattern)] Can.Pattern [(Path, Can.Pattern)]
453extract :: Path -> [(Path, Can.Pattern)] -> Extract
454extract selectedPath pathPatterns =
455 case pathPatterns of
456 [] ->
457 NotFound
459 first@(path, pattern) : rest ->
460 if path == selectedPath then
461 Found [] pattern rest
463 else
464 case extract selectedPath rest of
465 NotFound ->
466 NotFound
468 Found start foundPattern end ->
469 Found (first : start) foundPattern end
473-- FIND IRRELEVANT BRANCHES
476isIrrelevantTo :: Path -> Branch -> Bool
477isIrrelevantTo selectedPath (Branch _ pathPatterns) =
478 case List.lookup selectedPath pathPatterns of
479 Nothing ->
480 True
482 Just pattern ->
483 not (needsTests pattern)
486needsTests :: Can.Pattern -> Bool
487needsTests (A.At _ pattern) =
488 case pattern of
489 Can.PVar _ -> False
490 Can.PAnything -> False
491 Can.PRecord _ -> False
492 Can.PCtor _ _ _ _ _ _ -> True
493 Can.PList _ -> True
494 Can.PCons _ _ -> True
495 Can.PUnit -> True
496 Can.PTuple _ _ _ -> True
497 Can.PChr _ -> True
498 Can.PStr _ -> True
499 Can.PInt _ -> True
500 Can.PAlias _ _ ->
501 error "aliases should never reach 'isIrrelevantTo' function"
506-- PICK A PATH
509pickPath :: [Branch] -> Path
510pickPath branches =
511 let
512 allPaths =
513 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
514 in
515 case bests (addWeights (smallDefaults branches) allPaths) of
516 [path] ->
517 path
519 tiedPaths ->
520 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
523isChoicePath :: (Path, Can.Pattern) -> Maybe Path
524isChoicePath (path, pattern) =
525 if needsTests pattern then
526 Just path
527 else
528 Nothing
531addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
532addWeights toWeight paths =
533 map (\path -> (path, toWeight path)) paths
536bests :: [(Path, Int)] -> [Path]
537bests allPaths =
538 case allPaths of
539 [] ->
540 error "Cannot choose the best of zero paths. This should never happen."
542 (headPath, headWeight) : weightedPaths ->
543 let
544 gatherMinimum acc@(minWeight, paths) (path, weight) =
545 if weight == minWeight then
546 (minWeight, path : paths)
548 else if weight < minWeight then
549 (weight, [path])
551 else
552 acc
553 in
554 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
558-- PATH PICKING HEURISTICS
561smallDefaults :: [Branch] -> Path -> Int
562smallDefaults branches path =
563 length (filter (isIrrelevantTo path) branches)
566smallBranchingFactor :: [Branch] -> Path -> Int
567smallBranchingFactor branches path =
568 let
569 (edges, fallback) =
570 gatherEdges branches path
571 in
572 length edges + (if null fallback then 0 else 1)
576-- BINARY
579instance Binary Test where
580 put test =
581 case test of
582 IsCtor a b c d e -> putWord8 0 >> put a >> put b >> put c >> put d >> put e
583 IsCons -> putWord8 1
584 IsNil -> putWord8 2
585 IsTuple -> putWord8 3
586 IsChr a -> putWord8 4 >> put a
587 IsStr a -> putWord8 5 >> put a
588 IsInt a -> putWord8 6 >> put a
590 get =
591 do word <- getWord8
592 case word of
593 0 -> liftM5 IsCtor get get get get get
594 1 -> pure IsCons
595 2 -> pure IsNil
596 3 -> pure IsTuple
597 4 -> liftM IsChr get
598 5 -> liftM IsStr get
599 6 -> liftM IsInt get
600 _ -> error "problem getting DecisionTree.Test binary"
603instance Binary Path where
604 put path =
605 case path of
606 Index a b -> putWord8 0 >> put a >> put b
607 Unbox a -> putWord8 1 >> put a
608 Empty -> putWord8 2
610 get =
611 do word <- getWord8
612 case word of
613 0 -> liftM2 Index get get
614 1 -> liftM Unbox get
615 2 -> pure Empty
616 _ -> error "problem getting DecisionTree.Path binary"