git.iko.soy/elk
on 8a89c86 help sign in
elk/src/Optimize/DecisionTree.hs
607 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, liftM3)
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 Int N.Name Index.ZeroBased
75 | IsInt Int
76 | IsChr Text
77 | IsStr Text
78 deriving (Eq, Ord)
81data Path
82 = Index Int Path
83 | Field N.Name Path
84 | Empty
85 | Alias
86 deriving (Eq)
90-- PATH HELPERS
93add :: Path -> Path -> Path
94add path finalLink =
95 case path of
96 Empty ->
97 finalLink
99 Alias ->
100 error "nothing should be added to an alias path"
102 Index index subpath ->
103 Index index (add subpath finalLink)
105 Field name subpath ->
106 Field name (add subpath finalLink)
109subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
110subPositions path patterns =
111 zipWith
112 (\index pattern -> (add path (Index index Empty), pattern))
113 [0..]
114 patterns
118-- ACTUALLY BUILD DECISION TREES
121data Branch =
122 Branch
123 { _goal :: Int
124 , _patterns :: [(Path, Can.Pattern)]
125 }
128toDecisionTree :: [Branch] -> DecisionTree
129toDecisionTree rawBranches =
130 let
131 branches =
132 map flattenPatterns rawBranches
133 in
134 case checkForMatch branches of
135 Just goal ->
136 Match goal
138 Nothing ->
139 let
140 path =
141 pickPath branches
143 (edges, fallback) =
144 gatherEdges branches path
146 decisionEdges =
147 map (second toDecisionTree) edges
148 in
149 case (decisionEdges, fallback) of
150 ([(_tag, decisionTree)], []) ->
151 decisionTree
153 (_, []) ->
154 Decision path decisionEdges Nothing
156 ([], _ : _) ->
157 toDecisionTree fallback
159 (_, _) ->
160 Decision path decisionEdges (Just (toDecisionTree fallback))
163isComplete :: [Test] -> Bool
164isComplete tests =
165 case head tests of
166 IsCtor alts _ _ ->
167 alts == length tests
169 _ ->
170 False
174-- FLATTEN PATTERNS
177{-| Flatten type aliases and use the VariantDict to figure out when a tag is
178the only variant so we can skip doing any tests on it.
179-}
180flattenPatterns :: Branch -> Branch
181flattenPatterns (Branch goal pathPatterns) =
182 Branch goal (concatMap flatten pathPatterns)
185flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)]
186flatten pathPattern@(path, A.At region pattern) =
187 case pattern of
188 Can.PVar _ ->
189 [pathPattern]
191 Can.PAnything ->
192 [pathPattern]
194 Can.PAlias realPattern alias ->
195 (add path Alias, A.At region (Can.PVar alias))
196 :
197 flatten (path, realPattern)
199 Can.PRecord _ ->
200 [pathPattern]
202 Can.PUnit ->
203 [pathPattern]
205 Can.PTuple a b Nothing ->
206 concatMap flatten (subPositions path [a,b])
208 Can.PTuple a b (Just c) ->
209 concatMap flatten (subPositions path [a,b,c])
211 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) _ _ args ->
212 if numAlts == 1 then
213 concatMap flatten (subPositions path (map dearg args))
214 else
215 [pathPattern]
217 Can.PList _ ->
218 [pathPattern]
220 Can.PCons _ _ ->
221 [pathPattern]
223 Can.PChr _ ->
224 [pathPattern]
226 Can.PStr _ ->
227 [pathPattern]
229 Can.PInt _ ->
230 [pathPattern]
233dearg :: Can.PatternCtorArg -> Can.Pattern
234dearg (Can.PatternCtorArg _ _ pattern) =
235 pattern
239-- SUCCESSFULLY MATCH
242{-| If the first branch has no more "decision points" we can finally take that
243path. If that is the case we give the resulting label and a mapping from free
244variables to "how to get their value". So a pattern like (Just (x,_)) will give
245us something like ("x" => value.0.0)
246-}
247checkForMatch :: [Branch] -> Maybe Int
248checkForMatch branches =
249 case branches of
250 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
251 Just goal
253 _ ->
254 Nothing
258-- GATHER OUTGOING EDGES
261gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
262gatherEdges branches path =
263 let
264 relevantTests =
265 testsAtPath path branches
267 allEdges =
268 map (edgesFor path branches) relevantTests
270 fallbacks =
271 if isComplete relevantTests then
272 []
273 else
274 filter (isIrrelevantTo path) branches
275 in
276 ( allEdges, fallbacks )
280-- FIND RELEVANT TESTS
283testsAtPath :: Path -> [Branch] -> [Test]
284testsAtPath selectedPath branches =
285 let
286 allTests =
287 Maybe.mapMaybe (testAtPath selectedPath) branches
289 skipVisited test curr@(uniqueTests, visitedTests) =
290 if Set.member test visitedTests then
291 curr
292 else
293 ( test : uniqueTests
294 , Set.insert test visitedTests
295 )
296 in
297 fst (foldr skipVisited ([], Set.empty) allTests)
300testAtPath :: Path -> Branch -> Maybe Test
301testAtPath selectedPath (Branch _ pathPatterns) =
302 case List.lookup selectedPath pathPatterns of
303 Nothing ->
304 Nothing
306 Just (A.At _ pattern) ->
307 case pattern of
308 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) name index _ ->
309 Just (IsCtor numAlts name index)
311 Can.PInt int ->
312 Just (IsInt int)
314 Can.PStr str ->
315 Just (IsStr str)
317 Can.PChr chr ->
318 Just (IsChr chr)
320 Can.PList [] ->
321 Just nilTest
323 Can.PList _ ->
324 Just consTest
326 Can.PCons _ _ ->
327 Just consTest
329 Can.PVar _ ->
330 Nothing
332 Can.PAlias _ _ ->
333 error "aliases should never reach 'testAtPath' function"
335 Can.PAnything ->
336 Nothing
338 Can.PUnit ->
339 Nothing
341 Can.PTuple _ _ _ ->
342 Nothing
344 Can.PRecord _ ->
345 Nothing
348{-# NOINLINE nilTest #-}
349nilTest :: Test
350nilTest = IsCtor 2 "[]" (Index.unsafe (-1))
353{-# NOINLINE consTest #-}
354consTest :: Test
355consTest = IsCtor 2 "::" (Index.unsafe (-2))
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 Just (start, A.At region pattern, end) ->
373 case pattern of
374 Can.PCtor _ _ _ (Can.CtorAlts numAlts _) name index args ->
375 if test == IsCtor numAlts name index then
376 Just (Branch goal (start ++ subPositions path (map dearg args) ++ end))
377 else
378 Nothing
380 Can.PList [] ->
381 if test == nilTest then
382 Just (Branch goal (start ++ end))
383 else
384 Nothing
386 Can.PList (hd:tl) ->
387 if test == consTest then
388 let tl' = A.At region (Can.PList tl) in
389 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
390 else
391 Nothing
393 Can.PCons hd tl ->
394 if test == consTest then
395 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
396 else
397 Nothing
399 Can.PChr chr ->
400 if test == IsChr chr then
401 Just (Branch goal (start ++ end))
402 else
403 Nothing
405 Can.PStr str ->
406 if test == IsStr str then
407 Just (Branch goal (start ++ end))
408 else
409 Nothing
411 Can.PInt int ->
412 if test == IsInt int then
413 Just (Branch goal (start ++ end))
414 else
415 Nothing
417 _ ->
418 Just branch
420 _ ->
421 Just branch
424extract
425 :: Path
426 -> [(Path, Can.Pattern)]
427 -> Maybe ([(Path, Can.Pattern)], Can.Pattern, [(Path, Can.Pattern)])
428extract selectedPath pathPatterns =
429 case pathPatterns of
430 [] ->
431 Nothing
433 first@(path, pattern) : rest ->
434 if path == selectedPath then
435 Just ([], pattern, rest)
437 else
438 case extract selectedPath rest of
439 Nothing ->
440 Nothing
442 Just (start, foundPattern, end) ->
443 Just (first : start, foundPattern, end)
447-- FIND IRRELEVANT BRANCHES
450isIrrelevantTo :: Path -> Branch -> Bool
451isIrrelevantTo selectedPath (Branch _ pathPatterns) =
452 case List.lookup selectedPath pathPatterns of
453 Nothing ->
454 True
456 Just pattern ->
457 not (needsTests pattern)
460needsTests :: Can.Pattern -> Bool
461needsTests (A.At _ pattern) =
462 case pattern of
463 Can.PVar _ ->
464 False
466 Can.PAnything ->
467 False
469 Can.PAlias _ _ ->
470 error "aliases should never reach 'isIrrelevantTo' function"
472 Can.PRecord _ ->
473 False
475 Can.PCtor _ _ _ _ _ _ _ ->
476 True
478 Can.PList _ ->
479 True
481 Can.PCons _ _ ->
482 True
484 Can.PUnit ->
485 False
487 Can.PTuple _ _ _ ->
488 False
490 Can.PChr _ ->
491 True
493 Can.PStr _ ->
494 True
496 Can.PInt _ ->
497 True
501-- PICK A PATH
504pickPath :: [Branch] -> Path
505pickPath branches =
506 let
507 allPaths =
508 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
509 in
510 case bests (addWeights (smallDefaults branches) allPaths) of
511 [path] ->
512 path
514 tiedPaths ->
515 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
518isChoicePath :: (Path, Can.Pattern) -> Maybe Path
519isChoicePath (path, pattern) =
520 if needsTests pattern then
521 Just path
522 else
523 Nothing
526addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
527addWeights toWeight paths =
528 map (\path -> (path, toWeight path)) paths
531bests :: [(Path, Int)] -> [Path]
532bests allPaths =
533 case allPaths of
534 [] ->
535 error "Cannot choose the best of zero paths. This should never happen."
537 (headPath, headWeight) : weightedPaths ->
538 let
539 gatherMinimum acc@(minWeight, paths) (path, weight) =
540 if weight == minWeight then
541 (minWeight, path : paths)
543 else if weight < minWeight then
544 (weight, [path])
546 else
547 acc
548 in
549 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
553-- PATH PICKING HEURISTICS
556smallDefaults :: [Branch] -> Path -> Int
557smallDefaults branches path =
558 length (filter (isIrrelevantTo path) branches)
561smallBranchingFactor :: [Branch] -> Path -> Int
562smallBranchingFactor branches path =
563 let
564 (edges, fallback) =
565 gatherEdges branches path
566 in
567 length edges + (if null fallback then 0 else 1)
571-- BINARY
574instance Binary Test where
575 put test =
576 case test of
577 IsCtor a b c -> putWord8 0 >> put a >> put b >> put c
578 IsChr a -> putWord8 1 >> put a
579 IsStr a -> putWord8 2 >> put a
580 IsInt a -> putWord8 3 >> put a
582 get =
583 do word <- getWord8
584 case word of
585 0 -> liftM3 IsCtor get get get
586 1 -> liftM IsChr get
587 2 -> liftM IsStr get
588 3 -> liftM IsInt get
589 _ -> error "problem getting DecisionTree.Test binary"
592instance Binary Path where
593 put path =
594 case path of
595 Index a b -> putWord8 0 >> put a >> put b
596 Field a b -> putWord8 1 >> put a >> put b
597 Empty -> putWord8 2
598 Alias -> putWord8 3
600 get =
601 do word <- getWord8
602 case word of
603 0 -> liftM2 Index get get
604 1 -> liftM2 Field get get
605 2 -> pure Empty
606 3 -> pure Alias
607 _ -> error "problem getting DecisionTree.Path binary"