Support unlifted datatypes in GHCi

fixes #19628

9 files changed, 204 insertions, 61 deletions

diff --git a/compiler/GHC/ByteCode/Asm.hs b/compiler/GHC/ByteCode/Asm.hs
index 6b9d4a9223..6ef6ef474c 100644
--- a/compiler/GHC/ByteCode/Asm.hs
+++ b/compiler/GHC/ByteCode/Asm.hs
@@ -454,7 +454,7 @@ assembleI platform i = case i of
   JMP       l              -> emit bci_JMP [LabelOp l]
   ENTER                    -> emit bci_ENTER []
   RETURN                   -> emit bci_RETURN []
-  RETURN_UBX rep           -> emit (return_ubx rep) []
+  RETURN_UNLIFTED rep      -> emit (return_unlifted rep) []
   RETURN_TUPLE             -> emit bci_RETURN_T []
   CCALL off m_addr i       -> do np <- addr m_addr
                                  emit bci_CCALL [SmallOp off, Op np, SmallOp i]
@@ -522,16 +522,16 @@ push_alts V16 = error "push_alts: vector"
 push_alts V32 = error "push_alts: vector"
 push_alts V64 = error "push_alts: vector"
 
-return_ubx :: ArgRep -> Word16
-return_ubx V   = bci_RETURN_V
-return_ubx P   = bci_RETURN_P
-return_ubx N   = bci_RETURN_N
-return_ubx L   = bci_RETURN_L
-return_ubx F   = bci_RETURN_F
-return_ubx D   = bci_RETURN_D
-return_ubx V16 = error "return_ubx: vector"
-return_ubx V32 = error "return_ubx: vector"
-return_ubx V64 = error "return_ubx: vector"
+return_unlifted :: ArgRep -> Word16
+return_unlifted V   = bci_RETURN_V
+return_unlifted P   = bci_RETURN_P
+return_unlifted N   = bci_RETURN_N
+return_unlifted L   = bci_RETURN_L
+return_unlifted F   = bci_RETURN_F
+return_unlifted D   = bci_RETURN_D
+return_unlifted V16 = error "return_unlifted: vector"
+return_unlifted V32 = error "return_unlifted: vector"
+return_unlifted V64 = error "return_unlifted: vector"
 
 {-
   we can only handle up to a fixed number of words on the stack,
diff --git a/compiler/GHC/ByteCode/Instr.hs b/compiler/GHC/ByteCode/Instr.hs
index cb1fb4d335..1950289f79 100644
--- a/compiler/GHC/ByteCode/Instr.hs
+++ b/compiler/GHC/ByteCode/Instr.hs
@@ -172,9 +172,9 @@ data BCInstr
 
    -- To Infinity And Beyond
    | ENTER
-   | RETURN            -- return a lifted value
-   | RETURN_UBX ArgRep -- return an unlifted value, here's its rep
-   | RETURN_TUPLE      -- return an unboxed tuple (info already on stack)
+   | RETURN                 -- return a lifted value
+   | RETURN_UNLIFTED ArgRep -- return an unlifted value, here's its rep
+   | RETURN_TUPLE           -- return an unboxed tuple (info already on stack)
 
    -- Breakpoints
    | BRK_FUN          Word16 Unique (RemotePtr CostCentre)
@@ -310,7 +310,7 @@ instance Outputable BCInstr where
                                                <+> text "by" <+> ppr n
    ppr ENTER                 = text "ENTER"
    ppr RETURN                = text "RETURN"
-   ppr (RETURN_UBX pk)       = text "RETURN_UBX  " <+> ppr pk
+   ppr (RETURN_UNLIFTED pk)  = text "RETURN_UNLIFTED  " <+> ppr pk
    ppr (RETURN_TUPLE)        = text "RETURN_TUPLE"
    ppr (BRK_FUN index uniq _cc) = text "BRK_FUN" <+> ppr index <+> ppr uniq <+> text "<cc>"
 
@@ -390,7 +390,7 @@ bciStackUse CASEFAIL{}            = 0
 bciStackUse JMP{}                 = 0
 bciStackUse ENTER{}               = 0
 bciStackUse RETURN{}              = 0
-bciStackUse RETURN_UBX{}          = 1 -- pushes stg_ret_X for some X
+bciStackUse RETURN_UNLIFTED{}     = 1 -- pushes stg_ret_X for some X
 bciStackUse RETURN_TUPLE{}        = 1 -- pushes stg_ret_t header
 bciStackUse CCALL{}               = 0
 bciStackUse SWIZZLE{}             = 0
diff --git a/compiler/GHC/StgToByteCode.hs b/compiler/GHC/StgToByteCode.hs
index 37a6539fe6..64de0ff05e 100644
--- a/compiler/GHC/StgToByteCode.hs
+++ b/compiler/GHC/StgToByteCode.hs
@@ -439,7 +439,10 @@ schemeTopBind (id, rhs)
         -- by just re-using the single top-level definition.  So
         -- for the worker itself, we must allocate it directly.
     -- ioToBc (putStrLn $ "top level BCO")
-    emitBc (mkProtoBCO platform (getName id) (toOL [PACK data_con 0, ENTER])
+    let enter = if isUnliftedTypeKind (tyConResKind (dataConTyCon data_con))
+                then RETURN_UNLIFTED P
+                else ENTER
+    emitBc (mkProtoBCO platform (getName id) (toOL [PACK data_con 0, enter])
                        (Right rhs) 0 0 [{-no bitmap-}] False{-not alts-})
 
   | otherwise
@@ -575,36 +578,36 @@ fvsToEnv p rhs =  [v | v <- dVarSetElems $ freeVarsOfRhs rhs,
 
 -- Returning an unlifted value.
 -- Heave it on the stack, SLIDE, and RETURN.
-returnUnboxedAtom
+returnUnliftedAtom
     :: StackDepth
     -> Sequel
     -> BCEnv
     -> StgArg
     -> BcM BCInstrList
-returnUnboxedAtom d s p e = do
+returnUnliftedAtom d s p e = do
     let reps = case e of
                  StgLitArg lit -> typePrimRepArgs (literalType lit)
                  StgVarArg i   -> bcIdPrimReps i
     (push, szb) <- pushAtom d p e
-    ret <- returnUnboxedReps d s szb reps
+    ret <- returnUnliftedReps d s szb reps
     return (push `appOL` ret)
 
--- return an unboxed value from the top of the stack
-returnUnboxedReps
+-- return an unlifted value from the top of the stack
+returnUnliftedReps
     :: StackDepth
     -> Sequel
     -> ByteOff    -- size of the thing we're returning
     -> [PrimRep]  -- representations
     -> BcM BCInstrList
-returnUnboxedReps d s szb reps = do
+returnUnliftedReps d s szb reps = do
     profile <- getProfile
     let platform = profilePlatform profile
         non_void VoidRep = False
         non_void _ = True
     ret <- case filter non_void reps of
              -- use RETURN_UBX for unary representations
-             []    -> return (unitOL $ RETURN_UBX V)
-             [rep] -> return (unitOL $ RETURN_UBX (toArgRep platform rep))
+             []    -> return (unitOL $ RETURN_UNLIFTED V)
+             [rep] -> return (unitOL $ RETURN_UNLIFTED (toArgRep platform rep))
              -- otherwise use RETURN_TUPLE with a tuple descriptor
              nv_reps -> do
                let (tuple_info, args_offsets) = layoutTuple profile 0 (primRepCmmType platform) nv_reps
@@ -633,19 +636,19 @@ returnUnboxedTuple d s p es = do
                                          massert (off == dd + szb)
                                          go (dd + szb) (push:pushes) cs
     pushes <- go d [] tuple_components
-    ret <- returnUnboxedReps d
-                             s
-                             (wordsToBytes platform $ tupleSize tuple_info)
-                             (map atomPrimRep es)
+    ret <- returnUnliftedReps d
+                              s
+                              (wordsToBytes platform $ tupleSize tuple_info)
+                              (map atomPrimRep es)
     return (mconcat pushes `appOL` ret)
 
 -- Compile code to apply the given expression to the remaining args
 -- on the stack, returning a HNF.
 schemeE
     :: StackDepth -> Sequel -> BCEnv -> CgStgExpr -> BcM BCInstrList
-schemeE d s p (StgLit lit) = returnUnboxedAtom d s p (StgLitArg lit)
+schemeE d s p (StgLit lit) = returnUnliftedAtom d s p (StgLitArg lit)
 schemeE d s p (StgApp x [])
-   | isUnliftedType (idType x) = returnUnboxedAtom d s p (StgVarArg x)
+   | isUnliftedType (idType x) = returnUnliftedAtom d s p (StgVarArg x)
 -- Delegate tail-calls to schemeT.
 schemeE d s p e@(StgApp {}) = schemeT d s p e
 schemeE d s p e@(StgConApp {}) = schemeT d s p e
@@ -884,7 +887,9 @@ schemeT d s p (StgConApp con _ext args _tys)
         platform <- profilePlatform <$> getProfile
         return (alloc_con         `appOL`
                 mkSlideW 1 (bytesToWords platform $ d - s) `snocOL`
-                ENTER)
+                if isUnliftedTypeKind (tyConResKind (dataConTyCon con))
+                then RETURN_UNLIFTED P
+                else ENTER)
 
    -- Case 4: Tail call of function
 schemeT d s p (StgApp fn args)
@@ -952,7 +957,10 @@ doTailCall init_d s p fn args = do
         platform <- profilePlatform <$> getProfile
         assert (sz == wordSize platform ) return ()
         let slide = mkSlideB platform (d - init_d + wordSize platform) (init_d - s)
-        return (push_fn `appOL` (slide `appOL` unitOL ENTER))
+            enter = if isUnliftedType (idType fn)
+                    then RETURN_UNLIFTED P
+                    else ENTER
+        return (push_fn `appOL` (slide `appOL` unitOL enter))
   do_pushes !d args reps = do
       let (push_apply, n, rest_of_reps) = findPushSeq reps
           (these_args, rest_of_args) = splitAt n args
@@ -1049,9 +1057,9 @@ doCase d s p scrut bndr alts
         -- An unlifted value gets an extra info table pushed on top
         -- when it is returned.
         unlifted_itbl_size_b :: StackDepth
-        unlifted_itbl_size_b | isAlgCase       = 0
-                             | ubx_tuple_frame = 3 * wordSize platform
-                             | otherwise       = wordSize platform
+        unlifted_itbl_size_b | ubx_tuple_frame              = 3 * wordSize platform
+                             | not (isUnliftedType bndr_ty) = 0
+                             | otherwise                    = wordSize platform
 
         (bndr_size, tuple_info, args_offsets)
            | ubx_tuple_frame =
@@ -1072,7 +1080,7 @@ doCase d s p scrut bndr alts
         d_bndr =
             d + ret_frame_size_b + bndr_size
 
-        -- depth of stack after the extra info table for an unboxed return
+        -- depth of stack after the extra info table for an unlifted return
         -- has been pushed, if any.  This is the stack depth at the
         -- continuation.
         d_alts = d + ret_frame_size_b + bndr_size + unlifted_itbl_size_b
@@ -1082,7 +1090,7 @@ doCase d s p scrut bndr alts
         p_alts = Map.insert bndr d_bndr p
 
         bndr_ty = idType bndr
-        isAlgCase = not (isUnliftedType bndr_ty)
+        isAlgCase = isAlgType bndr_ty
 
         -- given an alt, return a discr and code for it.
         codeAlt (DEFAULT, _, rhs)
@@ -1131,11 +1139,17 @@ doCase d s p scrut bndr alts
                         [ (arg, stack_bot - ByteOff offset)
                         | (NonVoid arg, offset) <- args_offsets ]
                         p_alts
+
+                 -- unlifted datatypes have an infotable word on top
+                 unpack = if isUnliftedType bndr_ty
+                          then PUSH_L 1 `consOL`
+                               UNPACK (trunc16W size) `consOL`
+                               unitOL (SLIDE (trunc16W size) 1)
+                          else unitOL (UNPACK (trunc16W size))
              in do
              massert isAlgCase
              rhs_code <- schemeE stack_bot s p' rhs
-             return (my_discr alt,
-                     unitOL (UNPACK (trunc16W size)) `appOL` rhs_code)
+             return (my_discr alt, unpack `appOL` rhs_code)
            where
              real_bndrs = filterOut isTyVar bndrs
 
@@ -1224,7 +1238,7 @@ doCase d s p scrut bndr alts
               return (PUSH_ALTS_TUPLE alt_bco' tuple_info tuple_bco
                       `consOL` scrut_code)
        else let push_alts
-                  | isAlgCase
+                  | not (isUnliftedType bndr_ty)
                   = PUSH_ALTS alt_bco'
                   | otherwise
                   = let unlifted_rep =
@@ -1619,7 +1633,7 @@ generateCCall d0 s p (CCallSpec target cconv safety) result_ty args_r_to_l
          -- slide and return
          d_after_r_min_s = bytesToWords platform (d_after_r - s)
          wrapup       = mkSlideW (trunc16W r_sizeW) (d_after_r_min_s - r_sizeW)
-                        `snocOL` RETURN_UBX (toArgRep platform r_rep)
+                        `snocOL` RETURN_UNLIFTED (toArgRep platform r_rep)
          --trace (show (arg1_offW, args_offW  ,  (map argRepSizeW a_reps) )) $
      return (
          push_args `appOL`
diff --git a/compiler/GHC/StgToCmm/Closure.hs b/compiler/GHC/StgToCmm/Closure.hs
index c2a936f33a..3ff745a719 100644
--- a/compiler/GHC/StgToCmm/Closure.hs
+++ b/compiler/GHC/StgToCmm/Closure.hs
@@ -318,6 +318,10 @@ type DynTag = Int       -- The tag on a *pointer*
 -- x86-32 and 3 bits on x86-64.
 --
 -- Also see Note [Tagging big families] in GHC.StgToCmm.Expr
+--
+-- The interpreter also needs to be updated if we change the
+-- tagging strategy. See Note [Data constructor dynamic tags] in
+-- rts/Interpreter.c
 
 isSmallFamily :: Platform -> Int -> Bool
 isSmallFamily platform fam_size = fam_size <= mAX_PTR_TAG platform
diff --git a/rts/Interpreter.c b/rts/Interpreter.c
index d6478a0164..8c90678bb7 100644
--- a/rts/Interpreter.c
+++ b/rts/Interpreter.c
@@ -361,11 +361,11 @@ interpretBCO (Capability* cap)
     // ------------------------------------------------------------------------
     // Case 3:
     //
-    //       We have an unboxed value to return.  See comment before
-    //       do_return_unboxed, below.
+    //       We have an unlifted value to return.  See comment before
+    //       do_return_lifted, below.
     //
     else {
-        goto do_return_unboxed;
+        goto do_return_unlifted;
     }
 
     // Evaluate the object on top of the stack.
@@ -650,7 +650,7 @@ do_return:
     }
 
     // -------------------------------------------------------------------------
-    // Returning an unboxed value.  The stack looks like this:
+    // Returning an unlifted value.  The stack looks like this:
     //
     //    |     ....      |
     //    +---------------+
@@ -672,7 +672,7 @@ do_return:
     // We're only interested in the case when the real return address
     // is a BCO; otherwise we'll return to the scheduler.
 
-do_return_unboxed:
+do_return_unlifted:
     {
         int offset;
 
@@ -688,7 +688,7 @@ do_return_unboxed:
         IF_DEBUG(interpreter,
              debugBelch(
              "\n---------------------------------------------------------------\n");
-             debugBelch("Returning unboxed\n");
+             debugBelch("Returning unlifted\n");
              debugBelch("Sp = %p\n", Sp);
 #if defined(PROFILING)
              fprintCCS(stderr, cap->r.rCCCS);
@@ -710,7 +710,7 @@ do_return_unboxed:
             INTERP_TICK(it_retto_BCO);
             obj = (StgClosure*)SpW(offset+1);
             ASSERT(get_itbl(obj)->type == BCO);
-            goto run_BCO_return_unboxed;
+            goto run_BCO_return_unlifted;
 
         default:
         {
@@ -928,7 +928,7 @@ run_BCO_return:
 
     goto run_BCO;
 
-run_BCO_return_unboxed:
+run_BCO_return_unlifted:
     // Heap check
     if (doYouWantToGC(cap)) {
         RETURN_TO_SCHEDULER(ThreadInterpret, HeapOverflow);
@@ -940,7 +940,7 @@ run_BCO_return_unboxed:
     /*
        Restore the current cost centre stack if a tuple is being returned.
 
-       When a "simple" unboxed value is returned, the cccs is restored with
+       When a "simple" unlifted value is returned, the cccs is restored with
        an stg_restore_cccs frame on the stack, for example:
 
            ...
@@ -1654,7 +1654,7 @@ run_BCO:
             /* Unpack N ptr words from t.o.s constructor */
             int i;
             int n_words = BCO_NEXT;
-            StgClosure* con = (StgClosure*)SpW(0);
+            StgClosure* con = UNTAG_CLOSURE((StgClosure*)SpW(0));
             Sp_subW(n_words);
             for (i = 0; i < n_words; i++) {
                 SpW(i) = (W_)con->payload[i];
@@ -1679,10 +1679,25 @@ run_BCO:
             // No write barrier is needed here as this is a new allocation
             // visible only from our stack
             SET_HDR(con, (StgInfoTable*)BCO_LIT(o_itbl), cap->r.rCCCS);
-            SpW(0) = (W_)con;
+
+            // Note [Data constructor dynamic tags]
+            //
+            // compute the pointer tag for the constructor and tag the pointer
+            //
+            //     - 1..(TAG_MASK-1): for first TAG_MASK-1 constructors
+            //     - TAG_MASK:        look in info table
+            //
+            // Note: we need to update this if we change the tagging strategy
+            //
+            // For full details of the invariants on tagging, see
+            // https://gitlab.haskell.org/ghc/ghc/wikis/commentary/rts/haskell-execution/pointer-tagging
+
+            StgClosure* tagged_con = TAG_CLOSURE(stg_min(TAG_MASK, 1 + GET_TAG(con)), con);
+            SpW(0) = (W_)tagged_con;
+
             IF_DEBUG(interpreter,
                      debugBelch("\tBuilt ");
-                     printObj((StgClosure*)con);
+                     printObj((StgClosure*)tagged_con);
                 );
             goto nextInsn;
         }
@@ -1822,32 +1837,32 @@ run_BCO:
         case bci_RETURN_P:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_p_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_N:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_n_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_F:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_f_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_D:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_d_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_L:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_l_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_V:
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_v_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         case bci_RETURN_T: {
             /* tuple_info and tuple_bco must already be on the stack */
             Sp_subW(1);
             SpW(0) = (W_)&stg_ret_t_info;
-            goto do_return_unboxed;
+            goto do_return_unlifted;
         }
 
         case bci_SWIZZLE: {
diff --git a/testsuite/tests/ghci/should_run/T19628.hs b/testsuite/tests/ghci/should_run/T19628.hs
new file mode 100644
index 0000000000..74891c690f
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/T19628.hs
@@ -0,0 +1,69 @@
+{- test GHCi support for unlifted types -}
+
+{-# LANGUAGE UnliftedDatatypes #-}
+{-# OPTIONS_GHC -fbyte-code #-}
+
+module Main (main) where
+
+import GHC.Exts
+import GHC.Arr
+
+import Data.Kind
+import Control.Exception
+
+import T19628a
+
+x1 :: Int
+x1 = case test of Force _ -> 10
+  where
+    test :: Strict Int
+    test = Force undefined
+
+x2 :: Int
+x2 = case arr of _ -> 15
+    where
+        Array _ _ _ arr = listArray (1::Int, 10) [1..]
+
+x3 :: Int
+x3 = case test of Force2 y z -> y + z
+  where
+    test :: Strict2 Int
+    test = Force2 15 20
+
+x4 :: Int
+x4 = 40
+  where
+    test :: Maybe Int
+    test = undefined
+
+x5 :: Int
+x5 = 45
+  where
+    test :: Strict Int
+    test = undefined
+
+x6 :: Int
+x6 = case test of Force y -> y
+  where
+    test :: Strict Int
+    test = Force undefined
+
+x7 :: Int
+x7 = case addStrict (Force 4) (Force 5) of (Force y) -> y
+
+x8 :: (Int, Int)
+x8 = (y1, y2)
+  where
+    y1 = unStrict (Force 8)
+    y2 = case toStrict 7 of Force z -> z
+
+main :: IO ()
+main = do
+  print x1
+  print x2
+  print x3
+  print x4
+  print x5 `catch` \(e::SomeException) -> putStrLn "x5: exception"
+  print x6 `catch` \(e::SomeException) -> putStrLn "x6: exception"
+  print x7
+  print x8
diff --git a/testsuite/tests/ghci/should_run/T19628.stdout b/testsuite/tests/ghci/should_run/T19628.stdout
new file mode 100644
index 0000000000..9cd54030d2
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/T19628.stdout
@@ -0,0 +1,8 @@
+10
+15
+35
+40
+x5: exception
+x6: exception
+9
+(8,7)
diff --git a/testsuite/tests/ghci/should_run/T19628a.hs b/testsuite/tests/ghci/should_run/T19628a.hs
new file mode 100644
index 0000000000..ccecf4e4f6
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/T19628a.hs
@@ -0,0 +1,31 @@
+{- test GHCi support for unlifted types -}
+
+{-# LANGUAGE UnliftedDatatypes #-}
+{-# OPTIONS_GHC -fobject-code #-}
+
+module T19628a where
+
+import GHC.Exts
+import GHC.Arr
+
+import Data.Kind
+import Control.Exception
+
+-- unlifted but boxed datatypes
+type Strict :: Type -> TYPE ('BoxedRep 'Unlifted)
+data Strict a = Force a
+
+type Strict2 :: Type -> TYPE ('BoxedRep 'Unlifted)
+data Strict2 a = Force2 a a
+
+{-# NOINLINE addStrict #-}
+addStrict :: Strict Int -> Strict Int -> Strict Int
+addStrict (Force x) (Force y) = Force (x+y)
+
+{-# NOINLINE unStrict #-}
+unStrict :: Strict a -> a
+unStrict (Force x) = x
+
+{-# NOINLINE toStrict #-}
+toStrict :: a -> Strict a
+toStrict x = Force x
diff --git a/testsuite/tests/ghci/should_run/all.T b/testsuite/tests/ghci/should_run/all.T
index 9421768546..4b4b151963 100644
--- a/testsuite/tests/ghci/should_run/all.T
+++ b/testsuite/tests/ghci/should_run/all.T
@@ -79,3 +79,5 @@ test('T18594', just_ghci, ghci_script, ['T18594.script'])
 test('T18562', just_ghci, ghci_script, ['T18562.script'])
 test('T19460', just_ghci, ghci_script, ['T19460.script'])
 test('T19733', just_ghci, compile_and_run, [''])
+
+test('T19628', [extra_files(['T19628a.hs']), only_ways(['ghci']) ], compile_and_run, [''])