Fix and document Unique generation for sum TyCon and DataCons

Test Plan: validate

Reviewers: bgamari, austin

Reviewed By: bgamari

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D2420

2 files changed, 39 insertions, 8 deletions

diff --git a/compiler/basicTypes/Unique.hs b/compiler/basicTypes/Unique.hs
index 800198bc95..545ea9f1ba 100644
--- a/compiler/basicTypes/Unique.hs
+++ b/compiler/basicTypes/Unique.hs
@@ -329,11 +329,9 @@ mkAlphaTyVarUnique     :: Int -> Unique
 mkPreludeClassUnique   :: Int -> Unique
 mkPreludeTyConUnique   :: Int -> Unique
 mkTupleTyConUnique     :: Boxity -> Arity -> Unique
-mkSumTyConUnique       :: Arity -> Unique
 mkCTupleTyConUnique    :: Arity -> Unique
 mkPreludeDataConUnique :: Arity -> Unique
 mkTupleDataConUnique   :: Boxity -> Arity -> Unique
-mkSumDataConUnique     :: ConTagZ -> Arity -> Unique
 mkPrimOpIdUnique       :: Int -> Unique
 mkPreludeMiscIdUnique  :: Int -> Unique
 mkPArrDataConUnique    :: Int -> Unique
@@ -351,7 +349,6 @@ mkPreludeTyConUnique i                = mkUnique '3' (2*i)
 mkTupleTyConUnique Boxed           a  = mkUnique '4' (2*a)
 mkTupleTyConUnique Unboxed         a  = mkUnique '5' (2*a)
 mkCTupleTyConUnique                a  = mkUnique 'k' (2*a)
-mkSumTyConUnique                   a  = mkUnique 'z' (2*a)
 
 tyConRepNameUnique :: Unique -> Unique
 tyConRepNameUnique  u = incrUnique u
@@ -372,12 +369,35 @@ tyConRepNameUnique  u = incrUnique u
 mkPreludeDataConUnique i              = mkUnique '6' (3*i)    -- Must be alphabetic
 mkTupleDataConUnique Boxed          a = mkUnique '7' (3*a)    -- ditto (*may* be used in C labels)
 mkTupleDataConUnique Unboxed        a = mkUnique '8' (3*a)
+
+--------------------------------------------------
+-- Sum arities start from 2. A sum of arity N has N data constructors, so it
+-- occupies N+1 slots: 1 TyCon + N DataCons.
+--
+-- So arity 2 sum takes uniques 0 (tycon), 1, 2  (2 data cons)
+--    arity 3 sum takes uniques 3 (tycon), 4, 5, 6 (3 data cons)
+-- etc.
+
+mkSumTyConUnique :: Arity -> Unique
+mkSumTyConUnique arity = mkUnique 'z' (sumUniqsOccupied arity)
+
+mkSumDataConUnique :: ConTagZ -> Arity -> Unique
 mkSumDataConUnique alt arity
   | alt >= arity
   = panic ("mkSumDataConUnique: " ++ show alt ++ " >= " ++ show arity)
   | otherwise
-  = mkUnique 'z' (2 * alt * arity)
+  = mkUnique 'z' (sumUniqsOccupied arity + alt + 1 {- skip the tycon -})
+
+-- How many unique slots occupied by sum types (including constructors) up to
+-- the given arity?
+sumUniqsOccupied :: Arity -> Int
+sumUniqsOccupied arity
+  = ASSERT(arity >= 2)
+    -- 3 + 4 + ... + arity
+    ((arity * (arity + 1)) `div` 2) - 3
+{-# INLINE sumUniqsOccupied #-}
 
+--------------------------------------------------
 dataConRepNameUnique, dataConWorkerUnique :: Unique -> Unique
 dataConWorkerUnique  u = incrUnique u
 dataConRepNameUnique u = stepUnique u 2
diff --git a/compiler/prelude/TysWiredIn.hs b/compiler/prelude/TysWiredIn.hs
index 0775d06a11..102847833e 100644
--- a/compiler/prelude/TysWiredIn.hs
+++ b/compiler/prelude/TysWiredIn.hs
@@ -861,8 +861,15 @@ mkSumDataConOcc alt n = mkOccName dataName str
 
 -- | Type constructor for n-ary unboxed sum.
 sumTyCon :: Arity -> TyCon
-sumTyCon n | n > mAX_SUM_SIZE = fst (mk_sum n)  -- Build one specially
-sumTyCon n = fst (unboxedSumArr ! n)
+sumTyCon arity
+  | arity > mAX_SUM_SIZE
+  = fst (mk_sum arity)  -- Build one specially
+
+  | arity < 2
+  = panic ("sumTyCon: Arity starts from 2. (arity: " ++ show arity ++ ")")
+
+  | otherwise
+  = fst (unboxedSumArr ! arity)
 
 -- | Data constructor for i-th alternative of a n-ary unboxed sum.
 sumDataCon :: ConTag -- Alternative
@@ -870,13 +877,17 @@ sumDataCon :: ConTag -- Alternative
            -> DataCon
 sumDataCon alt arity
   | alt > arity
-  = panic ("sumDataCon: index out of bounds: alt "
+  = panic ("sumDataCon: index out of bounds: alt: "
            ++ show alt ++ " > arity " ++ show arity)
 
   | alt <= 0
   = panic ("sumDataCon: Alts start from 1. (alt: " ++ show alt
            ++ ", arity: " ++ show arity ++ ")")
 
+  | arity < 2
+  = panic ("sumDataCon: Arity starts from 2. (alt: " ++ show alt
+           ++ ", arity: " ++ show arity ++ ")")
+
   | arity > mAX_SUM_SIZE
   = snd (mk_sum arity) ! (alt - 1)  -- Build one specially
 
@@ -887,7 +898,7 @@ sumDataCon alt arity
 -- indexed by the arity of the sum and the inner array is indexed by
 -- the alternative.
 unboxedSumArr :: Array Int (TyCon, Array Int DataCon)
-unboxedSumArr = listArray (0,mAX_SUM_SIZE) [mk_sum i | i <- [0..mAX_SUM_SIZE]]
+unboxedSumArr = listArray (2,mAX_SUM_SIZE) [mk_sum i | i <- [2..mAX_SUM_SIZE]]
 
 -- | Create type constructor and data constructors for n-ary unboxed sum.
 mk_sum :: Arity -> (TyCon, Array ConTagZ DataCon)