{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}

-- | Provides an implementatio of a byte-addressable memory, intended for internal
-- usage in interpreters for the 'LibRISCV.Effects.Operations.Operations' effect.
module LibRISCV.Effects.Operations.Default.Machine.Memory (
    Memory,
    HalfStorage (..),
    WordStorage (..),
    mkMemory,
    memSize,
    loadByte,
    loadHalf,
    loadWord,
    storeByte,
    storeHalf,
    storeWord,
    storeByteString,
    mkWord,
    mkBytes,
) where

import Data.Array.IO (
    MArray (getBounds, newArray_),
    readArray,
    writeArray,
 )
import Data.BitVector (BV, bitVec)
import Data.Bits (Bits (shift, shiftR, (.&.), (.|.)))
import qualified Data.ByteString.Lazy as BSL
import Data.Int ()
import Data.Word (Word16, Word32, Word8)
import LibRISCV

-- | Since the memory is byte-addressable it requires converting values of a
-- larger size to bytes. This type class is responsible for a conversion of
-- 16-bit values (halfs). That is, it converts halfs to bytes (and vice versa)
-- in little endian.
class HalfStorage halfType byteType where
    -- | Convert a list of two bytes to a single half.
    toHalf :: [byteType] -> halfType

    -- | Convert a single half to a list of two bytes.
    halfToBytes :: halfType -> [byteType]

-- | Similar to 'HalfStorage' but handles conversion of 32-bit values (words).
class WordStorage wordType byteType where
    -- | Convert a list of four bytes to a single word.
    toWord :: [byteType] -> wordType

    -- | Convert a single word to a list of four bytes.
    wordToBytes :: wordType -> [byteType]

instance WordStorage Word32 Word8 where
    toWord :: [Word8] -> Word32
toWord = [Word8] -> Word32
mkWord
    wordToBytes :: Word32 -> [Word8]
wordToBytes = Word32 -> [Word8]
mkBytes

instance HalfStorage Word16 Word8 where
    toHalf :: [Word8] -> Word16
toHalf = Word32 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word16) -> ([Word8] -> Word32) -> [Word8] -> Word16
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> Word32
mkWord
    halfToBytes :: Word16 -> [Word8]
halfToBytes = Word32 -> [Word8]
mkBytes (Word32 -> [Word8]) -> (Word16 -> Word32) -> Word16 -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word16 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance WordStorage BV Word8 where
    toWord :: [Word8] -> BV
toWord = (Int -> Word32 -> BV
forall a. Integral a => Int -> a -> BV
bitVec Int
32 :: Word32 -> BV) (Word32 -> BV) -> ([Word8] -> Word32) -> [Word8] -> BV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word32) -> ([Word8] -> Word32) -> [Word8] -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> Word32
mkWord
    wordToBytes :: BV -> [Word8]
wordToBytes = Word32 -> [Word8]
mkBytes (Word32 -> [Word8]) -> (BV -> Word32) -> BV -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BV -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

instance HalfStorage BV Word8 where
    toHalf :: [Word8] -> BV
toHalf = Int -> Word32 -> BV
forall a. Integral a => Int -> a -> BV
bitVec Int
16 (Word32 -> BV) -> ([Word8] -> Word32) -> [Word8] -> BV
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> Word32
mkWord
    halfToBytes :: BV -> [Word8]
halfToBytes = Word32 -> [Word8]
mkBytes (Word32 -> [Word8]) -> (BV -> Word32) -> BV -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BV -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | Converts a list of bytes to a 'Word32' in little endian.
mkWord :: [Word8] -> Word32
mkWord :: [Word8] -> Word32
mkWord [Word8]
bytes =
    (Word32 -> (Word8, Int) -> Word32)
-> Word32 -> [(Word8, Int)] -> Word32
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
        (\Word32
x (Word8
byte, Int
idx) -> (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
byte Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shift` (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8)) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word32
x)
        Word32
0
        ([(Word8, Int)] -> Word32) -> [(Word8, Int)] -> Word32
forall a b. (a -> b) -> a -> b
$ [Word8] -> [Int] -> [(Word8, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Word8]
bytes [Int
0 .. [Word8] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
bytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

-- | Split a 32-bit word into four octets in little endian.
mkBytes :: Word32 -> [Word8]
mkBytes :: Word32 -> [Word8]
mkBytes Word32
w = (Int -> Word8) -> [Int] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
off -> Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word8) -> Word32 -> Word8
forall a b. (a -> b) -> a -> b
$ Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
shiftR Word32
w Int
off Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0xff) [Int]
offs
  where
    offs :: [Int]
offs = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
8) Int
0

------------------------------------------------------------------------

-- | Representation of a byte-addressable memory. The type is parameterized
-- over an array implementation (such as 'Data.Array.IO.IOUArray') and a
-- generic value type (used to represent instruction operands).
data Memory t a = Memory
    { forall (t :: * -> * -> *) a. Memory t a -> Word32
memStart :: Address
    , forall (t :: * -> * -> *) a. Memory t a -> t Word32 a
memBytes :: t Address a
    }

-- | Create a new memory of the given size starting at the given address.
mkMemory :: (MArray t a IO) => Address -> Word32 -> IO (Memory t a)
mkMemory :: forall (t :: * -> * -> *) a.
MArray t a IO =>
Word32 -> Word32 -> IO (Memory t a)
mkMemory Word32
addr Word32
size = (t Word32 a -> Memory t a) -> IO (t Word32 a) -> IO (Memory t a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Word32 -> t Word32 a -> Memory t a
forall (t :: * -> * -> *) a. Word32 -> t Word32 a -> Memory t a
Memory Word32
addr) ((Word32, Word32) -> IO (t Word32 a)
forall i. Ix i => (i, i) -> IO (t i a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_ (Word32
0, Word32
size Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1))

-- Translate global address to a memory-local address.
toMemAddr :: Memory t a -> Address -> Address
toMemAddr :: forall (t :: * -> * -> *) a. Memory t a -> Word32 -> Word32
toMemAddr Memory t a
mem Word32
addr = Word32
addr Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- (Memory t a -> Word32
forall (t :: * -> * -> *) a. Memory t a -> Word32
memStart Memory t a
mem)

-- | Returns the size of the memory in bytes.
memSize :: (MArray t a IO) => Memory t a -> IO Word32
memSize :: forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> IO Word32
memSize = ((Word32, Word32) -> Word32) -> IO (Word32, Word32) -> IO Word32
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
1) (Word32 -> Word32)
-> ((Word32, Word32) -> Word32) -> (Word32, Word32) -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word32, Word32) -> Word32
forall a b. (a, b) -> b
snd) (IO (Word32, Word32) -> IO Word32)
-> (Memory t a -> IO (Word32, Word32)) -> Memory t a -> IO Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t Word32 a -> IO (Word32, Word32)
forall i. Ix i => t i a -> IO (i, i)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds (t Word32 a -> IO (Word32, Word32))
-> (Memory t a -> t Word32 a) -> Memory t a -> IO (Word32, Word32)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Memory t a -> t Word32 a
forall (t :: * -> * -> *) a. Memory t a -> t Word32 a
memBytes

------------------------------------------------------------------------

-- TODO: Only provide load/store, remove all wrapper functions. Could
-- use the 'Size' type from the operations language for this purpose too.

-- | Load a single byte from memory at the given address.
loadByte :: (MArray t a IO) => Memory t a -> Address -> IO a
loadByte :: forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> Word32 -> IO a
loadByte Memory t a
mem = t Word32 a -> Word32 -> IO a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray (Memory t a -> t Word32 a
forall (t :: * -> * -> *) a. Memory t a -> t Word32 a
memBytes Memory t a
mem) (Word32 -> IO a) -> (Word32 -> Word32) -> Word32 -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Memory t a -> Word32 -> Word32
forall (t :: * -> * -> *) a. Memory t a -> Word32 -> Word32
toMemAddr Memory t a
mem

load :: (MArray t a IO) => ([a] -> b) -> Word32 -> Memory t a -> Address -> IO b
load :: forall (t :: * -> * -> *) a b.
MArray t a IO =>
([a] -> b) -> Word32 -> Memory t a -> Word32 -> IO b
load [a] -> b
proc Word32
bytesize Memory t a
mem Word32
addr =
    [a] -> b
proc
        ([a] -> b) -> IO [a] -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word32 -> IO a) -> [Word32] -> IO [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Word32
off -> Memory t a -> Word32 -> IO a
forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> Word32 -> IO a
loadByte Memory t a
mem (Word32 -> IO a) -> Word32 -> IO a
forall a b. (a -> b) -> a -> b
$ Word32
addr Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
off) [Word32
0 .. (Word32
bytesize Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1)]

-- | Load a half (16-bit) from memory at the given address.
loadHalf :: (MArray t a IO, HalfStorage b a) => Memory t a -> Address -> IO b
loadHalf :: forall (t :: * -> * -> *) a b.
(MArray t a IO, HalfStorage b a) =>
Memory t a -> Word32 -> IO b
loadHalf = ([a] -> b) -> Word32 -> Memory t a -> Word32 -> IO b
forall (t :: * -> * -> *) a b.
MArray t a IO =>
([a] -> b) -> Word32 -> Memory t a -> Word32 -> IO b
load [a] -> b
forall halfType byteType.
HalfStorage halfType byteType =>
[byteType] -> halfType
toHalf Word32
2

-- | Load a word (32-bit) from memory at the given address.
loadWord :: (MArray t a IO, WordStorage b a) => Memory t a -> Address -> IO b
loadWord :: forall (t :: * -> * -> *) a b.
(MArray t a IO, WordStorage b a) =>
Memory t a -> Word32 -> IO b
loadWord = ([a] -> b) -> Word32 -> Memory t a -> Word32 -> IO b
forall (t :: * -> * -> *) a b.
MArray t a IO =>
([a] -> b) -> Word32 -> Memory t a -> Word32 -> IO b
load [a] -> b
forall wordType byteType.
WordStorage wordType byteType =>
[byteType] -> wordType
toWord Word32
4

-- | Store a single byte in memory.
storeByte :: (MArray t a IO) => Memory t a -> Address -> a -> IO ()
storeByte :: forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> Word32 -> a -> IO ()
storeByte Memory t a
mem Word32
addr = t Word32 a -> Word32 -> a -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray (Memory t a -> t Word32 a
forall (t :: * -> * -> *) a. Memory t a -> t Word32 a
memBytes Memory t a
mem) (Word32 -> a -> IO ()) -> Word32 -> a -> IO ()
forall a b. (a -> b) -> a -> b
$ Memory t a -> Word32 -> Word32
forall (t :: * -> * -> *) a. Memory t a -> Word32 -> Word32
toMemAddr Memory t a
mem Word32
addr

store :: (MArray t a IO) => (b -> [a]) -> Word32 -> Memory t a -> Address -> b -> IO ()
store :: forall (t :: * -> * -> *) a b.
MArray t a IO =>
(b -> [a]) -> Word32 -> Memory t a -> Word32 -> b -> IO ()
store b -> [a]
proc Word32
bytesize Memory t a
mem Word32
addr =
    ((Word32, a) -> IO ()) -> [(Word32, a)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(Word32
off, a
val) -> Memory t a -> Word32 -> a -> IO ()
forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> Word32 -> a -> IO ()
storeByte Memory t a
mem (Word32
addr Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
off) a
val)
        ([(Word32, a)] -> IO ()) -> (b -> [(Word32, a)]) -> b -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word32] -> [a] -> [(Word32, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Word32
0 .. (Word32
bytesize Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
- Word32
1)]
        ([a] -> [(Word32, a)]) -> (b -> [a]) -> b -> [(Word32, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> [a]
proc

-- | Store a half (16-bit) in memory.
storeHalf :: (MArray t a IO, HalfStorage b a) => Memory t a -> Address -> b -> IO ()
storeHalf :: forall (t :: * -> * -> *) a b.
(MArray t a IO, HalfStorage b a) =>
Memory t a -> Word32 -> b -> IO ()
storeHalf = (b -> [a]) -> Word32 -> Memory t a -> Word32 -> b -> IO ()
forall (t :: * -> * -> *) a b.
MArray t a IO =>
(b -> [a]) -> Word32 -> Memory t a -> Word32 -> b -> IO ()
store b -> [a]
forall halfType byteType.
HalfStorage halfType byteType =>
halfType -> [byteType]
halfToBytes Word32
2

-- | Store a word (32-bit) in memory.
storeWord :: (MArray t a IO, WordStorage b a) => Memory t a -> Address -> b -> IO ()
storeWord :: forall (t :: * -> * -> *) a b.
(MArray t a IO, WordStorage b a) =>
Memory t a -> Word32 -> b -> IO ()
storeWord = (b -> [a]) -> Word32 -> Memory t a -> Word32 -> b -> IO ()
forall (t :: * -> * -> *) a b.
MArray t a IO =>
(b -> [a]) -> Word32 -> Memory t a -> Word32 -> b -> IO ()
store b -> [a]
forall wordType byteType.
WordStorage wordType byteType =>
wordType -> [byteType]
wordToBytes Word32
4

-- | Write a 'BSL.ByteString' to memory in little endian byteorder. Expects a
-- function to convert single bytes ('Word8') to the chosen value
-- representation, a t'Memory', as well the 'Address' where the string should be
-- stored and the 'BSL.ByteString' itself.
storeByteString ::
    (MArray t a IO) =>
    (Word8 -> a) ->
    Memory t a ->
    Address ->
    BSL.ByteString ->
    IO ()
storeByteString :: forall (t :: * -> * -> *) a.
MArray t a IO =>
(Word8 -> a) -> Memory t a -> Word32 -> ByteString -> IO ()
storeByteString Word8 -> a
convert Memory t a
mem Word32
addr ByteString
bs =
    ((Word32, Word8) -> IO ()) -> [(Word32, Word8)] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(Word32
off, Word8
val) -> Memory t a -> Word32 -> a -> IO ()
forall (t :: * -> * -> *) a.
MArray t a IO =>
Memory t a -> Word32 -> a -> IO ()
storeByte Memory t a
mem (Word32
addr Word32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
+ Word32
off) (Word8 -> a
convert Word8
val)) ([(Word32, Word8)] -> IO ()) -> [(Word32, Word8)] -> IO ()
forall a b. (a -> b) -> a -> b
$
        [Word32] -> [Word8] -> [(Word32, Word8)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Word32
0 ..] ([Word8] -> [(Word32, Word8)]) -> [Word8] -> [(Word32, Word8)]
forall a b. (a -> b) -> a -> b
$
            ByteString -> [Word8]
BSL.unpack ByteString
bs