Levity-polymorphic arrays and mutable variables

This patch makes the following types levity-polymorphic in their
last argument:

  - Array# a, SmallArray# a, Weak# b, StablePtr# a, StableName# a

  - MutableArray# s a, SmallMutableArray# s a,
    MutVar# s a, TVar# s a, MVar# s a, IOPort# s a

The corresponding primops are also made levity-polymorphic, e.g.
`newArray#`, `readArray#`, `writeMutVar#`, `writeIOPort#`, etc.

Additionally, exception handling functions such as `catch#`, `raise#`,
`maskAsyncExceptions#`,... are made levity/representation-polymorphic.

Now that Array# and MutableArray# also work with unlifted types,
we can simply re-define ArrayArray# and MutableArrayArray# in terms
of them. This means that ArrayArray# and MutableArrayArray# are no
longer primitive types, but simply unlifted newtypes around Array# and
MutableArrayArray#.

This completes the implementation of the Pointer Rep proposal
  https://github.com/ghc-proposals/ghc-proposals/pull/203

Fixes #20911

-------------------------
Metric Increase:
    T12545
-------------------------

-------------------------
Metric Decrease:
    T12545
-------------------------

5 files changed, 293 insertions, 19 deletions

diff --git a/libraries/base/GHC/ArrayArray.hs b/libraries/base/GHC/ArrayArray.hs
new file mode 100644
index 0000000000..ea84f13edf
--- /dev/null
+++ b/libraries/base/GHC/ArrayArray.hs
@@ -0,0 +1,156 @@
+{-# LANGUAGE Trustworthy #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  GHC.ArrayArray
+-- License     :  see libraries/base/LICENSE
+--
+-- Maintainer  :  cvs-ghc@haskell.org
+-- Stability   :  internal
+-- Portability :  non-portable (GHC Extensions)
+--
+-- Legacy interface for arrays of arrays.
+-- Deprecated, because the 'Array#' type can now store arrays directly.
+-- Consider simply using 'Array#' instead of 'ArrayArray#'.
+--
+-- Use GHC.Exts instead of importing this module directly.
+--
+----------------------------------------------------------------------------
+
+module GHC.ArrayArray
+  ( ArrayArray#(..), MutableArrayArray#(..)
+  , newArrayArray#
+  , unsafeFreezeArrayArray#
+  , sizeofArrayArray#
+  , sizeofMutableArrayArray#
+  , indexByteArrayArray#
+  , indexArrayArrayArray#
+  , readByteArrayArray#
+  , readMutableByteArrayArray#
+  , readArrayArrayArray#
+  , readMutableArrayArrayArray#
+  , writeByteArrayArray#
+  , writeMutableByteArrayArray#
+  , writeArrayArrayArray#
+  , writeMutableArrayArrayArray#
+  , copyArrayArray#
+  , copyMutableArrayArray#
+  , sameArrayArray#
+  , sameMutableArrayArray#
+  )
+  where
+
+import GHC.Prim
+import GHC.Types ( Type, UnliftedType, isTrue# )
+import Unsafe.Coerce ( unsafeCoerce, unsafeCoerceUnlifted )
+default ()
+
+{- **********************************************************************
+*                                                                       *
+*                 Arrays of arrays (legacy interface)                    *
+*                                                                       *
+********************************************************************** -}
+
+type ArrayArray# :: UnliftedType
+newtype ArrayArray# = ArrayArray# (Array# ByteArray#)
+
+type MutableArrayArray# :: Type -> UnliftedType
+newtype MutableArrayArray# s = MutableArrayArray# (MutableArray# s ByteArray#)
+
+-- | Create a new mutable array of arrays with the specified number of elements,
+--   in the specified state thread, with each element recursively referring to the
+--   newly created array.
+newArrayArray# :: Int# -> State# s -> (# State# s, MutableArrayArray# s #)
+newArrayArray# sz s1 =
+  -- Create a placeholder ByteArray to initialise the underlying MutableArray#.
+  case newByteArray# 0# s1 of
+    (# s2, placeholder #) ->
+      -- Create a new MutableArray# holding the placeholder ByteArray# value.
+      case newArray# sz (unsafeCoerceUnlifted placeholder) s2 of
+        (# s3, arr #) ->
+          -- Now update the MutableArray# so that the elements refer back
+          -- to the mutable array itself.
+          case write_array_to_array arr 0# s3 of
+            s4 -> (# s4, MutableArrayArray# (unsafeCoerceUnlifted arr) #)
+
+  where
+    write_array_to_array :: MutableArray# s ByteArray# -> Int# -> State# s -> State# s
+    write_array_to_array _ i s
+      | isTrue# (i >=# sz)
+      = s
+    write_array_to_array arr i s
+      = case writeArray# arr i (unsafeCoerceUnlifted arr) s of
+          s' -> write_array_to_array arr (i +# 1#) s'
+
+-- | Make a mutable array of arrays immutable, without copying.
+unsafeFreezeArrayArray# :: MutableArrayArray# s -> State# s -> (# State# s, ArrayArray# #)
+unsafeFreezeArrayArray# = unsafeCoerce unsafeFreezeArray#
+
+-- | Return the number of elements in the array.
+sizeofArrayArray# :: ArrayArray# -> Int#
+sizeofArrayArray# = unsafeCoerce sizeofArray#
+
+-- | Return the number of elements in the array.
+sizeofMutableArrayArray# :: MutableArrayArray# s -> Int#
+sizeofMutableArrayArray# = unsafeCoerce sizeofMutableArray#
+
+indexByteArrayArray# :: ArrayArray# -> Int# -> ByteArray#
+indexByteArrayArray# = unsafeCoerce indexArray#
+
+indexArrayArrayArray# :: ArrayArray# -> Int# -> ArrayArray#
+indexArrayArrayArray# = unsafeCoerce indexArray#
+
+readByteArrayArray# :: MutableArrayArray# s -> Int# -> State# s -> (# State# s, ByteArray# #)
+readByteArrayArray# = unsafeCoerce readArray#
+
+readMutableByteArrayArray# :: MutableArrayArray# s -> Int# -> State# s -> (# State# s, MutableByteArray# s #)
+readMutableByteArrayArray# = unsafeCoerce readArray#
+
+readArrayArrayArray# :: MutableArrayArray# s -> Int# -> State# s -> (# State# s, ArrayArray# #)
+readArrayArrayArray# = unsafeCoerce readArray#
+
+readMutableArrayArrayArray# :: MutableArrayArray# s -> Int# -> State# s -> (# State# s, MutableArrayArray# s #)
+readMutableArrayArrayArray# = unsafeCoerce readArray#
+
+writeByteArrayArray# :: MutableArrayArray# s -> Int# -> ByteArray# -> State# s -> State# s
+writeByteArrayArray# = unsafeCoerce writeArray#
+
+writeMutableByteArrayArray# :: MutableArrayArray# s -> Int# -> MutableByteArray# s -> State# s -> State# s
+writeMutableByteArrayArray# = unsafeCoerce writeArray#
+
+writeArrayArrayArray# :: MutableArrayArray# s -> Int# -> ArrayArray# -> State# s -> State# s
+writeArrayArrayArray# = unsafeCoerce writeArray#
+
+writeMutableArrayArrayArray# :: MutableArrayArray# s -> Int# -> MutableArrayArray# s -> State# s -> State# s
+writeMutableArrayArrayArray# = unsafeCoerce writeArray#
+
+-- | Copy a range of the 'ArrayArray#' to the specified region in the 'MutableArrayArray#'.
+--   Both arrays must fully contain the specified ranges, but this is not checked.
+--   The two arrays must not be the same array in different states, but this is not checked either.
+copyArrayArray# :: ArrayArray# -> Int# -> MutableArrayArray# s -> Int# -> Int# -> State# s -> State# s
+copyArrayArray# = unsafeCoerce copyArray#
+
+-- | Copy a range of the first MutableArrayArray# to the specified region in the second
+--   MutableArrayArray#.
+--   Both arrays must fully contain the specified ranges, but this is not checked.
+--   The regions are allowed to overlap, although this is only possible when the same
+--   array is provided as both the source and the destination.
+copyMutableArrayArray# :: MutableArrayArray# s -> Int# -> MutableArrayArray# s -> Int# -> Int# -> State# s -> State# s
+copyMutableArrayArray# = unsafeCoerce copyMutableArray#
+
+-- | Compare the underlying pointers of two arrays of arrays.
+sameArrayArray# :: ArrayArray# -> ArrayArray# -> Int#
+sameArrayArray# (ArrayArray# arr1) (ArrayArray# arr2) = reallyUnsafePtrEquality# arr1 arr2
+
+-- | Compare the underlying pointers of two mutable arrays of arrays.
+sameMutableArrayArray# :: MutableArrayArray# s -> MutableArrayArray# s -> Int#
+sameMutableArrayArray# (MutableArrayArray# marr1) (MutableArrayArray# marr2 ) = reallyUnsafePtrEquality# marr1 marr2
diff --git a/libraries/base/GHC/Exts.hs b/libraries/base/GHC/Exts.hs
index c03b397601..62c36cb0d5 100755
--- a/libraries/base/GHC/Exts.hs
+++ b/libraries/base/GHC/Exts.hs
@@ -34,6 +34,9 @@ module GHC.Exts
         -- ** Other primitive types
         module GHC.Types,
 
+        -- ** Legacy interface for arrays of arrays
+        module GHC.ArrayArray,
+
         -- * Primitive operations
 
         module GHC.Prim,
@@ -55,8 +58,6 @@ module GHC.Exts
         sameSmallMutableArray#,
         sameByteArray#,
         sameMutableByteArray#,
-        sameArrayArray#,
-        sameMutableArrayArray#,
         sameMVar#,
         sameMutVar#,
         sameTVar#,
@@ -136,6 +137,7 @@ import GHC.Types
            -- GHC's internal representation of 'TyCon's, for 'Typeable'
          , Module, TrName, TyCon, TypeLitSort, KindRep, KindBndr )
 import qualified GHC.Prim.Ext
+import GHC.ArrayArray
 import GHC.Base hiding ( coerce )
 import GHC.Ptr
 import GHC.Stack
diff --git a/libraries/base/base.cabal b/libraries/base/base.cabal
index 79d481ada6..017b97081d 100644
--- a/libraries/base/base.cabal
+++ b/libraries/base/base.cabal
@@ -187,6 +187,7 @@ Library
         Foreign.StablePtr
         Foreign.Storable
         GHC.Arr
+        GHC.ArrayArray
         GHC.Base
         GHC.Bits
         GHC.ByteOrder
diff --git a/libraries/ghc-prim/GHC/Prim/PtrEq.hs b/libraries/ghc-prim/GHC/Prim/PtrEq.hs
index 5e9d2e564b..5cc3e511e6 100644
--- a/libraries/ghc-prim/GHC/Prim/PtrEq.hs
+++ b/libraries/ghc-prim/GHC/Prim/PtrEq.hs
@@ -26,8 +26,6 @@ module GHC.Prim.PtrEq
     sameSmallMutableArray#,
     sameByteArray#,
     sameMutableByteArray#,
-    sameArrayArray#,
-    sameMutableArrayArray#,
     sameMutVar#,
     sameTVar#,
     sameMVar#,
@@ -100,14 +98,6 @@ sameByteArray# = reallyUnsafePtrEquality#
 sameMutableByteArray# :: MutableByteArray# s -> MutableByteArray# s -> Int#
 sameMutableByteArray# = reallyUnsafePtrEquality#
 
--- | Compare the underlying pointers of two arrays of arrays.
-sameArrayArray# :: ArrayArray# -> ArrayArray# -> Int#
-sameArrayArray# = reallyUnsafePtrEquality#
-
--- | Compare the underlying pointers of two mutable arrays of arrays.
-sameMutableArrayArray# :: MutableArrayArray# s -> MutableArrayArray# s -> Int#
-sameMutableArrayArray# = reallyUnsafePtrEquality#
-
 -- | Compare the underlying pointers of two 'MutVar#'s.
 sameMutVar# :: MutVar# s a -> MutVar# s a -> Int#
 sameMutVar# = reallyUnsafePtrEquality#
diff --git a/libraries/ghc-prim/changelog.md b/libraries/ghc-prim/changelog.md
index 63f2881dcb..372018290b 100644
--- a/libraries/ghc-prim/changelog.md
+++ b/libraries/ghc-prim/changelog.md
@@ -20,25 +20,150 @@
   Note that the explicit type applications are required, as the call to
   `withDict` would be ambiguous otherwise.
 
+- Primitive types and functions which handle boxed values are now levity-polymorphic,
+  meaning that they now also work with unlifted boxed values (i.e. values whose type
+  has kind `TYPE (BoxedRep Unlifted)`).
+
+  The following type constructors are now levity-polymorphic:
+
+    - `Array#`, `SmallArray#`, `Weak#`, `StablePtr#`, `StableName#`,
+
+    - `MutableArray#`, `SmallMutableArray#`, `MutVar#`,
+      `TVar#`, `MVar#`, `IOPort#`.
+
+  For example, `Array#` used to have kind:
+
+  ```
+  Type -> UnliftedType
+  ```
+
+  but it now has kind:
+
+  ```
+  forall {l :: Levity}. TYPE (BoxedRep l) -> UnliftedType
+  ```
+
+  Similarly, `MutVar#` used to have kind:
+
+  ```
+  Type -> Type -> UnliftedType
+  ```
+
+  but it now has kind:
+
+  ```
+  forall {l :: Levity}. Type -> TYPE (BoxedRep l) -> UnliftedType
+  ```
+
+  This means that in `Array# a`, `MutableArray# s a`, `MutVar# s a`, ...,
+  the element type `a`, must always be boxed, but it can now either be lifted
+  or unlifted.
+  In particular, arrays and mutable variables can now be used to store
+  other arrays and mutable variables.
+
+  All functions which use these updated primitive types are also levity-polymorphic:
+
+    - all array operations (reading/writing/copying/...), for both arrays and small arrays,
+      mutable and immutable:
+
+      - `newArray#`, `readArray#`, `writeArray#`, `sizeofArray#`, `sizeofMutableArray#`, `indexArray#`,
+        `unsafeFreezeArray#`, `unsafeThawArray#`, `copyArray#`, `copyMutableArray#`, `cloneArray#`,
+        `cloneMutableArray#`, `freezeArray#`, `thawArray#`, `casArray#`,
+
+      - `newSmallArray#`, `shrinkSmallMutableArray#`, `readSmallArray#`, `writeSmallArray#`, `sizeofSmallArray#`,
+        `getSizeofSmallMutableArray#`, `indexSmallArray#`, `unsafeFreezeSmallArray#`,
+        `unsafeThawSmallArray#`, `copySmallArray#`, `copySmallMutableArray#`, `cloneSmallArray#`,
+        `cloneSmallMutableArray#`, `freezeSmallArray#`, `thawSmallArray#`, `casSmallArray#`,
+
+    - `newMutVar#`,`readMutVar#`,`writeMutV#`,`casMutVar#`,
+
+    - operations on `MVar#` and `TVar#`:
+
+      - `newTVar#`, `readTVar#`, `readTVarIO#`, `writeTVar#`,
+
+      - `newMVar#`, `takeMVar#`, `tryTakeMVar#`, `putMVar#`,
+        `tryPutMVar#`, `readMVar#`, `tryReadMVar#`,
+
+    - `STM` operations `atomically#`, `retry#`, `catchRetry#` and `catchSTM#`.
+
+    - `newIOPort#`, `readIOPort#`, `writeIOPort#`,
+
+    - `mkWeak#`, `mkWeakNoFinalizer#`, `addCFinalizerToWeak#`, `deRefWeak#`, `finalizeWeak#`,
+
+    - `makeStablePtr#`, `deRefStablePtr#`, `eqStablePtr#`, `makeStableName#`, `stableNameToInt#`,
+
+  For example, the full type of `newMutVar#` is now:
+
+  ```
+  newMutVar#
+    :: forall s {l :: Levity} (a :: TYPE (BoxedRep l)).
+       a -> State# s -> (# State# s, MVar# s a #)
+  ```
+
+  and the full type of `writeSmallArray#` is:
+
+  ```
+  writeSmallArray#
+    :: forall s {l :: Levity} (a :: TYPE ('BoxedRep l)).
+       SmallMutableArray# s a -> Int# -> a -> State# s -> State# s
+  ```
+
+- `ArrayArray#` and `MutableArrayArray#` have been moved from `GHC.Prim` to `GHC.Exts`.
+  They are deprecated, because their functionality is now subsumed by `Array#`
+  and `MutableArray#`.
+
 - `mkWeak#`, `mkWeakNoFinalizer#`, `touch#` and `keepAlive#` are now
   levity-polymorphic instead of representation-polymorphic. For instance:
 
   ```
   mkWeakNoFinalizer#
-    :: forall {l :: Levity} (a :: TYPE (BoxedRep l)) (b :: Type)
-    .  a -> b -> State# RealWorld -> (# State# RealWorld, Weak# b #)
+    :: forall {l :: Levity} (a :: TYPE ('BoxedRep l))
+              {k :: Levity} (b :: TYPE ('BoxedRep k)).
+       a -> b -> State# RealWorld -> (# State# RealWorld, Weak# b #)
   ```
 
-  That is, the type signature now quantifies over a variable of type `Levity`
-  instead of `RuntimeRep`. In addition, this variable is now inferred,
+  That is, the type signature now quantifies over the `Levity` of `a`
+  instead of its `RuntimeRep`. In addition, this variable is now inferred,
   instead of specified, meaning that it is no longer eligible for visible type application.
+  Note that `b` is now also levity-polymorphic, due to the change outlined in the
+  previous point.
 
-- The `RuntimeRep` parameter to `raise#` is now inferred:
+- Primitive functions for throwing and catching exceptions are now more polymorphic
+  than before. For example, `catch#` now has type:
 
   ```
-  raise# :: forall (a :: Type) {r :: RuntimeRep} (b :: TYPE r). a -> b
+  catch#
+    :: forall {r :: RuntimeRep} (a :: TYPE r)
+              {l :: Levity} (b :: TYPE ('BoxedRep l)).
+        ( State# RealWorld -> (# State# RealWorld, a #) )
+    -> ( b -> State# RealWorld -> (# State# RealWorld, a #) )
+    -> State# RealWorld -> (# State# RealWorld, a #)
   ```
 
+  The following functions are concerned:
+
+    - `catch#`,
+
+    - `raise#`, `raiseIO#`,
+
+    - `maskAsyncExceptions#`, `maskUninterruptible#`, `unmaskAsyncExceptions#`.
+
+  Note in particular that `raise#` is now both representation-polymorphic
+  (with an inferred `RuntimeRep` argument) and levity-polymorphic, with type:
+
+  ```
+  raise# :: forall {l :: Levity} (a :: TYPE (BoxedRep l))
+                   {r :: RuntimeRep} (b :: TYPE r).
+            a -> b
+  ```
+
+- ``fork#`` and ``forkOn#`` are now representation-polymorphic. For example, ``fork#``
+  now has type: ::
+
+      fork# :: forall {r :: RuntimeRep} (a :: TYPE r).
+               (State# RealWorld -> (# State# RealWorld, a #))
+            -> (State# RealWorld -> (# State# RealWorld, a #))
+
 - `reallyUnsafePtrEquality#` has been made more general, as it is now
    both levity-polymorphic and heterogeneous:
 
@@ -61,7 +186,7 @@
   - `eqStableName#`.
 
 - The following functions have been added to `GHC.Exts`:
- 
+
   ```
   sameArray# :: Array# a -> Array# a -> Int#
   sameSmallArray# :: SmallArray# a -> SmallArray# a -> Int#