git.iko.soy/elk
on 49032cd help sign in
elk/src/Optimize/Cases.hs
415 lines 10 kB history raw
1module Optimize.Cases where
2{- To learn more about how this works, definitely read through:
4 "When Do Match-Compilation Heuristics Matter?"
6by Kevin Scott and Norman Ramsey. The rough idea is that we start with a simple
7list of patterns and expressions, and then turn that into a "decision tree"
8that requires as few tests as possible to make it to a leaf. Read the paper, it
9explains this extraordinarily well! We are currently using the same heuristics
10as SML/NJ to get nice trees.
11-}
14import Control.Arrow (second)
15import Data.Function (on)
16import qualified Data.List as List
17import qualified Data.Map as Map
18import Data.Map ((!))
19import qualified Data.Maybe as Maybe
21import qualified AST.Expression.Optimized as Opt
22import qualified AST.Literal as L
23import qualified AST.Pattern as P
24import qualified AST.Variable as Var
25import qualified Reporting.Annotation as A
28-- COMPILE CASES
30{-| Users of this module will mainly interact with this function. It takes
31some normal branches and gives out a decision tree that has "labels" at all
32the leafs and a dictionary that maps these "labels" to the code that should
33run.
35If 2 or more leaves point to the same label, we need to do some tricks in JS to
36make that work nicely. When is JS getting goto?! ;) That is outside the scope
37of this module though.
38-}
39compile
40 :: VariantDict
41 -> [(P.Optimized, Opt.Expr)]
42 -> (DecisionTree Int, Map.Map Int Opt.Expr)
43compile variantDict rawBranches =
44 let
45 format index (pattern, expr) =
46 ( Branch index [(Empty, pattern)]
47 , (index, expr)
48 )
50 (branches, resultList) =
51 unzip (zipWith format [0..] rawBranches)
52 in
53 ( toDecisionTree variantDict branches
54 , Map.fromList resultList
55 )
58{-| When a certain union type is defined, you specify a certain number of tags.
59This helps us do a few optimizations:
61 * If there is only one possible tag, we can always skip checking what it is.
62 Tuples are a common example of this.
63 * If we use all possible tags, we can skip doing the last test. If we have
64 checked N-1 of N tags, there is no need to test the Nth, we know its the
65 one we want
67So this dictionary maps tags to the number of variants that exist, so it'll
68contain things like [ ("Just", 2), ("Nothing", 2), ("_Tuple2", 1), ... ] which
69we can use for these optimizations.
70-}
71type VariantDict =
72 Map.Map Var.Canonical Int
75-- DECISION TREES
77data DecisionTree a
78 = Decision
79 { _test :: Path
80 , _edges :: [(Test, DecisionTree a)]
81 , _default :: Maybe (DecisionTree a)
82 }
83 | Match
84 { result :: a
85 , substitutions :: Map.Map String Path
86 }
89data Test
90 = Constructor Var.Canonical
91 | Literal L.Literal
92 deriving (Eq, Ord)
95data Path
96 = Position Int Path
97 | Field String Path
98 | Empty
99 | Alias
100 deriving (Eq)
103-- PATH HELPERS
105add :: Path -> Path -> Path
106add path finalLink =
107 case path of
108 Empty ->
109 finalLink
111 Alias ->
112 error "nothing should be added to an alias path"
114 Position index subpath ->
115 Position index (add subpath finalLink)
117 Field name subpath ->
118 Field name (add subpath finalLink)
121subPositions :: Path -> [P.Optimized] -> [(Path, P.Optimized)]
122subPositions path patterns =
123 zipWith
124 (\index pattern -> (add path (Position index Empty), pattern))
125 [0..]
126 patterns
129-- ACTUALLY BUILD DECISION TREES
131data Branch =
132 Branch
133 { _goal :: Int
134 , _patterns :: [(Path, P.Optimized)]
135 }
138toDecisionTree :: VariantDict -> [Branch] -> DecisionTree Int
139toDecisionTree variantDict rawBranches =
140 let
141 branches =
142 map (flattenPatterns variantDict) rawBranches
143 in
144 case checkForMatch branches of
145 Just (goal, substitutions) ->
146 Match goal substitutions
148 Nothing ->
149 let
150 path =
151 smallDefaults branches
153 relevantTests =
154 tagsAtPath path branches
156 allRawEdges =
157 map (edgesFor path branches) relevantTests
159 (rawEdges, unexplored) =
160 if isComplete variantDict relevantTests then
161 ( init allRawEdges
162 , snd (last allRawEdges)
163 )
165 else
166 ( allRawEdges
167 , filter (isIrrelevantTo path) branches
168 )
170 edges =
171 map (second (toDecisionTree variantDict)) rawEdges
172 in
173 case (edges, unexplored) of
174 ([(_tag, decisionTree)], []) ->
175 decisionTree
177 (_, []) ->
178 Decision path edges Nothing
180 (_, _) ->
181 Decision path edges (Just (toDecisionTree variantDict unexplored))
184isComplete :: VariantDict -> [Test] -> Bool
185isComplete variantDict tests =
186 case head tests of
187 Constructor name ->
188 (variantDict ! name) == length tests
190 Literal (L.Boolean _) ->
191 length tests == 2
193 _ ->
194 False
198-- FLATTEN PATTERNS
200{-| Flatten type aliases and use the VariantDict to figure out when a tag is
201the only variant so we can skip doing any tests on it.
202-}
203flattenPatterns :: VariantDict -> Branch -> Branch
204flattenPatterns variantDict (Branch goal pathPatterns) =
205 Branch goal (concatMap (flatten variantDict) pathPatterns)
208flatten :: VariantDict -> (Path, P.Optimized) -> [(Path, P.Optimized)]
209flatten variantDict pathPattern@(path, A.A ann pattern) =
210 case pattern of
211 P.Var _ ->
212 [pathPattern]
214 P.Anything ->
215 [pathPattern]
217 P.Alias alias realPattern ->
218 [ (add path Alias, A.A ann (P.Var alias))
219 , (path, realPattern)
220 ]
222 P.Record _ ->
223 [pathPattern]
225 P.Data tag patterns ->
226 if variantDict ! tag == 1 then
227 concatMap (flatten variantDict) (subPositions path patterns)
229 else
230 [pathPattern]
232 P.Literal _ ->
233 [pathPattern]
236-- SUCCESSFULLY MATCH
238{-| If the first branch has no more "decision points" we can finally take that
239path. If that is the case we give the resulting label and a mapping from free
240variables to "how to get their value". So a pattern like (Just (x,_)) will give
241us something like ("x" => value.0.0)
242-}
243checkForMatch :: [Branch] -> Maybe (Int, Map.Map String Path)
244checkForMatch branches =
245 case branches of
246 Branch goal patterns : _ ->
247 do subs <- concat `fmap` mapM getSubstitution patterns
248 return (goal, Map.fromList subs)
250 _ ->
251 Nothing
254getSubstitution :: (Path, P.Optimized) -> Maybe [(String, Path)]
255getSubstitution (path, A.A _ pattern) =
256 case pattern of
257 P.Var x ->
258 Just [(x, path)]
260 P.Anything ->
261 Just []
263 P.Alias _ _ ->
264 error "aliases should never reach 'getSubstitution' function"
266 P.Record fields ->
267 Just (map (\name -> (name, add path (Field name Empty))) fields)
269 P.Data _ _ ->
270 Nothing
272 P.Literal _ ->
273 Nothing
276-- FIND RELEVANT TESTS
278testsAtPath :: Path -> [Branch] -> [Test]
279testsAtPath selectedPath branches =
280 List.nub $ Maybe.mapMaybe (testAtPath selectedPath) branches
283testAtPath :: Path -> Branch -> Maybe Test
284testAtPath selectedPath (Branch _ pathPatterns) =
285 case List.lookup selectedPath pathPatterns of
286 Nothing ->
287 Nothing
289 Just (A.A _ pattern) ->
290 case pattern of
291 P.Data name _ ->
292 Just (Constructor name)
294 P.Literal lit ->
295 Just (Literal lit)
297 P.Var _ ->
298 Nothing
300 P.Alias _ _ ->
301 error "aliases should never reach 'testAtPath' function"
303 P.Anything ->
304 Nothing
306 P.Record _ ->
307 Nothing
310-- BUILD EDGES
312edgesFor :: Path -> [Branch] -> Test -> (Test, [Branch])
313edgesFor path branches test =
314 ( test
315 , Maybe.mapMaybe (toRelevantBranch test path) branches
316 )
319toRelevantBranch :: Test -> Path -> Branch -> Maybe Branch
320toRelevantBranch test path branch@(Branch goal pathPatterns) =
321 case extract path pathPatterns of
322 Just (start, A.A _ pattern, end) ->
323 case pattern of
324 P.Data name patterns | test == Constructor name ->
325 Just (Branch goal (start ++ subPositions path patterns ++ end))
327 P.Literal lit | test == Literal lit ->
328 Just (Branch goal (start ++ end))
330 _ ->
331 Nothing
333 _ ->
334 Just branch
337extract
338 :: Path
339 -> [(Path, P.Optimized)]
340 -> Maybe ([(Path, P.Optimized)], P.Optimized, [(Path, P.Optimized)])
341extract selectedPath pathPatterns =
342 case pathPatterns of
343 [] ->
344 Nothing
346 first@(path, pattern) : rest ->
347 if path == selectedPath then
348 Just ([], pattern, rest)
350 else
351 case extract selectedPath rest of
352 Nothing ->
353 Nothing
355 Just (start, foundPattern, end) ->
356 Just (first : start, foundPattern, end)
359-- FIND IRRELEVANT BRANCHES
361isIrrelevantTo :: Path -> Branch -> Bool
362isIrrelevantTo selectedPath (Branch _ pathPatterns) =
363 case List.lookup selectedPath pathPatterns of
364 Nothing ->
365 True
367 Just pattern ->
368 not (needsTests pattern)
371needsTests :: P.Optimized -> Bool
372needsTests (A.A _ pattern) =
373 case pattern of
374 P.Var _ ->
375 False
377 P.Anything ->
378 False
380 P.Alias _ _ ->
381 error "aliases should never reach 'isIrrelevantTo' function"
383 P.Record _ ->
384 False
386 P.Data _ _ ->
387 True
389 P.Literal _ ->
390 True
393-- PATH PICKING HEURISTICS
395smallDefaults :: [Branch] -> Path
396smallDefaults branches =
397 let
398 allPaths =
399 Maybe.mapMaybe isChoicePath (concatMap _patterns branches)
401 weightedPaths =
402 map (\path -> (path, length (filter (isIrrelevantTo path) branches))) allPaths
403 in
404 fst (List.minimumBy (compare `on` snd) weightedPaths)
407isChoicePath :: (Path, P.Optimized) -> Maybe Path
408isChoicePath (path, pattern) =
409 if needsTests pattern then
410 Just path
411 else
412 Nothing
415-- smallBranchingFactor :: [Branch] -> Path