git.iko.soy/elk
on befb3df help sign in
elk/src/Optimize/Case.hs
174 lines 4.4 kB history raw
1{-# OPTIONS_GHC -Wall #-}
2module Optimize.Case (optimize) where
4import Control.Arrow (second)
5import qualified Data.Map as Map
6import Data.Map ((!))
7import qualified Data.Maybe as Maybe
8import Data.Text (Text)
10import qualified AST.Expression.Optimized as Opt
11import qualified AST.Pattern as P
12import qualified Optimize.DecisionTree as DT
16-- OPTIMIZE A CASE EXPRESSION
19optimize :: DT.VariantDict -> Text -> [(P.Canonical, Opt.Expr)] -> Opt.Expr
20optimize variantDict exprName optBranches =
21 let
22 (patterns, indexedBranches) =
23 unzip (zipWith indexify [0..] optBranches)
25 decisionTree =
26 DT.compile variantDict patterns
27 in
28 treeToExpr exprName decisionTree indexedBranches
31indexify :: Int -> (a,b) -> ((a,Int), (Int,b))
32indexify index (pattern, branch) =
33 ( (pattern, index)
34 , (index, branch)
35 )
39-- CONVERT A TREE TO AN EXPRESSION
42treeToExpr :: Text -> DT.DecisionTree -> [(Int, Opt.Expr)] -> Opt.Expr
43treeToExpr name decisionTree allJumps =
44 let
45 decider =
46 treeToDecider decisionTree
48 targetCounts =
49 countTargets decider
51 (choices, maybeJumps) =
52 unzip (map (createChoices targetCounts) allJumps)
53 in
54 Opt.Case
55 name
56 (insertChoices (Map.fromList choices) decider)
57 (Maybe.catMaybes maybeJumps)
61-- TREE TO DECIDER
62--
63-- Decision trees may have some redundancies, so we convert them to a Decider
64-- which has special constructs to avoid code duplication when possible.
67treeToDecider :: DT.DecisionTree -> Opt.Decider Int
68treeToDecider tree =
69 case tree of
70 DT.Match target ->
71 Opt.Leaf target
73 -- zero options
74 DT.Decision _ [] Nothing ->
75 error "compiler bug, somehow created an empty decision tree"
77 -- one option
78 DT.Decision _ [(_, subTree)] Nothing ->
79 treeToDecider subTree
81 DT.Decision _ [] (Just subTree) ->
82 treeToDecider subTree
84 -- two options
85 DT.Decision path [(test, successTree)] (Just failureTree) ->
86 toChain path test successTree failureTree
88 DT.Decision path [(test, successTree), (_, failureTree)] Nothing ->
89 toChain path test successTree failureTree
91 -- many options
92 DT.Decision path edges Nothing ->
93 let
94 (necessaryTests, fallback) =
95 (init edges, snd (last edges))
96 in
97 Opt.FanOut
98 path
99 (map (second treeToDecider) necessaryTests)
100 (treeToDecider fallback)
102 DT.Decision path edges (Just fallback) ->
103 Opt.FanOut path (map (second treeToDecider) edges) (treeToDecider fallback)
106toChain :: DT.Path -> DT.Test -> DT.DecisionTree -> DT.DecisionTree -> Opt.Decider Int
107toChain path test successTree failureTree =
108 let
109 failure =
110 treeToDecider failureTree
111 in
112 case treeToDecider successTree of
113 Opt.Chain testChain success subFailure | failure == subFailure ->
114 Opt.Chain ((path, test) : testChain) success failure
116 success ->
117 Opt.Chain [(path, test)] success failure
121-- INSERT CHOICES
122--
123-- If a target appears exactly once in a Decider, the corresponding expression
124-- can be inlined. Whether things are inlined or jumps is called a "choice".
127countTargets :: Opt.Decider Int -> Map.Map Int Int
128countTargets decisionTree =
129 case decisionTree of
130 Opt.Leaf target ->
131 Map.singleton target 1
133 Opt.Chain _ success failure ->
134 Map.unionWith (+) (countTargets success) (countTargets failure)
136 Opt.FanOut _ tests fallback ->
137 Map.unionsWith (+) (map countTargets (fallback : map snd tests))
140createChoices
141 :: Map.Map Int Int
142 -> (Int, Opt.Expr)
143 -> ( (Int, Opt.Choice), Maybe (Int, Opt.Expr) )
144createChoices targetCounts (target, branch) =
145 if targetCounts ! target == 1 then
146 ( (target, Opt.Inline branch)
147 , Nothing
148 )
150 else
151 ( (target, Opt.Jump target)
152 , Just (target, branch)
153 )
156insertChoices
157 :: Map.Map Int Opt.Choice
158 -> Opt.Decider Int
159 -> Opt.Decider Opt.Choice
160insertChoices choiceDict decider =
161 let
162 go =
163 insertChoices choiceDict
164 in
165 case decider of
166 Opt.Leaf target ->
167 Opt.Leaf (choiceDict ! target)
169 Opt.Chain testChain success failure ->
170 Opt.Chain testChain (go success) (go failure)
172 Opt.FanOut path tests fallback ->
173 Opt.FanOut path (map (second go) tests) (go fallback)