git.iko.soy/elk
on 22bb2ee help sign in
elk/src/Optimize/DecisionTree.hs
605 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, 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 | IsInt Int
78 | IsChr Text
79 | IsStr Text
80 deriving (Eq, Ord)
83data Path
84 = Index Index.ZeroBased Path
85 | Unbox Path
86 | Empty
87 deriving (Eq)
91-- ACTUALLY BUILD DECISION TREES
94data Branch =
95 Branch
96 { _goal :: Int
97 , _patterns :: [(Path, Can.Pattern)]
98 }
101toDecisionTree :: [Branch] -> DecisionTree
102toDecisionTree rawBranches =
103 let
104 branches =
105 map flattenPatterns rawBranches
106 in
107 case checkForMatch branches of
108 Just goal ->
109 Match goal
111 Nothing ->
112 let
113 path =
114 pickPath branches
116 (edges, fallback) =
117 gatherEdges branches path
119 decisionEdges =
120 map (second toDecisionTree) edges
121 in
122 case (decisionEdges, fallback) of
123 ([(_tag, decisionTree)], []) ->
124 decisionTree
126 (_, []) ->
127 Decision path decisionEdges Nothing
129 ([], _ : _) ->
130 toDecisionTree fallback
132 (_, _) ->
133 Decision path decisionEdges (Just (toDecisionTree fallback))
136isComplete :: [Test] -> Bool
137isComplete tests =
138 case head tests of
139 IsCtor _ _ numAlts _ ->
140 numAlts == length tests
142 IsCons ->
143 length tests == 2
145 IsNil ->
146 length tests == 2
148 _ ->
149 False
153-- FLATTEN PATTERNS
156{-| Flatten type aliases and use the VariantDict to figure out when a tag is
157the only variant so we can skip doing any tests on it.
158-}
159flattenPatterns :: Branch -> Branch
160flattenPatterns (Branch goal pathPatterns) =
161 Branch goal (concatMap flatten pathPatterns)
164flatten :: (Path, Can.Pattern) -> [(Path, Can.Pattern)]
165flatten pathPattern@(path, A.At region pattern) =
166 case pattern of
167 Can.PVar _ ->
168 [pathPattern]
170 Can.PAnything ->
171 [pathPattern]
173 Can.PTuple a b Nothing ->
174 concatMap flatten (subPositions path [a,b])
176 Can.PTuple a b (Just c) ->
177 concatMap flatten (subPositions path [a,b,c])
179 Can.PCtor _ _ (Can.Union _ _ 1 _) _ _ ctorArgs ->
180 case map dearg ctorArgs of
181 [arg] ->
182 flatten (Unbox path, arg)
184 args ->
185 concatMap flatten (subPositions path args)
187 Can.PCtor _ _ _ _ _ _ ->
188 [pathPattern]
190 Can.PAlias realPattern alias ->
191 (path, A.At region (Can.PVar alias)) : flatten (path, realPattern)
193 Can.PRecord _ ->
194 [pathPattern]
196 Can.PUnit ->
197 [pathPattern]
199 Can.PList _ ->
200 [pathPattern]
202 Can.PCons _ _ ->
203 [pathPattern]
205 Can.PChr _ ->
206 [pathPattern]
208 Can.PStr _ ->
209 [pathPattern]
211 Can.PInt _ ->
212 [pathPattern]
215subPositions :: Path -> [Can.Pattern] -> [(Path, Can.Pattern)]
216subPositions path patterns =
217 Index.indexedMap (\index pattern -> (Index index path, pattern)) patterns
220dearg :: Can.PatternCtorArg -> Can.Pattern
221dearg (Can.PatternCtorArg _ _ pattern) =
222 pattern
226-- SUCCESSFULLY MATCH
229{-| If the first branch has no more "decision points" we can finally take that
230path. If that is the case we give the resulting label and a mapping from free
231variables to "how to get their value". So a pattern like (Just (x,_)) will give
232us something like ("x" => value.0.0)
233-}
234checkForMatch :: [Branch] -> Maybe Int
235checkForMatch branches =
236 case branches of
237 Branch goal patterns : _ | all (not . needsTests . snd) patterns ->
238 Just goal
240 _ ->
241 Nothing
245-- GATHER OUTGOING EDGES
248gatherEdges :: [Branch] -> Path -> ([(Test, [Branch])], [Branch])
249gatherEdges branches path =
250 let
251 relevantTests =
252 testsAtPath path branches
254 allEdges =
255 map (edgesFor path branches) relevantTests
257 fallbacks =
258 if isComplete relevantTests then
259 []
260 else
261 filter (isIrrelevantTo path) branches
262 in
263 ( allEdges, fallbacks )
267-- FIND RELEVANT TESTS
270testsAtPath :: Path -> [Branch] -> [Test]
271testsAtPath selectedPath branches =
272 let
273 allTests =
274 Maybe.mapMaybe (testAtPath selectedPath) branches
276 skipVisited test curr@(uniqueTests, visitedTests) =
277 if Set.member test visitedTests then
278 curr
279 else
280 ( test : uniqueTests
281 , Set.insert test visitedTests
282 )
283 in
284 fst (foldr skipVisited ([], Set.empty) allTests)
287testAtPath :: Path -> Branch -> Maybe Test
288testAtPath selectedPath (Branch _ pathPatterns) =
289 case List.lookup selectedPath pathPatterns of
290 Nothing ->
291 Nothing
293 Just (A.At _ pattern) ->
294 case pattern of
295 Can.PCtor _ _ (Can.Union _ _ numAlts opts) name index _ ->
296 Just (IsCtor name index numAlts opts)
298 Can.PInt int ->
299 Just (IsInt int)
301 Can.PStr str ->
302 Just (IsStr str)
304 Can.PChr chr ->
305 Just (IsChr chr)
307 Can.PList [] ->
308 Just IsNil
310 Can.PList _ ->
311 Just IsCons
313 Can.PCons _ _ ->
314 Just IsCons
316 Can.PVar _ ->
317 Nothing
319 Can.PAlias _ _ ->
320 error "aliases should never reach 'testAtPath' function"
322 Can.PAnything ->
323 Nothing
325 Can.PUnit ->
326 Nothing
328 Can.PTuple _ _ _ ->
329 Nothing
331 Can.PRecord _ ->
332 Nothing
336-- BUILD EDGES
339edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
340edgesFor path branches test =
341 ( test
342 , Maybe.mapMaybe (toRelevantBranch test path) branches
343 )
346toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
347toRelevantBranch test path branch@(Branch goal pathPatterns) =
348 case extract path pathPatterns of
349 Just (start, A.At region pattern, end) ->
350 case pattern of
351 Can.PCtor _ _ (Can.Union _ _ numAlts _) name _ ctorArgs ->
352 case test of
353 IsCtor testName _ _ _ | name == testName ->
354 Just $ Branch goal $
355 case map dearg ctorArgs of
356 [arg] | numAlts == 1 ->
357 start ++ [(Unbox path, arg)] ++ end
359 args ->
360 start ++ subPositions path args ++ end
362 _ ->
363 Nothing
365 Can.PList [] ->
366 case test of
367 IsNil ->
368 Just (Branch goal (start ++ end))
370 _ ->
371 Nothing
373 Can.PList (hd:tl) ->
374 case test of
375 IsCons ->
376 let tl' = A.At region (Can.PList tl) in
377 Just (Branch goal (start ++ subPositions path [ hd, tl' ] ++ end))
379 _ ->
380 Nothing
382 Can.PCons hd tl ->
383 case test of
384 IsCons ->
385 Just (Branch goal (start ++ subPositions path [hd,tl] ++ end))
387 _ ->
388 Nothing
390 Can.PChr chr ->
391 case test of
392 IsChr testChr | chr == testChr ->
393 Just (Branch goal (start ++ end))
394 _ ->
395 Nothing
397 Can.PStr str ->
398 case test of
399 IsStr testStr | str == testStr ->
400 Just (Branch goal (start ++ end))
402 _ ->
403 Nothing
405 Can.PInt int ->
406 case test of
407 IsInt testInt | int == testInt ->
408 Just (Branch goal (start ++ end))
410 _ ->
411 Nothing
413 _ ->
414 Just branch
416 _ ->
417 Just branch
420extract
421 :: Path
422 -> [(Path, Can.Pattern)]
423 -> Maybe ([(Path, Can.Pattern)], Can.Pattern, [(Path, Can.Pattern)])
424extract selectedPath pathPatterns =
425 case pathPatterns of
426 [] ->
427 Nothing
429 first@(path, pattern) : rest ->
430 if path == selectedPath then
431 Just ([], pattern, rest)
433 else
434 case extract selectedPath rest of
435 Nothing ->
436 Nothing
438 Just (start, foundPattern, end) ->
439 Just (first : start, foundPattern, end)
443-- FIND IRRELEVANT BRANCHES
446isIrrelevantTo :: Path -> Branch -> Bool
447isIrrelevantTo selectedPath (Branch _ pathPatterns) =
448 case List.lookup selectedPath pathPatterns of
449 Nothing ->
450 True
452 Just pattern ->
453 not (needsTests pattern)
456needsTests :: Can.Pattern -> Bool
457needsTests (A.At _ pattern) =
458 case pattern of
459 Can.PVar _ ->
460 False
462 Can.PAnything ->
463 False
465 Can.PAlias _ _ ->
466 error "aliases should never reach 'isIrrelevantTo' function"
468 Can.PRecord _ ->
469 False
471 Can.PCtor _ _ _ _ _ _ ->
472 True
474 Can.PList _ ->
475 True
477 Can.PCons _ _ ->
478 True
480 Can.PUnit ->
481 False
483 Can.PTuple _ _ _ ->
484 False
486 Can.PChr _ ->
487 True
489 Can.PStr _ ->
490 True
492 Can.PInt _ ->
493 True
497-- PICK A PATH
500pickPath :: [Branch] -> Path
501pickPath branches =
502 let
503 allPaths =
504 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
505 in
506 case bests (addWeights (smallDefaults branches) allPaths) of
507 [path] ->
508 path
510 tiedPaths ->
511 head (bests (addWeights (smallBranchingFactor branches) tiedPaths))
514isChoicePath :: (Path, Can.Pattern) -> Maybe Path
515isChoicePath (path, pattern) =
516 if needsTests pattern then
517 Just path
518 else
519 Nothing
522addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
523addWeights toWeight paths =
524 map (\path -> (path, toWeight path)) paths
527bests :: [(Path, Int)] -> [Path]
528bests allPaths =
529 case allPaths of
530 [] ->
531 error "Cannot choose the best of zero paths. This should never happen."
533 (headPath, headWeight) : weightedPaths ->
534 let
535 gatherMinimum acc@(minWeight, paths) (path, weight) =
536 if weight == minWeight then
537 (minWeight, path : paths)
539 else if weight < minWeight then
540 (weight, [path])
542 else
543 acc
544 in
545 snd (List.foldl' gatherMinimum (headWeight, [headPath]) weightedPaths)
549-- PATH PICKING HEURISTICS
552smallDefaults :: [Branch] -> Path -> Int
553smallDefaults branches path =
554 length (filter (isIrrelevantTo path) branches)
557smallBranchingFactor :: [Branch] -> Path -> Int
558smallBranchingFactor branches path =
559 let
560 (edges, fallback) =
561 gatherEdges branches path
562 in
563 length edges + (if null fallback then 0 else 1)
567-- BINARY
570instance Binary Test where
571 put test =
572 case test of
573 IsCtor a b c d -> putWord8 0 >> put a >> put b >> put c >> put d
574 IsCons -> putWord8 1
575 IsNil -> putWord8 2
576 IsChr a -> putWord8 3 >> put a
577 IsStr a -> putWord8 4 >> put a
578 IsInt a -> putWord8 5 >> put a
580 get =
581 do word <- getWord8
582 case word of
583 0 -> liftM4 IsCtor get get get get
584 1 -> pure IsCons
585 2 -> pure IsNil
586 3 -> liftM IsChr get
587 4 -> liftM IsStr get
588 5 -> 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 Unbox a -> putWord8 1 >> put a
597 Empty -> putWord8 2
599 get =
600 do word <- getWord8
601 case word of
602 0 -> liftM2 Index get get
603 1 -> liftM Unbox get
604 2 -> pure Empty
605 _ -> error "problem getting DecisionTree.Path binary"