Add an (optional) range field to the SwitchTargets data type

As there is one in the Cmm syntax, and we might be able to exploit that
during code generation.

3 files changed, 40 insertions, 31 deletions

diff --git a/compiler/cmm/CmmNode.hs b/compiler/cmm/CmmNode.hs
index a3edc3f0e2..4b3dfd2d40 100644
--- a/compiler/cmm/CmmNode.hs
+++ b/compiler/cmm/CmmNode.hs
@@ -698,25 +698,34 @@ combineTickScopes s1 s2
 
 -- See Note [Switch Table]
 data SwitchTargets =
-    SwitchTargets (Maybe Label) (M.Map Integer Label)
+    SwitchTargets (Maybe (Integer, Integer)) (Maybe Label) (M.Map Integer Label)
     deriving Eq
 
-mkSwitchTargets :: Maybe Label -> M.Map Integer Label -> SwitchTargets
-mkSwitchTargets = SwitchTargets
+-- mkSwitchTargets normalises the map a bit:
+mkSwitchTargets :: Maybe (Integer, Integer) -> Maybe Label -> M.Map Integer Label -> SwitchTargets
+mkSwitchTargets mbrange mbdef ids
+    = SwitchTargets mbrange mbdef $ dropDefault $ restrict ids
+  where
+    -- It drops entries outside the range, if there is a range
+    restrict | Just (lo,hi) <- mbrange = M.filterWithKey (\x _ -> lo <= x && x <= hi)
+             | otherwise               = id
+    -- It entries that equal the default, if there is a default
+    dropDefault | Just l <- mbdef = M.filter (/= l)
+                | otherwise       = id
 
 mapSwitchTargets :: (Label -> Label) -> SwitchTargets -> SwitchTargets
-mapSwitchTargets f (SwitchTargets mbdef branches)
-    = SwitchTargets (fmap f mbdef) (fmap f branches)
+mapSwitchTargets f (SwitchTargets range mbdef branches)
+    = SwitchTargets range (fmap f mbdef) (fmap f branches)
 
 switchTargetsCases :: SwitchTargets -> [(Integer, Label)]
-switchTargetsCases (SwitchTargets _ branches) = M.toList branches
+switchTargetsCases (SwitchTargets _ _ branches) = M.toList branches
 
 switchTargetsDefault :: SwitchTargets -> Maybe Label
-switchTargetsDefault (SwitchTargets mbdef _) = mbdef
+switchTargetsDefault (SwitchTargets _ mbdef _) = mbdef
 
 switchTargetsToTable :: SwitchTargets -> [Maybe Label]
-switchTargetsToTable (SwitchTargets mbdef branches)
-    | min < 0 =   pprPanic "mapSwitchTargets" empty
+switchTargetsToTable (SwitchTargets _ mbdef branches)
+    | min < 0 = pprPanic "mapSwitchTargets" empty
     | otherwise = [ labelFor i | i <- [0..max] ]
   where
     min = fst (M.findMin branches)
@@ -725,20 +734,21 @@ switchTargetsToTable (SwitchTargets mbdef branches)
                                              Nothing -> mbdef
 
 switchTargetsToList :: SwitchTargets -> [Label]
-switchTargetsToList (SwitchTargets mbdef branches) =  maybeToList mbdef ++ M.elems branches
+switchTargetsToList (SwitchTargets _ mbdef branches)
+    = maybeToList mbdef ++ M.elems branches
 
 -- | Groups cases with equal targets, suitable for pretty-printing to a
 -- c-like switch statement with fall-through semantics.
 switchTargetsFallThrough :: SwitchTargets -> ([([Integer], Label)], Maybe Label)
-switchTargetsFallThrough (SwitchTargets mbdef branches) = (groups, mbdef)
+switchTargetsFallThrough (SwitchTargets _ mbdef branches) = (groups, mbdef)
   where
     groups = map (\xs -> (map fst xs, snd (head xs))) $
              groupBy ((==) `on` snd) $
              M.toList branches
 
 eqSwitchTargetWith :: (Label -> Label -> Bool) -> SwitchTargets -> SwitchTargets -> Bool
-eqSwitchTargetWith eq (SwitchTargets mbdef1 ids1) (SwitchTargets mbdef2 ids2) =
-    goMB mbdef1 mbdef2 && goList (M.toList ids1) (M.toList ids2)
+eqSwitchTargetWith eq (SwitchTargets range1 mbdef1 ids1) (SwitchTargets range2 mbdef2 ids2) =
+    range1 == range2 && goMB mbdef1 mbdef2 && goList (M.toList ids1) (M.toList ids2)
   where
     goMB Nothing Nothing = True
     goMB (Just l1) (Just l2) = l1 `eq` l2
diff --git a/compiler/cmm/CmmParse.y b/compiler/cmm/CmmParse.y
index 4f286f5e32..7ec1e4a30b 100644
--- a/compiler/cmm/CmmParse.y
+++ b/compiler/cmm/CmmParse.y
@@ -677,24 +677,24 @@ globals :: { [GlobalReg] }
         : GLOBALREG                     { [$1] }
         | GLOBALREG ',' globals         { $1 : $3 }
 
-maybe_range :: { Maybe (Int,Int) }
-        : '[' INT '..' INT ']'  { Just (fromIntegral $2, fromIntegral $4) }
+maybe_range :: { Maybe (Integer,Integer) }
+        : '[' INT '..' INT ']'  { Just ($2, $4) }
         | {- empty -}           { Nothing }
 
-arms    :: { [CmmParse ([Int],Either BlockId (CmmParse ()))] }
+arms    :: { [CmmParse ([Integer],Either BlockId (CmmParse ()))] }
         : {- empty -}                   { [] }
         | arm arms                      { $1 : $2 }
 
-arm     :: { CmmParse ([Int],Either BlockId (CmmParse ())) }
+arm     :: { CmmParse ([Integer],Either BlockId (CmmParse ())) }
         : 'case' ints ':' arm_body      { do b <- $4; return ($2, b) }
 
 arm_body :: { CmmParse (Either BlockId (CmmParse ())) }
         : '{' body '}'                  { return (Right (withSourceNote $1 $3 $2)) }
         | 'goto' NAME ';'               { do l <- lookupLabel $2; return (Left l) }
 
-ints    :: { [Int] }
-        : INT                           { [ fromIntegral $1 ] }
-        | INT ',' ints                  { fromIntegral $1 : $3 }
+ints    :: { [Integer] }
+        : INT                           { [ $1 ] }
+        | INT ',' ints                  { $1 : $3 }
 
 default :: { Maybe (CmmParse ()) }
         : 'default' ':' '{' body '}'    { Just (withSourceNote $3 $5 $4) }
@@ -1308,7 +1308,9 @@ withSourceNote a b parse = do
 -- optional range on the switch (eg. switch [0..7] {...}), or by
 -- the minimum/maximum values from the branches.
 
-doSwitch :: Maybe (Int,Int) -> CmmParse CmmExpr -> [([Int],Either BlockId (CmmParse ()))]
+doSwitch :: Maybe (Integer,Integer)
+         -> CmmParse CmmExpr
+         -> [([Integer],Either BlockId (CmmParse ()))]
          -> Maybe (CmmParse ()) -> CmmParse ()
 doSwitch mb_range scrut arms deflt
    = do
@@ -1326,13 +1328,13 @@ doSwitch mb_range scrut arms deflt
 
         expr <- scrut
         -- ToDo: check for out of range and jump to default if necessary
-        emit $ mkSwitch expr (mkSwitchTargets dflt_entry table)
+        emit $ mkSwitch expr (mkSwitchTargets mb_range dflt_entry table)
    where
-        emitArm :: ([Int],Either BlockId (CmmParse ())) -> CmmParse [(Integer,BlockId)]
-        emitArm (ints,Left blockid) = return [ (fromIntegral i,blockid) | i <- ints ]
+        emitArm :: ([Integer],Either BlockId (CmmParse ())) -> CmmParse [(Integer,BlockId)]
+        emitArm (ints,Left blockid) = return [ (i,blockid) | i <- ints ]
         emitArm (ints,Right code) = do
            blockid <- forkLabelledCode code
-           return [ (fromIntegral i,blockid) | i <- ints ]
+           return [ (i,blockid) | i <- ints ]
 
 forkLabelledCode :: CmmParse () -> CmmParse BlockId
 forkLabelledCode p = do
diff --git a/compiler/codeGen/StgCmmUtils.hs b/compiler/codeGen/StgCmmUtils.hs
index df62df0356..11864d7b5c 100644
--- a/compiler/codeGen/StgCmmUtils.hs
+++ b/compiler/codeGen/StgCmmUtils.hs
@@ -544,15 +544,12 @@ mk_switch tag_expr branches mb_deflt lo_tag hi_tag via_C
         -- either end of the range (see below), so the first
         -- tag of a real branch is real_lo_tag (not lo_tag).
         arms :: M.Map Integer BlockId
-        arms = M.fromList [ (fromIntegral (i - real_lo_tag), l)
-                          | (i,l) <- branches
-                          , real_lo_tag <= i
-                          , i <= real_hi_tag
-                          ]
+        arms = M.fromList [ (fromIntegral i, l) | (i,l) <- branches ]
+       
        dflags <- getDynFlags
        return $ mkSwitch
            (cmmOffset dflags tag_expr (- real_lo_tag))
-           (mkSwitchTargets mb_deflt arms)
+           (mkSwitchTargets (Just (0, fromIntegral (real_hi_tag-real_lo_tag))) mb_deflt arms)
 
   -- if we can knock off a bunch of default cases with one if, then do so
   | Just deflt <- mb_deflt, (lowest_branch - lo_tag) >= n_branches