Introduce unboxedSum{Data,Type}Name to template-haskell

Summary:
In D2448 (which introduced Template Haskell support for unboxed
sums), I neglected to add `unboxedSumDataName` and `unboxedSumTypeName`
functions, since there wasn't any way you could write unboxed sum data or type
constructors in prefix form to begin with (see #12514). But even if you can't
write these `Name`s directly in source code, it would still be nice to be able
to use these `Name`s in Template Haskell (for instance, to be able to treat
unboxed sum type constructors like any other type constructors).

Along the way, this uncovered a minor bug in `isBuiltInOcc_maybe` in
`TysWiredIn`, which was calculating the arity of unboxed sum data constructors
incorrectly.

Test Plan: make test TEST=T12478_5

Reviewers: osa1, goldfire, austin, bgamari

Subscribers: thomie

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

GHC Trac Issues: #12478, #12514

6 files changed, 81 insertions, 4 deletions

diff --git a/compiler/basicTypes/Lexeme.hs b/compiler/basicTypes/Lexeme.hs
index 7012f5afed..dadc79ce21 100644
--- a/compiler/basicTypes/Lexeme.hs
+++ b/compiler/basicTypes/Lexeme.hs
@@ -156,8 +156,10 @@ okConIdOcc :: String -> Bool
 okConIdOcc str = okIdOcc str ||
                  is_tuple_name1 True  str ||
                    -- Is it a boxed tuple...
-                 is_tuple_name1 False str
-                   -- ...or an unboxed tuple (Trac #12407)?
+                 is_tuple_name1 False str ||
+                   -- ...or an unboxed tuple (Trac #12407)...
+                 is_sum_name1 str
+                   -- ...or an unboxed sum (Trac #12514)?
   where
     -- check for tuple name, starting at the beginning
     is_tuple_name1 True  ('(' : rest)       = is_tuple_name2 True  rest
@@ -172,6 +174,18 @@ okConIdOcc str = okIdOcc str ||
       | isSpace ws                    = is_tuple_name2 boxed rest
     is_tuple_name2 _     _            = False
 
+    -- check for sum name, starting at the beginning
+    is_sum_name1 ('(' : '#' : rest) = is_sum_name2 False rest
+    is_sum_name1 _                  = False
+
+    -- check for sum tail, only allowing at most one underscore
+    is_sum_name2 _          "#)"         = True
+    is_sum_name2 underscore ('|' : rest) = is_sum_name2 underscore rest
+    is_sum_name2 False      ('_' : rest) = is_sum_name2 True rest
+    is_sum_name2 underscore (ws  : rest)
+      | isSpace ws                       = is_sum_name2 underscore rest
+    is_sum_name2 _          _            = False
+
 -- | Is this an acceptable symbolic constructor name, assuming it
 -- starts with an acceptable character?
 okConSymOcc :: String -> Bool
diff --git a/compiler/prelude/TysWiredIn.hs b/compiler/prelude/TysWiredIn.hs
index 6e028fcf34..ce89e029e4 100644
--- a/compiler/prelude/TysWiredIn.hs
+++ b/compiler/prelude/TysWiredIn.hs
@@ -709,7 +709,7 @@ isBuiltInOcc_maybe occ =
         , Just rest'' <- "_" `stripPrefix` rest'
         , (pipes2, rest''') <- BS.span (=='|') rest''
         , "#)" <- rest'''
-             -> let arity = BS.length pipes1 + BS.length pipes2
+             -> let arity = BS.length pipes1 + BS.length pipes2 + 1
                     alt = BS.length pipes1 + 1
                 in Just $ dataConName $ sumDataCon alt arity
       _ -> Nothing
diff --git a/libraries/template-haskell/Language/Haskell/TH.hs b/libraries/template-haskell/Language/Haskell/TH.hs
index 5a497936ae..fd5c06f2f1 100644
--- a/libraries/template-haskell/Language/Haskell/TH.hs
+++ b/libraries/template-haskell/Language/Haskell/TH.hs
@@ -60,6 +60,8 @@ module Language.Haskell.TH(
         -- ** Built-in names
         tupleTypeName, tupleDataName,   -- Int -> Name
         unboxedTupleTypeName, unboxedTupleDataName, -- :: Int -> Name
+        unboxedSumTypeName, -- :: SumArity -> Name
+        unboxedSumDataName, -- :: SumAlt -> SumArity -> Name
 
     -- * The algebraic data types
     -- | The lowercase versions (/syntax operators/) of these constructors are
diff --git a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
index c9bccf665f..9de531ab9e 100644
--- a/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
+++ b/libraries/template-haskell/Language/Haskell/TH/Syntax.hs
@@ -1199,6 +1199,49 @@ mk_unboxed_tup_name n space
     n_commas = n - 1
     tup_mod  = mkModName "GHC.Tuple"
 
+-- Unboxed sum data and type constructors
+-- | Unboxed sum data constructor
+unboxedSumDataName :: SumAlt -> SumArity -> Name
+-- | Unboxed sum type constructor
+unboxedSumTypeName :: SumArity -> Name
+
+unboxedSumDataName alt arity
+  | alt > arity
+  = error $ prefix ++ "Index out of bounds." ++ debug_info
+
+  | alt <= 0
+  = error $ prefix ++ "Alt must be > 0." ++ debug_info
+
+  | arity < 2
+  = error $ prefix ++ "Arity must be >= 2." ++ debug_info
+
+  | otherwise
+  = Name (mkOccName sum_occ)
+         (NameG DataName (mkPkgName "ghc-prim") (mkModName "GHC.Prim"))
+
+  where
+    prefix     = "unboxedSumDataName: "
+    debug_info = " (alt: " ++ show alt ++ ", arity: " ++ show arity ++ ")"
+
+    -- Synced with the definition of mkSumDataConOcc in TysWiredIn
+    sum_occ = '(' : '#' : bars nbars_before ++ '_' : bars nbars_after ++ "#)"
+    bars i = replicate i '|'
+    nbars_before = alt - 1
+    nbars_after  = arity - alt
+
+unboxedSumTypeName arity
+  | arity < 2
+  = error $ "unboxedSumTypeName: Arity must be >= 2."
+         ++ " (arity: " ++ show arity ++ ")"
+
+  | otherwise
+  = Name (mkOccName sum_occ)
+         (NameG TcClsName (mkPkgName "ghc-prim") (mkModName "GHC.Prim"))
+
+  where
+    -- Synced with the definition of mkSumTyConOcc in TysWiredIn
+    sum_occ = '(' : '#' : replicate (arity - 1) '|' ++ "#)"
+
 -----------------------------------------------------
 --              Locations
 -----------------------------------------------------
diff --git a/testsuite/tests/th/T12478_5.hs b/testsuite/tests/th/T12478_5.hs
new file mode 100644
index 0000000000..bbbcb5502e
--- /dev/null
+++ b/testsuite/tests/th/T12478_5.hs
@@ -0,0 +1,17 @@
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE UnboxedSums #-}
+module T12478_5 where
+
+import Language.Haskell.TH
+
+foo :: $(conT (unboxedSumTypeName 2) `appT` conT ''() `appT` conT ''())
+    -> $(conT (unboxedSumTypeName 2) `appT` conT ''() `appT` conT ''())
+foo $(conP (unboxedSumDataName 1 2) [conP '() []])
+  = $(conE (unboxedSumDataName 2 2) `appE` conE '())
+foo $(conP (unboxedSumDataName 2 2) [conP '() []])
+  = $(conE (unboxedSumDataName 2 2) `appE` conE '())
+
+foo2 :: (# () | () #)
+     -> $(conT (unboxedSumTypeName 2) `appT` conT ''() `appT` conT ''())
+foo2 (# () | #) = $(conE (unboxedSumDataName 2 2) `appE` conE '())
+foo2 $(conP (unboxedSumDataName 2 2) [conP '() []]) = (# | () #)
diff --git a/testsuite/tests/th/all.T b/testsuite/tests/th/all.T
index ce4c5f58c8..66a7a9fc87 100644
--- a/testsuite/tests/th/all.T
+++ b/testsuite/tests/th/all.T
@@ -438,6 +438,7 @@ test('T12478_1', omit_ways(['ghci']), compile_and_run,
 test('T12478_2', omit_ways(['ghci']), compile_and_run, ['-v0'])
 test('T12478_3', omit_ways(['ghci']), compile, ['-v0'])
 test('T12478_4', omit_ways(['ghci']), compile_fail, ['-v0'])
+test('T12478_5', omit_ways(['ghci']), compile, ['-v0'])
 test('T12513', omit_ways(['ghci']), compile_fail, ['-v0'])
 test('T12530', normal, compile, ['-v0 -ddump-splices -dsuppress-uniques'])
 test('T12646', normal, compile, ['-v0'])
@@ -445,4 +446,4 @@ test('T12788', extra_clean(['T12788_Lib.hi', 'T12788_Lib.o']),
                multimod_compile_fail,
                ['T12788.hs', '-v0 ' + config.ghc_th_way_flags])
 test('T12977', normal, compile, ['-v0'])
-test('T12993', normal, multimod_compile, ['T12993.hs', '-v0'])
\ No newline at end of file
+test('T12993', normal, multimod_compile, ['T12993.hs', '-v0'])