git.iko.soy/elk
on 144a15a help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
445 lines 11 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 Data.Function (on)
16import qualified Data.List as List
17import qualified Data.Map as Map
18import qualified Data.Set as Set
19import qualified Data.Maybe as Maybe
21import qualified AST.Helpers as Help
22import qualified AST.Literal as L
23import qualified AST.Pattern as P
24import qualified AST.Variable as Var
25import qualified Reporting.Annotation as A
28type CPattern = P.CanonicalPattern
31-- COMPILE CASES
33{-| Users of this module will mainly interact with this function. It takes
34some normal branches and gives out a decision tree that has "labels" at all
35the leafs and a dictionary that maps these "labels" to the code that should
36run.
38If 2 or more leaves point to the same label, we need to do some tricks in JS to
39make that work nicely. When is JS getting goto?! ;) That is outside the scope
40of this module though.
41-}
42compile :: VariantDict -> [(CPattern, Int)] -> DecisionTree Jump
43compile variantDict rawBranches =
44 let
45 format (pattern, index) =
46 Branch index [(Empty, pattern)]
47 in
48 toDecisionTree variantDict (map format rawBranches)
51{-| When a certain union type is defined, you specify a certain number of tags.
52This helps us do a few optimizations:
54 * If there is only one possible tag, we can always skip checking what it is.
55 Tuples are a common example of this.
56 * If we use all possible tags, we can skip doing the last test. If we have
57 checked N-1 of N tags, there is no need to test the Nth, we know its the
58 one we want
60So this dictionary maps tags to the number of variants that exist, so it'll
61contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
62we can use for these optimizations.
63-}
64type VariantDict =
65 Map.Map Var.Home (Map.Map String Int)
68-- DECISION TREES
70data DecisionTree a
71 = Match a
72 | Decision
73 { _test :: Path
74 , _edges :: [(Test, DecisionTree a)]
75 , _default :: Maybe (DecisionTree a)
76 }
79data Jump = Jump
80 { _result :: Int
81 , _substitutions :: [(String, Path)]
82 }
85data Test
86 = Constructor Var.Canonical
87 | Literal L.Literal
88 deriving (Eq, Ord)
91data Path
92 = Position Int Path
93 | Field String Path
94 | Empty
95 | Alias
96 deriving (Eq)
99-- PATH HELPERS
101add :: Path -> Path -> Path
102add path finalLink =
103 case path of
104 Empty ->
105 finalLink
107 Alias ->
108 error "nothing should be added to an alias path"
110 Position index subpath ->
111 Position index (add subpath finalLink)
113 Field name subpath ->
114 Field name (add subpath finalLink)
117subPositions :: Path -> [CPattern] -> [(Path, CPattern)]
118subPositions path patterns =
119 zipWith
120 (\index pattern -> (add path (Position index Empty), pattern))
121 [0..]
122 patterns
125-- ACTUALLY BUILD DECISION TREES
127data Branch =
128 Branch
129 { _goal :: Int
130 , _patterns :: [(Path, CPattern)]
131 }
134toDecisionTree :: VariantDict -> [Branch] -> DecisionTree Jump
135toDecisionTree variantDict rawBranches =
136 let
137 branches =
138 map (flattenPatterns variantDict) rawBranches
139 in
140 case checkForMatch branches of
141 Just (goal, substitutions) ->
142 Match (Jump goal substitutions)
144 Nothing ->
145 let
146 path =
147 smallDefaults branches
149 relevantTests =
150 testsAtPath path branches
152 allRawEdges =
153 map (edgesFor path branches) relevantTests
155 (rawEdges, unexplored) =
156 if isComplete variantDict relevantTests then
157 ( init allRawEdges
158 , snd (last allRawEdges)
159 )
161 else
162 ( allRawEdges
163 , filter (isIrrelevantTo path) branches
164 )
166 edges =
167 map (second (toDecisionTree variantDict)) rawEdges
168 in
169 case (edges, unexplored) of
170 ([(_tag, decisionTree)], []) ->
171 decisionTree
173 (_, []) ->
174 Decision path edges Nothing
176 (_, _) ->
177 Decision path edges (Just (toDecisionTree variantDict unexplored))
180isComplete :: VariantDict -> [Test] -> Bool
181isComplete variantDict tests =
182 case head tests of
183 Constructor var ->
184 getArity variantDict var == length tests
186 Literal (L.Boolean _) ->
187 length tests == 2
189 _ ->
190 False
193getArity :: VariantDict -> Var.Canonical -> Int
194getArity variantDict (Var.Canonical home name) =
195 case Map.lookup name =<< Map.lookup home variantDict of
196 Just arity ->
197 arity
199 Nothing ->
200 if Help.isTuple name then
201 read (drop 6 name)
203 else
204 error
205 "Since the Optimize phase happens after canonicalization and type \
206 \inference, it is impossible that a pattern cannot be found."
209-- FLATTEN PATTERNS
211{-| Flatten type aliases and use the VariantDict to figure out when a tag is
212the only variant so we can skip doing any tests on it.
213-}
214flattenPatterns :: VariantDict -> Branch -> Branch
215flattenPatterns variantDict (Branch goal pathPatterns) =
216 Branch goal (concatMap (flatten variantDict) pathPatterns)
219flatten :: VariantDict -> (Path, CPattern) -> [(Path, CPattern)]
220flatten variantDict pathPattern@(path, A.A ann pattern) =
221 case pattern of
222 P.Var _ ->
223 [pathPattern]
225 P.Anything ->
226 [pathPattern]
228 P.Alias alias realPattern ->
229 [ (add path Alias, A.A ann (P.Var alias))
230 , (path, realPattern)
231 ]
233 P.Record _ ->
234 [pathPattern]
236 P.Data tag patterns ->
237 if getArity variantDict tag == 1 then
238 concatMap (flatten variantDict) (subPositions path patterns)
240 else
241 [pathPattern]
243 P.Literal _ ->
244 [pathPattern]
247-- SUCCESSFULLY MATCH
249{-| If the first branch has no more "decision points" we can finally take that
250path. If that is the case we give the resulting label and a mapping from free
251variables to "how to get their value". So a pattern like (Just (x,_)) will give
252us something like ("x" => value.0.0)
253-}
254checkForMatch :: [Branch] -> Maybe (Int, [(String, Path)])
255checkForMatch branches =
256 case branches of
257 Branch goal patterns : _ ->
258 (,) goal `fmap` concat `fmap` mapM getSubstitution patterns
260 _ ->
261 Nothing
264getSubstitution :: (Path, CPattern) -> Maybe [(String, Path)]
265getSubstitution (path, A.A _ pattern) =
266 case pattern of
267 P.Var x ->
268 Just [(x, path)]
270 P.Anything ->
271 Just []
273 P.Alias _ _ ->
274 error "aliases should never reach 'getSubstitution' function"
276 P.Record fields ->
277 Just (map (\name -> (name, add path (Field name Empty))) fields)
279 P.Data _ _ ->
280 Nothing
282 P.Literal _ ->
283 Nothing
286-- FIND RELEVANT TESTS
288testsAtPath :: Path -> [Branch] -> [Test]
289testsAtPath selectedPath branches =
290 let
291 allTests =
292 Maybe.mapMaybe (testAtPath selectedPath) branches
294 skipVisited test curr@(uniqueTests, visitedTests) =
295 if Set.member test visitedTests then
296 curr
297 else
298 ( test : uniqueTests
299 , Set.insert test visitedTests
300 )
301 in
302 fst (foldr skipVisited ([], Set.empty) allTests)
305testAtPath :: Path -> Branch -> Maybe Test
306testAtPath selectedPath (Branch _ pathPatterns) =
307 case List.lookup selectedPath pathPatterns of
308 Nothing ->
309 Nothing
311 Just (A.A _ pattern) ->
312 case pattern of
313 P.Data name _ ->
314 Just (Constructor name)
316 P.Literal lit ->
317 Just (Literal lit)
319 P.Var _ ->
320 Nothing
322 P.Alias _ _ ->
323 error "aliases should never reach 'testAtPath' function"
325 P.Anything ->
326 Nothing
328 P.Record _ ->
329 Nothing
332-- BUILD EDGES
334edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
335edgesFor path branches test =
336 ( test
337 , Maybe.mapMaybe (toRelevantBranch test path) branches
338 )
341toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
342toRelevantBranch test path branch@(Branch goal pathPatterns) =
343 case extract path pathPatterns of
344 Just (start, A.A _ pattern, end) ->
345 case pattern of
346 P.Data name patterns ->
347 if test == Constructor name then
348 Just (Branch goal (start ++ subPositions path patterns ++ end))
350 else
351 Nothing
353 P.Literal lit ->
354 if test == Literal lit then
355 Just (Branch goal (start ++ end))
357 else
358 Nothing
360 _ ->
361 Just branch
363 _ ->
364 Just branch
367extract
368 :: Path
369 -> [(Path, CPattern)]
370 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
371extract selectedPath pathPatterns =
372 case pathPatterns of
373 [] ->
374 Nothing
376 first@(path, pattern) : rest ->
377 if path == selectedPath then
378 Just ([], pattern, rest)
380 else
381 case extract selectedPath rest of
382 Nothing ->
383 Nothing
385 Just (start, foundPattern, end) ->
386 Just (first : start, foundPattern, end)
389-- FIND IRRELEVANT BRANCHES
391isIrrelevantTo :: Path -> Branch -> Bool
392isIrrelevantTo selectedPath (Branch _ pathPatterns) =
393 case List.lookup selectedPath pathPatterns of
394 Nothing ->
395 True
397 Just pattern ->
398 not (needsTests pattern)
401needsTests :: CPattern -> Bool
402needsTests (A.A _ pattern) =
403 case pattern of
404 P.Var _ ->
405 False
407 P.Anything ->
408 False
410 P.Alias _ _ ->
411 error "aliases should never reach 'isIrrelevantTo' function"
413 P.Record _ ->
414 False
416 P.Data _ _ ->
417 True
419 P.Literal _ ->
420 True
423-- PATH PICKING HEURISTICS
425smallDefaults :: [Branch] -> Path
426smallDefaults branches =
427 let
428 allPaths =
429 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
431 weightedPaths =
432 map (\path -> (path, length (filter (isIrrelevantTo path) branches))) allPaths
433 in
434 fst (List.minimumBy (compare `on` snd) weightedPaths)
437isChoicePath :: (Path, CPattern) -> Maybe Path
438isChoicePath (path, pattern) =
439 if needsTests pattern then
440 Just path
441 else
442 Nothing
445-- smallBranchingFactor :: [Branch] -> Path