git.iko.soy/elk
on 2297bbf help sign in
elk/src/Optimize/Patterns/DecisionTree.hs
430 lines 10 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 Data.Map ((!))
20import qualified Data.Maybe as Maybe
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.Canonical 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 name ->
184 (variantDict ! name) == length tests
186 Literal (L.Boolean _) ->
187 length tests == 2
189 _ ->
190 False
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 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-- FIND RELEVANT TESTS
273testsAtPath :: Path -> [Branch] -> [Test]
274testsAtPath selectedPath branches =
275 let
276 allTests =
277 Maybe.mapMaybe (testAtPath selectedPath) branches
279 skipVisited test curr@(uniqueTests, visitedTests) =
280 if Set.member test visitedTests then
281 curr
282 else
283 ( test : uniqueTests
284 , Set.insert test visitedTests
285 )
286 in
287 fst (foldr skipVisited ([], Set.empty) allTests)
290testAtPath :: Path -> Branch -> Maybe Test
291testAtPath selectedPath (Branch _ pathPatterns) =
292 case List.lookup selectedPath pathPatterns of
293 Nothing ->
294 Nothing
296 Just (A.A _ pattern) ->
297 case pattern of
298 P.Data name _ ->
299 Just (Constructor name)
301 P.Literal lit ->
302 Just (Literal lit)
304 P.Var _ ->
305 Nothing
307 P.Alias _ _ ->
308 error "aliases should never reach 'testAtPath' function"
310 P.Anything ->
311 Nothing
313 P.Record _ ->
314 Nothing
317-- BUILD EDGES
319edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
320edgesFor path branches test =
321 ( test
322 , Maybe.mapMaybe (toRelevantBranch test path) branches
323 )
326toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
327toRelevantBranch test path branch@(Branch goal pathPatterns) =
328 case extract path pathPatterns of
329 Just (start, A.A _ pattern, end) ->
330 case pattern of
331 P.Data name patterns ->
332 if test == Constructor name then
333 Just (Branch goal (start ++ subPositions path patterns ++ end))
335 else
336 Nothing
338 P.Literal lit ->
339 if test == Literal lit then
340 Just (Branch goal (start ++ end))
342 else
343 Nothing
345 _ ->
346 Just branch
348 _ ->
349 Just branch
352extract
353 :: Path
354 -> [(Path, CPattern)]
355 -> Maybe ([(Path, CPattern)], CPattern, [(Path, CPattern)])
356extract selectedPath pathPatterns =
357 case pathPatterns of
358 [] ->
359 Nothing
361 first@(path, pattern) : rest ->
362 if path == selectedPath then
363 Just ([], pattern, rest)
365 else
366 case extract selectedPath rest of
367 Nothing ->
368 Nothing
370 Just (start, foundPattern, end) ->
371 Just (first : start, foundPattern, end)
374-- FIND IRRELEVANT BRANCHES
376isIrrelevantTo :: Path -> Branch -> Bool
377isIrrelevantTo selectedPath (Branch _ pathPatterns) =
378 case List.lookup selectedPath pathPatterns of
379 Nothing ->
380 True
382 Just pattern ->
383 not (needsTests pattern)
386needsTests :: CPattern -> Bool
387needsTests (A.A _ pattern) =
388 case pattern of
389 P.Var _ ->
390 False
392 P.Anything ->
393 False
395 P.Alias _ _ ->
396 error "aliases should never reach 'isIrrelevantTo' function"
398 P.Record _ ->
399 False
401 P.Data _ _ ->
402 True
404 P.Literal _ ->
405 True
408-- PATH PICKING HEURISTICS
410smallDefaults :: [Branch] -> Path
411smallDefaults branches =
412 let
413 allPaths =
414 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
416 weightedPaths =
417 map (\path -> (path, length (filter (isIrrelevantTo path) branches))) allPaths
418 in
419 fst (List.minimumBy (compare `on` snd) weightedPaths)
422isChoicePath :: (Path, CPattern) -> Maybe Path
423isChoicePath (path, pattern) =
424 if needsTests pattern then
425 Just path
426 else
427 Nothing
430-- smallBranchingFactor :: [Branch] -> Path