git.iko.soy/elk
on 57490fd help sign in
elk/src/Optimize/DecisionTree.hs
606 lines 13 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)
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 Elm.Name as N
33import qualified Reporting.Annotation as A
37-- COMPILE CASES
40{-| Users of this module will mainly interact with this function. It takes
41some normal branches and gives out a decision tree that has "labels" at all
42the leafs and a dictionary that maps these "labels" to the code that should
43run.
45If 2 or more leaves point to the same label, we need to do some tricks in JS to
46make that work nicely. When is JS getting goto?! ;) That is outside the scope
47of this module though.
48-}
49compile :: [(Can.Pattern, Int)] -> DecisionTree
50compile rawBranches =
51 let
52 format (pattern, index) =
53 Branch index [(Empty, pattern)]
54 in
55 toDecisionTree (map format rawBranches)
59-- DECISION TREES
62data DecisionTree
63 = Match Int
64 | Decision
65 { _path :: Path
66 , _edges :: [(Test, DecisionTree)]
67 , _default :: Maybe DecisionTree
68 }
69 deriving (Eq)
72data Test
73 = IsCtor Int N.Name
74 | IsInt Int
75 | IsChr Text
76 | IsStr Text
77 deriving (Eq, Ord)
80data Path
81 = Index Int Path
82 | Field N.Name Path
83 | Empty
84 | Alias
85 deriving (Eq)
89-- PATH HELPERS
92add :: Path -> Path -> Path
93add path finalLink =
94 case path of
95 Empty ->
96 finalLink
98 Alias ->
99 error "nothing should be added to an alias path"
101 Index index subpath ->
102 Index index (add subpath finalLink)
104 Field name subpath ->
105 Field name (add subpath finalLink)
108subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
109subPositions path patterns =
110 zipWith
111 (\index pattern -> (add path (Index index Empty), pattern))
112 [0..]
113 patterns
117-- ACTUALLY BUILD DECISION TREES
120data Branch =
121 Branch
122 { _goal :: Int
123 , _patterns :: [(Path, Can.Pattern)]
124 }
127toDecisionTree :: [Branch] -> DecisionTree
128toDecisionTree rawBranches =
129 let
130 branches =
131 map flattenPatterns rawBranches
132 in
133 case checkForMatch branches of
134 Just goal ->
135 Match goal
137 Nothing ->
138 let
139 path =
140 pickPath branches
142 (edges, fallback) =
143 gatherEdges branches path
145 decisionEdges =
146 map (second toDecisionTree) edges
147 in
148 case (decisionEdges, fallback) of
149 ([(_tag, decisionTree)], []) ->
150 decisionTree
152 (_, []) ->
153 Decision path decisionEdges Nothing
155 ([], _ : _) ->
156 toDecisionTree fallback
158 (_, _) ->
159 Decision path decisionEdges (Just (toDecisionTree fallback))
162isComplete :: [Test] -> Bool
163isComplete tests =
164 case head tests of
165 IsCtor alts _ ->
166 alts == length tests
168 _ ->
169 False
173-- FLATTEN PATTERNS
176{-| Flatten type aliases and use the VariantDict to figure out when a tag is
177the only variant so we can skip doing any tests on it.
178-}
179flattenPatterns :: Branch -> Branch
180flattenPatterns (Branch goal pathPatterns) =
181 Branch goal (concatMap flatten pathPatterns)
184flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)]
185flatten pathPattern@(path, A.At region pattern) =
186 case pattern of
187 Can.PVar _ ->
188 [pathPattern]
190 Can.PAnything ->
191 [pathPattern]
193 Can.PAlias realPattern alias ->
194 (add path Alias, A.At region (Can.PVar alias))
195 :
196 flatten (path, realPattern)
198 Can.PRecord _ ->
199 [pathPattern]
201 Can.PUnit ->
202 [pathPattern]
204 Can.PTuple a b Nothing ->
205 concatMap flatten (subPositions path [a,b])
207 Can.PTuple a b (Just c) ->
208 concatMap flatten (subPositions path [a,b,c])
210 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) _ args ->
211 if numAlts == 1 then
212 concatMap flatten (subPositions path (map dearg args))
213 else
214 [pathPattern]
216 Can.PList _ ->
217 [pathPattern]
219 Can.PCons _ _ ->
220 [pathPattern]
222 Can.PChr _ ->
223 [pathPattern]
225 Can.PStr _ ->
226 [pathPattern]
228 Can.PInt _ ->
229 [pathPattern]
232dearg :: Can.PatternCtorArg -> Can.Pattern
233dearg (Can.PatternCtorArg _ _ pattern) =
234 pattern
238-- SUCCESSFULLY MATCH
241{-| If the first branch has no more "decision points" we can finally take that
242path. If that is the case we give the resulting label and a mapping from free
243variables to "how to get their value". So a pattern like (Just (x,_)) will give
244us something like ("x" => value.0.0)
245-}
246checkForMatch :: [Branch] -> Maybe Int
247checkForMatch branches =
248 case branches of
249 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
250 Just goal
252 _ ->
253 Nothing
257-- GATHER OUTGOING EDGES
260gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
261gatherEdges branches path =
262 let
263 relevantTests =
264 testsAtPath path branches
266 allEdges =
267 map (edgesFor path branches) relevantTests
269 fallbacks =
270 if isComplete relevantTests then
271 []
272 else
273 filter (isIrrelevantTo path) branches
274 in
275 ( allEdges, fallbacks )
279-- FIND RELEVANT TESTS
282testsAtPath :: Path -> [Branch] -> [Test]
283testsAtPath selectedPath branches =
284 let
285 allTests =
286 Maybe.mapMaybe (testAtPath selectedPath) branches
288 skipVisited test curr@(uniqueTests, visitedTests) =
289 if Set.member test visitedTests then
290 curr
291 else
292 ( test : uniqueTests
293 , Set.insert test visitedTests
294 )
295 in
296 fst (foldr skipVisited ([], Set.empty) allTests)
299testAtPath :: Path -> Branch -> Maybe Test
300testAtPath selectedPath (Branch _ pathPatterns) =
301 case List.lookup selectedPath pathPatterns of
302 Nothing ->
303 Nothing
305 Just (A.At _ pattern) ->
306 case pattern of
307 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) name _ ->
308 Just (IsCtor numAlts name)
310 Can.PInt int ->
311 Just (IsInt int)
313 Can.PStr str ->
314 Just (IsStr str)
316 Can.PChr chr ->
317 Just (IsChr chr)
319 Can.PList [] ->
320 Just nilTest
322 Can.PList _ ->
323 Just consTest
325 Can.PCons _ _ ->
326 Just consTest
328 Can.PVar _ ->
329 Nothing
331 Can.PAlias _ _ ->
332 error "aliases should never reach 'testAtPath' function"
334 Can.PAnything ->
335 Nothing
337 Can.PUnit ->
338 Nothing
340 Can.PTuple _ _ _ ->
341 Nothing
343 Can.PRecord _ ->
344 Nothing
347{-# NOINLINE nilTest #-}
348nilTest :: Test
349nilTest = IsCtor 2 "[]"
352{-# NOINLINE consTest #-}
353consTest :: Test
354consTest = IsCtor 2 "::"
358-- BUILD EDGES
361edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
362edgesFor path branches test =
363 ( test
364 , Maybe.mapMaybe (toRelevantBranch test path) branches
365 )
368toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
369toRelevantBranch test path branch@(Branch goal pathPatterns) =
370 case extract path pathPatterns of
371 Just (start, A.At region pattern, end) ->
372 case pattern of
373 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) name args ->
374 if test == IsCtor numAlts name then
375 Just (Branch goal (start ++ subPositions path (map dearg args) ++ end))
376 else
377 Nothing
379 Can.PList [] ->
380 if test == IsCtor 2 "[]" then
381 Just (Branch goal (start ++ end))
382 else
383 Nothing
385 Can.PList (hd:tl) ->
386 if test == IsCtor 2 "::" then
387 let tl' = A.At region (Can.PList tl) in
388 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
389 else
390 Nothing
392 Can.PCons hd tl ->
393 if test == IsCtor 2 "::" then
394 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
395 else
396 Nothing
398 Can.PChr chr ->
399 if test == IsChr chr then
400 Just (Branch goal (start ++ end))
401 else
402 Nothing
404 Can.PStr str ->
405 if test == IsStr str then
406 Just (Branch goal (start ++ end))
407 else
408 Nothing
410 Can.PInt int ->
411 if test == IsInt int then
412 Just (Branch goal (start ++ end))
413 else
414 Nothing
416 _ ->
417 Just branch
419 _ ->
420 Just branch
423extract
424 :: Path
425 -> [(Path, Can.Pattern)]
426 -> Maybe ([(Path, Can.Pattern)], Can.Pattern, [(Path, Can.Pattern)])
427extract selectedPath pathPatterns =
428 case pathPatterns of
429 [] ->
430 Nothing
432 first@(path, pattern) : rest ->
433 if path == selectedPath then
434 Just ([], pattern, rest)
436 else
437 case extract selectedPath rest of
438 Nothing ->
439 Nothing
441 Just (start, foundPattern, end) ->
442 Just (first : start, foundPattern, end)
446-- FIND IRRELEVANT BRANCHES
449isIrrelevantTo :: Path -> Branch -> Bool
450isIrrelevantTo selectedPath (Branch _ pathPatterns) =
451 case List.lookup selectedPath pathPatterns of
452 Nothing ->
453 True
455 Just pattern ->
456 not (needsTests pattern)
459needsTests :: Can.Pattern -> Bool
460needsTests (A.At _ pattern) =
461 case pattern of
462 Can.PVar _ ->
463 False
465 Can.PAnything ->
466 False
468 Can.PAlias _ _ ->
469 error "aliases should never reach 'isIrrelevantTo' function"
471 Can.PRecord _ ->
472 False
474 Can.PCtor _ _ _ _ _ _ ->
475 True
477 Can.PList _ ->
478 True
480 Can.PCons _ _ ->
481 True
483 Can.PUnit ->
484 False
486 Can.PTuple _ _ _ ->
487 False
489 Can.PChr _ ->
490 True
492 Can.PStr _ ->
493 True
495 Can.PInt _ ->
496 True
500-- PICK A PATH
503pickPath :: [Branch] -> Path
504pickPath branches =
505 let
506 allPaths =
507 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
508 in
509 case bests (addWeights (smallDefaults branches) allPaths) of
510 [path] ->
511 path
513 tiedPaths ->
514 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
517isChoicePath :: (Path, Can.Pattern) -> Maybe Path
518isChoicePath (path, pattern) =
519 if needsTests pattern then
520 Just path
521 else
522 Nothing
525addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
526addWeights toWeight paths =
527 map (\path -> (path, toWeight path)) paths
530bests :: [(Path, Int)] -> [Path]
531bests allPaths =
532 case allPaths of
533 [] ->
534 error "Cannot choose the best of zero paths. This should never happen."
536 (headPath, headWeight) : weightedPaths ->
537 let
538 gatherMinimum acc@(minWeight, paths) (path, weight) =
539 if weight == minWeight then
540 (minWeight, path : paths)
542 else if weight < minWeight then
543 (weight, [path])
545 else
546 acc
547 in
548 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
552-- PATH PICKING HEURISTICS
555smallDefaults :: [Branch] -> Path -> Int
556smallDefaults branches path =
557 length (filter (isIrrelevantTo path) branches)
560smallBranchingFactor :: [Branch] -> Path -> Int
561smallBranchingFactor branches path =
562 let
563 (edges, fallback) =
564 gatherEdges branches path
565 in
566 length edges + (if null fallback then 0 else 1)
570-- BINARY
573instance Binary Test where
574 put test =
575 case test of
576 IsCtor a b -> putWord8 0 >> put a >> put b
577 IsChr a -> putWord8 1 >> put a
578 IsStr a -> putWord8 2 >> put a
579 IsInt a -> putWord8 3 >> put a
581 get =
582 do word <- getWord8
583 case word of
584 0 -> liftM2 IsCtor get get
585 1 -> liftM IsChr get
586 2 -> liftM IsStr get
587 3 -> liftM IsInt get
588 _ -> error "problem getting DecisionTree.Test binary"
591instance Binary Path where
592 put path =
593 case path of
594 Index a b -> putWord8 0 >> put a >> put b
595 Field a b -> putWord8 1 >> put a >> put b
596 Empty -> putWord8 2
597 Alias -> putWord8 3
599 get =
600 do word <- getWord8
601 case word of
602 0 -> liftM2 Index get get
603 1 -> liftM2 Field get get
604 2 -> pure Empty
605 3 -> pure Alias
606 _ -> error "problem getting DecisionTree.Path binary"