Fix unsound behavior of unlifted datatypes in ghci (#20194)

Previously, directly calling a function that pattern matches on an
unlifted data type which has at least two constructors in GHCi resulted
 in a segfault.

This happened due to unaccounted return frame info table pointer. The fix is
to pop the above mentioned frame info table pointer when unlifted things are
returned. See Note [Popping return frame for unlifted things]

authors: bgamari, nineonine

9 files changed, 195 insertions, 21 deletions

diff --git a/compiler/GHC/StgToByteCode.hs b/compiler/GHC/StgToByteCode.hs
index ab5d0fb5bc..c574327665 100644
--- a/compiler/GHC/StgToByteCode.hs
+++ b/compiler/GHC/StgToByteCode.hs
@@ -215,11 +215,11 @@ mkProtoBCO
    -> name
    -> BCInstrList
    -> Either  [CgStgAlt] (CgStgRhs)
-        -- ^ original expression; for debugging only
-   -> Int
-   -> Word16
-   -> [StgWord]
-   -> Bool      -- True <=> is a return point, rather than a function
+                -- ^ original expression; for debugging only
+   -> Int       -- ^ arity
+   -> Word16    -- ^ bitmap size
+   -> [StgWord] -- ^ bitmap
+   -> Bool      -- ^ True <=> is a return point, rather than a function
    -> [FFIInfo]
    -> ProtoBCO name
 mkProtoBCO platform nm instrs_ordlist origin arity bitmap_size bitmap is_ret ffis
@@ -736,10 +736,10 @@ doTailCall init_d s p fn args = do
    do_pushes init_d args (map (atomRep platform) args)
   where
   do_pushes !d [] reps = do
-        assert (null reps ) return ()
+        assert (null reps) return ()
         (push_fn, sz) <- pushAtom d p (StgVarArg fn)
         platform <- profilePlatform <$> getProfile
-        assert (sz == wordSize platform ) return ()
+        assert (sz == wordSize platform) return ()
         let slide = mkSlideB platform (d - init_d + wordSize platform) (init_d - s)
             enter = if isUnliftedType (idType fn)
                     then RETURN_UNLIFTED P
@@ -817,6 +817,8 @@ doCase d s p scrut bndr alts
           (isUnboxedTupleType bndr_ty || isUnboxedSumType bndr_ty) &&
           length non_void_arg_reps > 1
 
+        unlifted_alg_ty = isUnliftedType bndr_ty && isAlgCase
+
         non_void_arg_reps = non_void (typeArgReps platform bndr_ty)
 
         profiling
@@ -838,11 +840,12 @@ doCase d s p scrut bndr alts
                           not ubx_tuple_frame = 2 * wordSize platform
                         | otherwise = 0
 
-        -- An unlifted value gets an extra info table pushed on top
-        -- when it is returned.
+        -- The size of the return frame info table pointer if one exists
         unlifted_itbl_size_b :: StackDepth
-        unlifted_itbl_size_b | ubx_tuple_frame              = 3 * wordSize platform
-                             | not (isUnliftedType bndr_ty) = 0
+        unlifted_itbl_size_b | ubx_tuple_frame              = wordSize platform
+                             | not (isUnliftedType bndr_ty)
+                             -- See Note [Popping return frame for unlifted things]
+                             || unlifted_alg_ty             = 0
                              | otherwise                    = wordSize platform
 
         (bndr_size, tuple_info, args_offsets)
@@ -877,6 +880,7 @@ doCase d s p scrut bndr alts
         isAlgCase = isAlgType bndr_ty
 
         -- given an alt, return a discr and code for it.
+        codeAlt :: CgStgAlt -> BcM (Discr, BCInstrList)
         codeAlt (DEFAULT, _, rhs)
            = do rhs_code <- schemeE d_alts s p_alts rhs
                 return (NoDiscr, rhs_code)
@@ -924,16 +928,11 @@ doCase d s p scrut bndr alts
                         | (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, unpack `appOL` rhs_code)
+             return (my_discr alt,
+                     unitOL (UNPACK (trunc16W size)) `appOL` rhs_code)
            where
              real_bndrs = filterOut isTyVar bndrs
 
@@ -1003,7 +1002,23 @@ doCase d s p scrut bndr alts
         bitmap = intsToReverseBitmap platform bitmap_size'{-size-} pointers
 
      alt_stuff <- mapM codeAlt alts
-     alt_final <- mkMultiBranch maybe_ncons alt_stuff
+     alt_final0 <- mkMultiBranch maybe_ncons alt_stuff
+     -- Note [Popping return frame for unlifted things]
+     -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+     -- When an unlifted value is returned, a special stg_ret_XXX_info frame will
+     -- be sitting on top of the stack. This mechanism is used to aid in switching
+     -- execution contexts between object code and interpreter. However, mkMultiBranch,
+     -- which produces the bytecode to discriminate the case alternatives, does not
+     -- account for that frame header and does branching based on the top of the stack.
+     -- Therefore, we must compensate for this by popping the frame header (2 words
+     -- for tuples and 1 word for other unlifted things) before passing control to
+     -- the case discrimination continuation. This ensures we are looking at the
+     -- right word and it also saves some stack space. Failing to account for this
+     -- was the cause of #20194.
+     let alt_final
+           | ubx_tuple_frame    = mkSlideW 0 2 `mappend` alt_final0
+           | unlifted_alg_ty    = mkSlideW 0 1 `mappend` alt_final0
+           | otherwise          = alt_final0
 
      let
          alt_bco_name = getName bndr
diff --git a/rts/Interpreter.c b/rts/Interpreter.c
index c911d99367..8c2195b6e9 100644
--- a/rts/Interpreter.c
+++ b/rts/Interpreter.c
@@ -1702,7 +1702,7 @@ run_BCO:
         case bci_TESTLT_P: {
             unsigned int discr  = BCO_NEXT;
             int failto = BCO_GET_LARGE_ARG;
-            StgClosure* con = (StgClosure*)SpW(0);
+            StgClosure* con = UNTAG_CLOSURE((StgClosure*)SpW(0));
             if (GET_TAG(con) >= discr) {
                 bciPtr = failto;
             }
@@ -1712,7 +1712,7 @@ run_BCO:
         case bci_TESTEQ_P: {
             unsigned int discr  = BCO_NEXT;
             int failto = BCO_GET_LARGE_ARG;
-            StgClosure* con = (StgClosure*)SpW(0);
+            StgClosure* con = UNTAG_CLOSURE((StgClosure*)SpW(0));
             if (GET_TAG(con) != discr) {
                 bciPtr = failto;
             }
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/ByteCode.hs b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/ByteCode.hs
new file mode 100644
index 0000000000..44fe504bde
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/ByteCode.hs
@@ -0,0 +1,8 @@
+{-# LANGUAGE CPP #-}
+{-# OPTIONS_GHC -fbyte-code #-}
+
+module ByteCode where
+
+import Types
+
+#include "Common.hs-incl"
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Common.hs-incl b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Common.hs-incl
new file mode 100644
index 0000000000..7dc6beb569
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Common.hs-incl
@@ -0,0 +1,37 @@
+showT :: T -> String
+showT T0_1 = "T0_1"
+showT (T1I i) = "T1I " ++ show i
+showT T0_2 = "T0_2"
+showT (T2BI b i) = "T2BI " ++ (if b then show i else "0")
+showT T0_3 = "T0_3"
+showT (T3CIB c i b) = "T3CIB " ++ show [c] ++ " " ++ (if b then show i else "0")
+showT T0_4 = "T0_4"
+
+showT0_1 = showT T0_1
+showT1I = showT (T1I 909)
+showT0_2 = showT T0_2
+showT2BI = showT (T2BI True 808)
+showT0_3 = showT T0_3
+showT3CIB = showT (T3CIB 'X' 707 True)
+showT0_4 = "T0_4"
+
+inc :: T -> T
+inc T0_1 = T0_2
+inc (T1I i) = T1I (i+1)
+inc T0_2 = T0_3
+inc (T2BI b i) = T2BI b (i+1)
+inc T0_3 = T0_4
+inc (T3CIB c i b) = T3CIB c (i+1) b
+inc T0_4 = T0_1
+
+t 1 = T0_1
+t 2 = T1I 999
+t 3 = T0_2
+t 4 = T2BI True 899
+t 5 = T0_3
+t 6 = T3CIB 'X' 799 True
+t _ = T0_4
+
+show_inc :: Int -> (Int -> T) -> String
+show_inc i f = let r = inc (f i)
+               in showT r
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Obj.hs b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Obj.hs
new file mode 100644
index 0000000000..7c4bbf16b1
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Obj.hs
@@ -0,0 +1,8 @@
+{-# LANGUAGE CPP #-}
+{-# OPTIONS_GHC -fobject-code #-}
+
+module Obj where
+
+import Types
+
+#include "Common.hs-incl"
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Types.hs b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Types.hs
new file mode 100644
index 0000000000..81500dd667
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/Types.hs
@@ -0,0 +1,14 @@
+{-# LANGUAGE UnliftedDatatypes, StandaloneKindSignatures #-}
+{-# OPTIONS_GHC -fobject-code #-}
+module Types where
+
+import GHC.Exts
+
+type T :: UnliftedType
+data T = T0_1
+       | T1I Int
+       | T0_2
+       | T2BI Bool Int
+       | T0_3
+       | T3CIB Char Int Bool
+       | T0_4
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.hs b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.hs
new file mode 100644
index 0000000000..e39bf884af
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.hs
@@ -0,0 +1,49 @@
+{-# OPTIONS_GHC -fbyte-code #-}
+
+module Main where
+
+{-
+  Test pattern matching on unlifted data types in ghci
+ -}
+
+import Data.Foldable (forM_)
+
+import qualified Obj      as O
+import qualified ByteCode as B
+import Types
+
+main :: IO ()
+main = do
+    testO O.showT0_1
+    testB B.showT0_1
+    testO O.showT1I
+    testB B.showT1I
+    testO O.showT0_2
+    testB B.showT0_2
+    testO O.showT2BI
+    testB B.showT2BI
+    testO O.showT0_3
+    testB B.showT0_3
+    testO O.showT3CIB
+    testB B.showT3CIB
+    testO O.showT0_4
+    testB B.showT0_4
+
+    -- testing calls between BCO and object code (object code function with an unlifted
+    -- value allocated from bytecode and vice-versa)
+    let a = testX [1..7] O.t B.show_inc
+    let b = testX [1..7] B.t O.show_inc
+    putStrLn "____"
+    print $ a == b
+    putStrLn "____"
+    putStrLn "Obj data Bytecode function"
+    forM_ a putStrLn
+    putStrLn "Bytecode data Object function"
+    forM_ b putStrLn
+
+
+testO v = putStrLn $ "Obj: " ++ v
+testB v = putStrLn $ "Bc: " ++  v
+
+testX :: [Int] -> (Int -> T) -> (Int -> (Int -> T) -> String) -> [String]
+testX is get_T show_inc_T = map (\i -> show_inc_T i get_T) is
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.stdout b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.stdout
new file mode 100644
index 0000000000..ade24383fc
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/UnliftedDataTypeInterp.stdout
@@ -0,0 +1,33 @@
+Obj: T0_1
+Bc: T0_1
+Obj: T1I 909
+Bc: T1I 909
+Obj: T0_2
+Bc: T0_2
+Obj: T2BI 808
+Bc: T2BI 808
+Obj: T0_3
+Bc: T0_3
+Obj: T3CIB "X" 707
+Bc: T3CIB "X" 707
+Obj: T0_4
+Bc: T0_4
+____
+True
+____
+Obj data Bytecode function
+T0_2
+T1I 1000
+T0_3
+T2BI 900
+T0_4
+T3CIB "X" 800
+T0_1
+Bytecode data Object function
+T0_2
+T1I 1000
+T0_3
+T2BI 900
+T0_4
+T3CIB "X" 800
+T0_1
diff --git a/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/unlifteddatatypeinterp.T b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/unlifteddatatypeinterp.T
new file mode 100644
index 0000000000..d31c394e9e
--- /dev/null
+++ b/testsuite/tests/ghci/should_run/UnliftedDataTypeInterp/unlifteddatatypeinterp.T
@@ -0,0 +1,10 @@
+test('UnliftedDataTypeInterp',
+     [ extra_files(['Obj.hs', 'ByteCode.hs', 'Types.hs', 'Common.hs-incl']),
+       req_interp,
+       extra_ways(['ghci']),
+       when(config.have_ext_interp, extra_ways(['ghci', 'ghci-ext'])),
+       when(config.have_ext_interp and config.have_profiling, extra_ways(['ghci', 'ghci-ext', 'ghci-ext-prof']))
+     ],
+     compile_and_run,
+     ['']
+    )