{-# LANGUAGE UnboxedTuples #-}
-- |
-- Copyright: (c) 2021 Xy Ren
-- License: BSD3
-- Maintainer: xy.r@outlook.com
-- Stability: unstable
-- Portability: non-portable (GHC only)
--
-- 'Mem' is a data structure that is a simulation of an array of thread-local pointers. This structure supports:
--
-- * \( O(n) \) creation of a new pointer;
-- * \( O(n) \) changing the pointer in an array cell;
-- * \( O(1) \) modification of the memory a pointer points to;
-- * \( O(1) \) read.
--
-- __This is an /internal/ module and its API may change even between minor versions.__ Therefore you should be
-- extra careful if you're to depend on this module.
module Cleff.Internal.Data.Mem (Mem, MemPtr, empty, adjust, alloca, read, write, replace, append, update) where

import           Cleff.Internal.Data.Any (Any, fromAny, toAny)
import           Data.IntMap.Strict      (IntMap)
import qualified Data.IntMap.Strict      as Map
import           Data.Kind               (Type)
import           Data.Rec.SmallArray     (Rec, pattern (:~:))
import qualified Data.Rec.SmallArray     as Rec
import           Prelude                 hiding (read)

-- | The representation of a pointer in a 'Mem'.
type role MemPtr representational nominal
newtype MemPtr (f :: k -> Type) (a :: k) = MemPtr { MemPtr f a -> Int
unMemPtr :: Int }
  deriving newtype
    ( Eq  -- ^ Pointer equality.
    , Ord -- ^ An arbitrary total order on the pointers.
    )

-- | A simulated array of thread-local pointers. This means for each array cell, you can either change the pointer or
-- change the memory the pointer points to.
--
-- Note that like real memory, any of the operations provided is not generally safe and it is your responsibility to
-- ensure the correctness of your calls.
type role Mem representational nominal
data Mem (f :: k -> Type) (es :: [k]) = Mem
  {-# UNPACK #-} !(Rec (MemPtr f) es) -- ^ The array.
  {-# UNPACK #-} !Int -- ^ The next memory address to allocate.
  !(IntMap Any) -- ^ The simulated memory.

-- | Create a 'Mem' with no pointers.
empty :: Mem f '[]
empty :: Mem f '[]
empty = Rec (MemPtr f) '[] -> Int -> IntMap Any -> Mem f '[]
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem Rec (MemPtr f) '[]
forall k (f :: k -> Type). Rec f '[]
Rec.empty Int
0 IntMap Any
forall a. IntMap a
Map.empty
{-# INLINE empty #-}

-- | Adjust the array of pointers.
adjust :: ∀ es' es f. (Rec (MemPtr f) es -> Rec (MemPtr f) es') -> Mem f es -> Mem f es'
adjust :: (Rec (MemPtr f) es -> Rec (MemPtr f) es') -> Mem f es -> Mem f es'
adjust Rec (MemPtr f) es -> Rec (MemPtr f) es'
f (Mem Rec (MemPtr f) es
re Int
n IntMap Any
mem) = Rec (MemPtr f) es' -> Int -> IntMap Any -> Mem f es'
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem (Rec (MemPtr f) es -> Rec (MemPtr f) es'
f Rec (MemPtr f) es
re) Int
n IntMap Any
mem
{-# INLINE adjust #-}

-- | Allocate a new address. \( O(1) \).
alloca :: ∀ e es f. Mem f es -> (# MemPtr f e, Mem f es #)
alloca :: Mem f es -> (# MemPtr f e, Mem f es #)
alloca (Mem Rec (MemPtr f) es
re Int
n IntMap Any
mem) = (# Int -> MemPtr f e
forall k (f :: k -> Type) (a :: k). Int -> MemPtr f a
MemPtr Int
n, Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem Rec (MemPtr f) es
re (Int -> Int
forall a. Enum a => a -> a
succ Int
n) IntMap Any
mem #)
{-# INLINE alloca #-}

-- | Read a pointer. \( O(1) \).
read :: ∀ e es f. Rec.Elem e es => Mem f es -> f e
read :: Mem f es -> f e
read (Mem Rec (MemPtr f) es
re Int
_ IntMap Any
mem) = Any -> f e
forall a. Any -> a
fromAny (Any -> f e) -> Any -> f e
forall a b. (a -> b) -> a -> b
$ IntMap Any
mem IntMap Any -> Int -> Any
forall a. IntMap a -> Int -> a
Map.! MemPtr f e -> Int
forall k (f :: k -> Type) (a :: k). MemPtr f a -> Int
unMemPtr (Rec (MemPtr f) es -> MemPtr f e
forall k (e :: k) (es :: [k]) (f :: k -> Type).
Elem e es =>
Rec f es -> f e
Rec.index @e Rec (MemPtr f) es
re)
{-# INLINE read #-}

-- | Write to the memory a pointer points to. \( O(1) \).
write :: ∀ e es f. MemPtr f e -> f e -> Mem f es -> Mem f es
write :: MemPtr f e -> f e -> Mem f es -> Mem f es
write (MemPtr Int
m) f e
x (Mem Rec (MemPtr f) es
re Int
n IntMap Any
mem) = Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem Rec (MemPtr f) es
re Int
n (Int -> Any -> IntMap Any -> IntMap Any
forall a. Int -> a -> IntMap a -> IntMap a
Map.insert Int
m (f e -> Any
forall a. a -> Any
toAny f e
x) IntMap Any
mem)
{-# INLINE write #-}

-- | Replace a pointer with a new one. \( O(n) \).
replace :: ∀ e es f. Rec.Elem e es => MemPtr f e -> f e -> Mem f es -> Mem f es
replace :: MemPtr f e -> f e -> Mem f es -> Mem f es
replace (MemPtr Int
m) f e
x (Mem Rec (MemPtr f) es
re Int
n IntMap Any
mem) = Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem (MemPtr f e -> Rec (MemPtr f) es -> Rec (MemPtr f) es
forall k (e :: k) (es :: [k]) (f :: k -> Type).
Elem e es =>
f e -> Rec f es -> Rec f es
Rec.update @e (Int -> MemPtr f e
forall k (f :: k -> Type) (a :: k). Int -> MemPtr f a
MemPtr Int
m) Rec (MemPtr f) es
re) Int
n (Int -> Any -> IntMap Any -> IntMap Any
forall a. Int -> a -> IntMap a -> IntMap a
Map.insert Int
m (f e -> Any
forall a. a -> Any
toAny f e
x) IntMap Any
mem)
{-# INLINE replace #-}

-- | Add a new pointer to the array. \( O(n) \).
append :: ∀ e es f. MemPtr f e -> f e -> Mem f es -> Mem f (e ': es)
append :: MemPtr f e -> f e -> Mem f es -> Mem f (e : es)
append (MemPtr Int
m) f e
x (Mem Rec (MemPtr f) es
re Int
n IntMap Any
mem) = Rec (MemPtr f) (e : es) -> Int -> IntMap Any -> Mem f (e : es)
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem (Int -> MemPtr f e
forall k (f :: k -> Type) (a :: k). Int -> MemPtr f a
MemPtr Int
m MemPtr f e -> Rec (MemPtr f) es -> Rec (MemPtr f) (e : es)
forall a (f :: a -> Type) (e :: a) (es :: [a]).
f e -> Rec f es -> Rec f (e : es)
:~: Rec (MemPtr f) es
re) Int
n (Int -> Any -> IntMap Any -> IntMap Any
forall a. Int -> a -> IntMap a -> IntMap a
Map.insert Int
m (f e -> Any
forall a. a -> Any
toAny f e
x) IntMap Any
mem)
{-# INLINE append #-}

-- | Use the memory of LHS as a newer version for the memory of RHS. \( O(1) \).
update :: ∀ es es' f. Mem f es' -> Mem f es -> Mem f es
update :: Mem f es' -> Mem f es -> Mem f es
update (Mem Rec (MemPtr f) es'
_ Int
n IntMap Any
mem) (Mem Rec (MemPtr f) es
re' Int
_ IntMap Any
_) = Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
forall k (f :: k -> Type) (es :: [k]).
Rec (MemPtr f) es -> Int -> IntMap Any -> Mem f es
Mem Rec (MemPtr f) es
re' Int
n IntMap Any
mem
{-# INLINE update #-}