Use pointer tag in dataToTag#

While looking at !2873 I noticed that dataToTag# previously didn't look
at a pointer's tag to determine its constructor. To be fair, there is a
bit of a trade-off here: using the pointer tag requires a bit more code
and another branch. On the other hand, it allows us to eliminate looking
at the info table in many cases (especially now since we tag large
constructor families; see #14373).

1 files changed, 35 insertions, 7 deletions

diff --git a/compiler/GHC/StgToCmm/Expr.hs b/compiler/GHC/StgToCmm/Expr.hs
index eb56a6ad09..aa10815330 100644
--- a/compiler/GHC/StgToCmm/Expr.hs
+++ b/compiler/GHC/StgToCmm/Expr.hs
@@ -36,7 +36,7 @@ import GHC.Cmm.Graph
 import GHC.Cmm.BlockId
 import GHC.Cmm hiding ( succ )
 import GHC.Cmm.Info
-import GHC.Cmm.Utils ( mAX_PTR_TAG )
+import GHC.Cmm.Utils ( zeroExpr, cmmTagMask, mAX_PTR_TAG )
 import GHC.Core
 import GHC.Core.DataCon
 import GHC.Types.ForeignCall
@@ -70,15 +70,43 @@ cgExpr (StgOpApp (StgPrimOp SeqOp) [StgVarArg a, _] _res_ty) =
   cgIdApp a []
 
 -- dataToTag# :: a -> Int#
--- See Note [dataToTag#] in primops.txt.pp
+-- See Note [dataToTag# magic] in primops.txt.pp
 cgExpr (StgOpApp (StgPrimOp DataToTagOp) [StgVarArg a] _res_ty) = do
   platform <- getPlatform
   emitComment (mkFastString "dataToTag#")
-  tmp <- newTemp (bWord platform)
-  _ <- withSequel (AssignTo [tmp] False) (cgIdApp a [])
-  -- TODO: For small types look at the tag bits instead of reading info table
-  ptr_opts <- getPtrOpts
-  emitReturn [getConstrTag ptr_opts (cmmUntag platform (CmmReg (CmmLocal tmp)))]
+  info <- getCgIdInfo a
+  tag_reg <- assignTemp $ cmmConstrTag1 dflags (idInfoToAmode info)
+  result_reg <- newTemp (bWord dflags)
+  let tag = CmmReg $ CmmLocal tag_reg
+  -- Here we will first check the tag bits of the pointer we were given;
+  -- if this doesn't work then enter the closure and use the info table
+  -- to determine the constructor. Note that all tag bits set means that
+  -- the constructor index is too large to fit in the pointer and therefore
+  -- we must look in the info table. See Note [Tagging big families].
+
+  slow_path <- getCode $ do
+      tmp <- newTemp (bWord platform)
+      _ <- withSequel (AssignTo [tmp] False) (cgIdApp a [])
+      -- TODO: For small types look at the tag bits instead of reading info table
+      emitAssign (CmmLocal result_reg)
+        $ getConstrTag dflags (cmmUntag platform (CmmReg (CmmLocal tmp)))
+
+  fast_path <- getCode $ do
+      emitAssign (CmmLocal result_reg)
+        $ cmmSubWord dflags tag (CmmLit $ mkWordCLit platform 1)
+
+  let zero = zeroExpr platform
+      too_big_tag = cmmTagMask platform
+      is_tagged =
+        cmmNeWord platform
+          (cmmOrWord platform
+            (cmmEqWord platform tag zero)         -- not evaluated
+            (cmmEqWord platform tag too_big_tag)) -- tag too big
+          zero
+
+  emit =<< mkCmmIfThenElse' is_tagged slow_path fast_path (Just False)
+  emitReturn [CmmReg $ CmmLocal result_reg]
+
 
 cgExpr (StgOpApp op args ty) = cgOpApp op args ty
 cgExpr (StgConApp con args _)= cgConApp con args