Merge pull request #5423 from unisonweb/cp/fix-segfault

Fix interpreter segfault

Author: dolio

unison-runtime/src/Unison/Runtime/ANF.hs

@@ -97,7 +97,6 @@ import Data.Bits (shiftL, shiftR, (.&.), (.|.))
 import Data.Functor.Compose (Compose (..))
 import Data.List hiding (and, or)
 import Data.Map qualified as Map
-import Data.Primitive qualified as PA
 import Data.Set qualified as Set
 import Data.Text qualified as Data.Text
 import GHC.Stack (CallStack, callStack)
@@ -112,6 +111,7 @@ import Unison.Pattern qualified as P
 import Unison.Prelude
 import Unison.Reference (Id, Reference, Reference' (Builtin, DerivedId))
 import Unison.Referent (Referent, pattern Con, pattern Ref)
+import Unison.Runtime.Array qualified as PA
 import Unison.Symbol (Symbol)
 import Unison.Term hiding (List, Ref, Text, float, fresh, resolve)
 import Unison.Type qualified as Ty

unison-runtime/src/Unison/Runtime/Array.hs

@@ -141,44 +141,67 @@ checkIBArray name a f arr i
 checkIMBArray
   :: CheckCtx
   => Prim a
+  => PrimMonad m
   => String
   -> a
-  -> (MutableByteArray s -> Int -> r)
-  -> MutableByteArray s -> Int -> r
-checkIMBArray name a f arr i
-  | i < 0 || sizeofMutableByteArray arr `quot` sizeOf a <= i
-  = error $ name ++ " unsafe check out of bounds: " ++ show i
-  | otherwise = f arr i
+  -> (MutableByteArray (PrimState m) -> Int ->  m r)
+  -> MutableByteArray (PrimState m) -> Int ->  m r
+checkIMBArray name a f arr i = do
+  sz <- getSizeofMutableByteArray arr
+  if (i < 0 || sz `quot` sizeOf a <= i)
+     then error $ name ++ " unsafe check out of bounds: " ++ show i
+     else f arr i
 {-# inline checkIMBArray #-}
 
+-- check write mutable byte array
+checkWMBArray
+  :: CheckCtx
+  => Prim a
+  => PrimMonad m
+  => String
+  -> (MutableByteArray (PrimState m) -> Int -> a -> m r)
+  -> MutableByteArray (PrimState m) -> Int -> a -> m r
+checkWMBArray name f arr i a = do
+  sz <- getSizeofMutableByteArray arr
+  if (i < 0 || sz `quot` sizeOf a <= i)
+     then error $ name ++ " unsafe check out of bounds: " ++ show i
+     else f arr i a
+{-# inline checkWMBArray #-}
+
+
 -- check copy byte array
 checkCBArray
   :: CheckCtx
+  => PrimMonad m
   => String
-  -> (MBA s -> Int -> BA -> Int -> Int -> r)
-  -> MBA s -> Int -> BA -> Int -> Int -> r
-checkCBArray name f dst d src s l
-  | d < 0
-  || s < 0
-  || sizeofMutableByteArray dst < d + l
-  || sizeofByteArray src < s + l
-  = error $ name ++ " unsafe check out of bounds: " ++ show (d, s, l)
-  | otherwise = f dst d src s l
+  -> (MBA (PrimState m) -> Int -> BA -> Int -> Int -> m r)
+  -> MBA (PrimState m) -> Int -> BA -> Int -> Int -> m r
+checkCBArray name f dst d src s l = do
+  szd <- getSizeofMutableByteArray dst
+  if (d < 0
+      || s < 0
+      || szd < d + l
+      || sizeofByteArray src < s + l
+      ) then error $ name ++ " unsafe check out of bounds: " ++ show (d, s, l)
+      else f dst d src s l
 {-# inline checkCBArray #-}
 
 -- check copy mutable byte array
 checkCMBArray
   :: CheckCtx
+  => PrimMonad m
   => String
-  -> (MBA s -> Int -> MBA s -> Int -> Int -> r)
-  -> MBA s -> Int -> MBA s -> Int -> Int -> r
-checkCMBArray name f dst d src s l
-  | d < 0
-  || s < 0
-  || sizeofMutableByteArray dst < d + l
-  || sizeofMutableByteArray src < s + l
-  = error $ name ++ " unsafe check out of bounds: " ++ show (d, s, l)
-  | otherwise = f dst d src s l
+  -> (MBA (PrimState m) -> Int -> MBA (PrimState m) -> Int -> Int -> m r)
+  -> MBA (PrimState m) -> Int -> MBA (PrimState m) -> Int -> Int -> m r
+checkCMBArray name f dst d src s l = do
+  szd <- getSizeofMutableByteArray dst
+  szs <- getSizeofMutableByteArray src
+  if ( d < 0
+      || s < 0
+      || szd < d + l
+      || szs < s + l
+    ) then error $ name ++ " unsafe check out of bounds: " ++ show (d, s, l)
+      else f dst d src s l
 {-# inline checkCMBArray #-}
 
 -- check index prim array
@@ -197,35 +220,57 @@ checkIPArray name f arr i
 -- check index mutable prim array
 checkIMPArray
   :: CheckCtx
+  => PrimMonad m
   => Prim a
   => String
-  -> (MutablePrimArray s a -> Int -> r)
-  -> MutablePrimArray s a -> Int -> r
-checkIMPArray name f arr i
-  | i < 0 || sizeofMutablePrimArray arr <= i
-  = error $ name ++ " unsafe check out of bounds: " ++ show i
-  | otherwise = f arr i
+  -> (MutablePrimArray (PrimState m) a -> Int -> m r)
+  -> MutablePrimArray (PrimState m) a -> Int -> m r
+checkIMPArray name f arr i = do
+  asz <- getSizeofMutablePrimArray arr
+  if (i < 0 || asz <= i)
+     then error $ name ++ " unsafe check out of bounds: " ++ show i
+     else f arr i
 {-# inline checkIMPArray #-}
 
+-- check write mutable prim array
+checkWMPArray
+  :: CheckCtx
+  => PrimMonad m
+  => Prim a
+  => String
+  -> (MutablePrimArray (PrimState m) a -> Int -> a -> m r)
+  -> MutablePrimArray (PrimState m) a -> Int -> a -> m r
+checkWMPArray name f arr i a = do
+  asz <- getSizeofMutablePrimArray arr
+  if (i < 0 || asz <= i)
+    then error $ name ++ " unsafe check out of bounds: " ++ show i
+    else f arr i a
+{-# inline checkWMPArray #-}
+
+
 #else
 type CheckCtx :: Constraint
 type CheckCtx = ()
 
-checkIMArray, checkIMPArray, checkIPArray :: String -> r -> r
+checkIMArray, checkIMPArray, checkWMPArray, checkIPArray :: String -> r -> r
 checkCArray, checkCMArray, checkRMArray :: String -> r -> r
 checkIMArray _ = id
 checkIMPArray _ = id
+checkWMPArray _ = id
 checkCArray _ = id
 checkCMArray _ = id
 checkRMArray _ = id
 checkIPArray _ = id
 
-checkIBArray, checkIMBArray :: String -> a -> r -> r
+checkIBArray, checkIMBArray:: String -> a -> r -> r
 checkCBArray, checkCMBArray :: String -> r -> r
 checkIBArray _ _ = id
 checkIMBArray _ _ = id
 checkCBArray _ = id
 checkCMBArray _ = id
+
+checkWMBArray :: String -> r -> r
+checkWMBArray _ = id
 #endif
 
 readArray ::
@@ -301,7 +346,7 @@ writeByteArray ::
   Int ->
   a ->
   m ()
-writeByteArray = checkIMBArray @a "writeByteArray" undefined PA.writeByteArray
+writeByteArray = checkWMBArray "writeByteArray" PA.writeByteArray
 {-# INLINE writeByteArray #-}
 
 indexByteArray ::
@@ -368,7 +413,7 @@ writePrimArray ::
   Int ->
   a ->
   m ()
-writePrimArray = checkIMPArray "writePrimArray" PA.writePrimArray
+writePrimArray = checkWMPArray "writePrimArray" PA.writePrimArray
 {-# INLINE writePrimArray #-}
 
 indexPrimArray ::

unison-runtime/src/Unison/Runtime/Foreign.hs

@@ -23,7 +23,6 @@ import Control.Concurrent (MVar, ThreadId)
 import Control.Concurrent.STM (TVar)
 import Crypto.Hash qualified as Hash
 import Data.IORef (IORef)
-import Data.Primitive (ByteArray, MutableArray, MutableByteArray)
 import Data.Tagged (Tagged (..))
 import Data.X509 qualified as X509
 import Network.Socket (Socket)
@@ -35,6 +34,7 @@ import System.Process (ProcessHandle)
 import Unison.Reference (Reference)
 import Unison.Referent (Referent)
 import Unison.Runtime.ANF (Code, Value)
+import Unison.Runtime.Array
 import Unison.Type qualified as Ty
 import Unison.Util.Bytes (Bytes)
 import Unison.Util.Text (Text)
@@ -256,6 +256,7 @@ instance BuiltinForeign FilePath where foreignRef = Tagged Ty.filePathRef
 instance BuiltinForeign TLS.Context where foreignRef = Tagged Ty.tlsRef
 
 instance BuiltinForeign Code where foreignRef = Tagged Ty.codeRef
+
 instance BuiltinForeign Value where foreignRef = Tagged Ty.valueRef
 
 instance BuiltinForeign TimeSpec where foreignRef = Tagged Ty.timeSpecRef

unison-runtime/src/Unison/Runtime/Foreign/Function.hs

@@ -20,8 +20,6 @@ import Data.Atomics (Ticket)
 import Data.Char qualified as Char
 import Data.Foldable (toList)
 import Data.IORef (IORef)
-import Data.Primitive.Array as PA
-import Data.Primitive.ByteArray as PA
 import Data.Sequence qualified as Sq
 import Data.Time.Clock.POSIX (POSIXTime)
 import Data.Word (Word16, Word32, Word64, Word8)
@@ -32,6 +30,7 @@ import System.IO (BufferMode (..), Handle, IOMode, SeekMode)
 import Unison.Builtin.Decls qualified as Ty
 import Unison.Reference (Reference)
 import Unison.Runtime.ANF (Code, Value, internalBug)
+import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.Exception
 import Unison.Runtime.Foreign
 import Unison.Runtime.MCode

unison-runtime/src/Unison/Runtime/MCode.hs

@@ -57,9 +57,6 @@ import Data.Bits (shiftL, shiftR, (.|.))
 import Data.Coerce
 import Data.Functor ((<&>))
 import Data.Map.Strict qualified as M
-import Data.Primitive.ByteArray
-import Data.Primitive.PrimArray
-import Data.Primitive.PrimArray qualified as PA
 import Data.Void (Void, absurd)
 import Data.Word (Word16, Word64)
 import GHC.Stack (HasCallStack)
@@ -91,6 +88,8 @@ import Unison.Runtime.ANF
     pattern TVar,
   )
 import Unison.Runtime.ANF qualified as ANF
+import Unison.Runtime.Array
+import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.Builtin.Types (builtinTypeNumbering)
 import Unison.Util.EnumContainers as EC
 import Unison.Util.Text (Text)

unison-runtime/src/Unison/Runtime/MCode/Serialize.hs

@@ -14,10 +14,10 @@ import Data.Bytes.Get
 import Data.Bytes.Put
 import Data.Bytes.Serial
 import Data.Bytes.VarInt
-import Data.Primitive.PrimArray
 import Data.Void (Void)
 import Data.Word (Word64)
 import GHC.Exts (IsList (..))
+import Unison.Runtime.Array (PrimArray)
 import Unison.Runtime.MCode hiding (MatchT)
 import Unison.Runtime.Serialize
 import Unison.Util.Text qualified as Util.Text

unison-runtime/src/Unison/Runtime/Serialize.hs

@@ -14,7 +14,6 @@ import Data.Bytes.VarInt
 import Data.Foldable (traverse_)
 import Data.Int (Int64)
 import Data.Map.Strict as Map (Map, fromList, toList)
-import Data.Primitive qualified as PA
 import Data.Text (Text)
 import Data.Text.Encoding (decodeUtf8, encodeUtf8)
 import Data.Vector.Primitive qualified as BA
@@ -26,6 +25,7 @@ import Unison.Hash (Hash)
 import Unison.Hash qualified as Hash
 import Unison.Reference (Id' (..), Reference, Reference' (Builtin, DerivedId), pattern Derived)
 import Unison.Referent (Referent, pattern Con, pattern Ref)
+import Unison.Runtime.Array qualified as PA
 import Unison.Runtime.Exception
 import Unison.Runtime.MCode
   ( BPrim1 (..),

unison-runtime/src/Unison/Runtime/Stack.hs

@@ -96,6 +96,7 @@ module Unison.Runtime.Stack
 where
 
 import Control.Monad.Primitive
+import Data.Primitive (sizeOf)
 import Data.Word
 import GHC.Exts as L (IsList (..))
 import Unison.Prelude
@@ -246,7 +247,7 @@ splitData = \case
 ints :: ByteArray -> [Int]
 ints ba = fmap (indexByteArray ba) [n - 1, n - 2 .. 0]
   where
-    n = sizeofByteArray ba `div` 8
+    n = sizeofByteArray ba `div` intSize
 
 -- | Converts a list of integers representing an unboxed segment back into the
 -- appropriate segment. Segments are stored backwards in the runtime, so this
@@ -348,11 +349,14 @@ type UA = MutableByteArray (PrimState IO)
 
 type BA = MutableArray (PrimState IO) Closure
 
+intSize :: Int
+intSize = sizeOf (0 :: Int)
+
 words :: Int -> Int
-words n = n `div` 8
+words n = n `div` intSize
 
 bytes :: Int -> Int
-bytes n = n * 8
+bytes n = n * intSize
 
 type Arrs = (UA, BA)
 
@@ -666,7 +670,7 @@ augSeg mode (Stack ap fp sp ustk bstk) (useg, bseg) margs = do
       cop <- newByteArray $ ssz + upsz + asz
       copyByteArray cop soff useg 0 ssz
       copyMutableByteArray cop 0 ustk (bytes $ ap + 1) upsz
-      for_ margs $ uargOnto ustk sp cop (words poff + upsz - 1)
+      for_ margs $ uargOnto ustk sp cop (words poff + bpsz - 1)
       unsafeFreezeByteArray cop
       where
         ssz = sizeofByteArray useg
@@ -675,9 +679,9 @@ augSeg mode (Stack ap fp sp ustk bstk) (useg, bseg) margs = do
           | otherwise = (0, upsz + asz)
         upsz = bytes bpsz
         asz = case margs of
-          Nothing -> 0
-          Just (Arg1 _) -> 8
-          Just (Arg2 _ _) -> 16
+          Nothing -> bytes 0
+          Just (Arg1 _) -> bytes 1
+          Just (Arg2 _ _) -> bytes 2
           Just (ArgN v) -> bytes $ sizeofPrimArray v
           Just (ArgR _ l) -> bytes l
     boxedSeg = do