git.iko.soy/elk
on d3c72d1 help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
488 lines 12 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Patterns.DecisionTree where
3{- To learn more about how this works, definitely read through:
5 "When Do Match-Compilation Heuristics Matter?"
7by Kevin Scott and Norman Ramsey. The rough idea is that we start with a simple
8list of patterns and expressions, and then turn that into a "decision tree"
9that requires as few tests as possible to make it to a leaf. Read the paper, it
10explains this extraordinarily well! We are currently using the same heuristics
11as SML/NJ to get nice trees.
12-}
14import Control.Arrow (second)
15import qualified Data.List as List
16import qualified Data.Map as Map
17import qualified Data.Set as Set
18import qualified Data.Maybe as Maybe
20import qualified AST.Helpers as Help
21import qualified AST.Literal as L
22import qualified AST.Pattern as P
23import qualified AST.Variable as Var
24import qualified Reporting.Annotation as A
27type CPattern = P.CanonicalPattern
30-- COMPILE CASES
32{-| Users of this module will mainly interact with this function. It takes
33some normal branches and gives out a decision tree that has "labels" at all
34the leafs and a dictionary that maps these "labels" to the code that should
35run.
37If 2 or more leaves point to the same label, we need to do some tricks in JS to
38make that work nicely. When is JS getting goto?! ;) That is outside the scope
39of this module though.
40-}
41compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree Jump
42compile variantDict rawBranches =
43 let
44 format (pattern, index) =
45 Branch index [(Empty, pattern)]
46 in
47 toDecisionTree variantDict (map format rawBranches)
50{-| When a certain union type is defined, you specify a certain number of tags.
51This helps us do a few optimizations:
53 * If there is only one possible tag, we can always skip checking what it is.
54 Tuples are a common example of this.
55 * If we use all possible tags, we can skip doing the last test. If we have
56 checked N-1 of N tags, there is no need to test the Nth, we know its the
57 one we want
59So this dictionary maps tags to the number of variants that exist, so it'll
60contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
61we can use for these optimizations.
62-}
63type VariantDict =
64 Map.Map Var.Home (Map.Map String Int)
67-- DECISION TREES
69data DecisionTree a
70 = Match a
71 | Decision
72 { _test :: Path
73 , _edges :: [(Test, DecisionTree a)]
74 , _default :: Maybe (DecisionTree a)
75 }
78data Jump = Jump
79 { _result :: Int
80 , _substitutions :: [(String, Path)]
81 }
84data Test
85 = Constructor Var.Canonical
86 | Literal L.Literal
87 deriving (Eq, Ord)
90data Path
91 = Position Int Path
92 | Field String Path
93 | Empty
94 | Alias
95 deriving (Eq)
98-- PATH HELPERS
100add :: Path -> Path -> Path
101add path finalLink =
102 case path of
103 Empty ->
104 finalLink
106 Alias ->
107 error "nothing should be added to an alias path"
109 Position index subpath ->
110 Position index (add subpath finalLink)
112 Field name subpath ->
113 Field name (add subpath finalLink)
116subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
117subPositions path patterns =
118 zipWith
119 (\index pattern -> (add path (Position index Empty), pattern))
120 [0..]
121 patterns
124-- ACTUALLY BUILD DECISION TREES
126data Branch =
127 Branch
128 { _goal :: Int
129 , _patterns :: [(Path, CPattern)]
130 }
133toDecisionTree :: VariantDict -> [Branch] -> DecisionTree Jump
134toDecisionTree variantDict rawBranches =
135 let
136 branches =
137 map (flattenPatterns variantDict) rawBranches
138 in
139 case checkForMatch branches of
140 Just (goal, substitutions) ->
141 Match (Jump goal substitutions)
143 Nothing ->
144 let
145 path =
146 pickPath variantDict branches
148 (edges, fallback) =
149 gatherEdges variantDict branches path
151 decisionEdges =
152 map (second (toDecisionTree variantDict)) edges
153 in
154 case (decisionEdges, fallback) of
155 ([(_tag, decisionTree)], []) ->
156 decisionTree
158 (_, []) ->
159 Decision path decisionEdges Nothing
161 (_, _) ->
162 Decision path decisionEdges (Just (toDecisionTree variantDict fallback))
165isComplete :: VariantDict -> [Test] -> Bool
166isComplete variantDict tests =
167 case head tests of
168 Constructor var ->
169 getArity variantDict var == length tests
171 Literal (L.Boolean _) ->
172 length tests == 2
174 _ ->
175 False
178getArity :: VariantDict -> Var.Canonical -> Int
179getArity variantDict (Var.Canonical home name) =
180 case Map.lookup name =<< Map.lookup home variantDict of
181 Just arity ->
182 arity
184 Nothing ->
185 if Help.isTuple name then
186 read (drop 6 name)
188 else
189 error
190 "Since the Optimize phase happens after canonicalization and type \
191 \inference, it is impossible that a pattern cannot be found."
194-- FLATTEN PATTERNS
196{-| Flatten type aliases and use the VariantDict to figure out when a tag is
197the only variant so we can skip doing any tests on it.
198-}
199flattenPatterns :: VariantDict -> Branch -> Branch
200flattenPatterns variantDict (Branch goal pathPatterns) =
201 Branch goal (concatMap (flatten variantDict) pathPatterns)
204flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
205flatten variantDict pathPattern@(path, A.A ann pattern) =
206 case pattern of
207 P.Var _ ->
208 [pathPattern]
210 P.Anything ->
211 [pathPattern]
213 P.Alias alias realPattern ->
214 [ (add path Alias, A.A ann (P.Var alias))
215 , (path, realPattern)
216 ]
218 P.Record _ ->
219 [pathPattern]
221 P.Data tag patterns ->
222 if getArity variantDict tag == 1 then
223 concatMap (flatten variantDict) (subPositions path patterns)
225 else
226 [pathPattern]
228 P.Literal _ ->
229 [pathPattern]
232-- SUCCESSFULLY MATCH
234{-| If the first branch has no more "decision points" we can finally take that
235path. If that is the case we give the resulting label and a mapping from free
236variables to "how to get their value". So a pattern like (Just (x,_)) will give
237us something like ("x" => value.0.0)
238-}
239checkForMatch :: [Branch] -> Maybe (Int, [(String, Path)])
240checkForMatch branches =
241 case branches of
242 Branch goal patterns : _ ->
243 (,) goal `fmap` concat `fmap` mapM getSubstitution patterns
245 _ ->
246 Nothing
249getSubstitution :: (Path, CPattern) -> Maybe [(String, Path)]
250getSubstitution (path, A.A _ pattern) =
251 case pattern of
252 P.Var x ->
253 Just [(x, path)]
255 P.Anything ->
256 Just []
258 P.Alias _ _ ->
259 error "aliases should never reach 'getSubstitution' function"
261 P.Record fields ->
262 Just (map (\name -> (name, add path (Field name Empty))) fields)
264 P.Data _ _ ->
265 Nothing
267 P.Literal _ ->
268 Nothing
271-- GATHER OUTGOING EDGES
273gatherEdges :: VariantDict -> [Branch] -> Path -> ([(Test, [Branch])], [Branch])
274gatherEdges variantDict branches path =
275 let
276 relevantTests =
277 testsAtPath path branches
279 allRawEdges =
280 map (edgesFor path branches) relevantTests
281 in
282 if isComplete variantDict relevantTests then
283 ( init allRawEdges
284 , snd (last allRawEdges)
285 )
287 else
288 ( allRawEdges
289 , filter (isIrrelevantTo path) branches
290 )
293-- FIND RELEVANT TESTS
295testsAtPath :: Path -> [Branch] -> [Test]
296testsAtPath selectedPath branches =
297 let
298 allTests =
299 Maybe.mapMaybe (testAtPath selectedPath) branches
301 skipVisited test curr@(uniqueTests, visitedTests) =
302 if Set.member test visitedTests then
303 curr
304 else
305 ( test : uniqueTests
306 , Set.insert test visitedTests
307 )
308 in
309 fst (foldr skipVisited ([], Set.empty) allTests)
312testAtPath :: Path -> Branch -> Maybe Test
313testAtPath selectedPath (Branch _ pathPatterns) =
314 case List.lookup selectedPath pathPatterns of
315 Nothing ->
316 Nothing
318 Just (A.A _ pattern) ->
319 case pattern of
320 P.Data name _ ->
321 Just (Constructor name)
323 P.Literal lit ->
324 Just (Literal lit)
326 P.Var _ ->
327 Nothing
329 P.Alias _ _ ->
330 error "aliases should never reach 'testAtPath' function"
332 P.Anything ->
333 Nothing
335 P.Record _ ->
336 Nothing
339-- BUILD EDGES
341edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
342edgesFor path branches test =
343 ( test
344 , Maybe.mapMaybe (toRelevantBranch test path) branches
345 )
348toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
349toRelevantBranch test path branch@(Branch goal pathPatterns) =
350 case extract path pathPatterns of
351 Just (start, A.A _ pattern, end) ->
352 case pattern of
353 P.Data name patterns ->
354 if test == Constructor name then
355 Just (Branch goal (start ++ subPositions path patterns ++ end))
357 else
358 Nothing
360 P.Literal lit ->
361 if test == Literal lit then
362 Just (Branch goal (start ++ end))
364 else
365 Nothing
367 _ ->
368 Just branch
370 _ ->
371 Just branch
374extract
375 :: Path
376 -> [(Path, CPattern)]
377 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
378extract selectedPath pathPatterns =
379 case pathPatterns of
380 [] ->
381 Nothing
383 first@(path, pattern) : rest ->
384 if path == selectedPath then
385 Just ([], pattern, rest)
387 else
388 case extract selectedPath rest of
389 Nothing ->
390 Nothing
392 Just (start, foundPattern, end) ->
393 Just (first : start, foundPattern, end)
396-- FIND IRRELEVANT BRANCHES
398isIrrelevantTo :: Path -> Branch -> Bool
399isIrrelevantTo selectedPath (Branch _ pathPatterns) =
400 case List.lookup selectedPath pathPatterns of
401 Nothing ->
402 True
404 Just pattern ->
405 not (needsTests pattern)
408needsTests :: CPattern -> Bool
409needsTests (A.A _ pattern) =
410 case pattern of
411 P.Var _ ->
412 False
414 P.Anything ->
415 False
417 P.Alias _ _ ->
418 error "aliases should never reach 'isIrrelevantTo' function"
420 P.Record _ ->
421 False
423 P.Data _ _ ->
424 True
426 P.Literal _ ->
427 True
430-- PICK A PATH
432pickPath :: VariantDict -> [Branch] -> Path
433pickPath variantDict branches =
434 let
435 allPaths =
436 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
437 in
438 case bests (addWeights (smallDefaults branches) allPaths) of
439 [path] ->
440 path
442 tiedPaths ->
443 head (bests (addWeights (smallBranchingFactor variantDict branches) tiedPaths))
446isChoicePath :: (Path, CPattern) -> Maybe Path
447isChoicePath (path, pattern) =
448 if needsTests pattern then
449 Just path
450 else
451 Nothing
454addWeights :: (Path -> Int) -> [Path] -> [(Path, Int)]
455addWeights toWeight paths =
456 map (\path -> (path, toWeight path)) paths
459bests :: [(Path, Int)] -> [Path]
460bests ((path,weight):weightedPaths) =
461 let
462 gatherMinimum acc@(minWeight, paths) (path, weight) =
463 if weight == minWeight then
464 (minWeight, path : paths)
466 else if weight < minWeight then
467 (weight, [path])
469 else
470 acc
471 in
472 snd (List.foldl' gatherMinimum (weight, [path]) weightedPaths)
475-- PATH PICKING HEURISTICS
477smallDefaults :: [Branch] -> Path -> Int
478smallDefaults branches path =
479 length (filter (isIrrelevantTo path) branches)
482smallBranchingFactor :: VariantDict -> [Branch] -> Path -> Int
483smallBranchingFactor variantDict branches path =
484 let
485 (edges, fallback) =
486 gatherEdges variantDict branches path
487 in
488 length edges + (if null fallback then 0 else 1)