module Lava.Ref
  ( Ref
  , ref
  , deref
  , memoRef

  , TableIO
  , tableIO
  , extendIO
  , findIO
  , memoRefIO

  , TableST
  , tableST
  , extendST
  , findST
  , memoRefST
  )
 where

import Lava.MyST

import System.IO
import System.IO.Unsafe
import Data.IORef

unsafeCoerce :: a -> b
unsafeCoerce :: a -> b
unsafeCoerce a :: a
a = IO b -> b
forall a. IO a -> a
unsafePerformIO (IO b -> b) -> IO b -> b
forall a b. (a -> b) -> a -> b
$
  do IORef a -> a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef a
forall a. IORef a
ref a
a
     IORef b -> IO b
forall a. IORef a -> IO a
readIORef IORef b
forall a. IORef a
ref
 where
  ref :: IORef a
ref = IO (IORef a) -> IORef a
forall a. IO a -> a
unsafePerformIO (IO (IORef a) -> IORef a) -> IO (IORef a) -> IORef a
forall a b. (a -> b) -> a -> b
$ a -> IO (IORef a)
forall a. a -> IO (IORef a)
newIORef a
forall a. HasCallStack => a
undefined
  -- Defined here because Unsafe.Coerce doesn't exist in Hugs.

{-

Warning! One should regard this module as a portable
extension to the Haskell language. It is not Haskell.

-}

{-

Here is how we implement Tables of Refs:

A Table is nothing but a unique tag, of type TableTag.
TableTag can be anything, as long as it is easy
to create new ones, and we can compare them for
equality. (I chose IORef ()).

So how do we store Refs in a Table? We do not
want the Tables keeping track of their Refs
(which would be disastrous when the table
becomes big, and we would not have any garbage
collection).

Instead, every Ref keeps track of the value it
has in each table it is in. This has the advantage
that we have a constant lookup time (if the number of
Tables we are using is small), and we get garbage
collection of table entries for free.

The disadvantage is that, since the types of the
Tables vary, the Ref has no idea what type of
values it is supposed to store. So we use dynamic
types.

A Ref is implemented as follows: it has two pieces
of information. The first one is an updatable
list of entries for each table it is a member in.
Since it is an updatable list, it is an IORef, which
we also use to compare two Refs. The second part is
just the value the Ref is pointing at (this can never
change anyway).

-}

-----------------------------------------------------------------
-- Ref

data Ref a
  = Ref (IORef [(TableTag, Dyn)]) a

instance Eq (Ref a) where
  Ref r1 :: IORef [(TableTag, Dyn)]
r1 _ == :: Ref a -> Ref a -> Bool
== Ref r2 :: IORef [(TableTag, Dyn)]
r2 _ = IORef [(TableTag, Dyn)]
r1 IORef [(TableTag, Dyn)] -> IORef [(TableTag, Dyn)] -> Bool
forall a. Eq a => a -> a -> Bool
== IORef [(TableTag, Dyn)]
r2

instance Show a => Show (Ref a) where
  showsPrec :: Int -> Ref a -> ShowS
showsPrec _ (Ref _ a :: a
a) = Char -> ShowS
showChar '{' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ShowS
forall a. Show a => a -> ShowS
shows a
a ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar '}'

ref :: a -> Ref a
ref :: a -> Ref a
ref a :: a
a = IO (Ref a) -> Ref a
forall a. IO a -> a
unsafePerformIO (IO (Ref a) -> Ref a) -> IO (Ref a) -> Ref a
forall a b. (a -> b) -> a -> b
$
  do IORef [(TableTag, Dyn)]
r <- [(TableTag, Dyn)] -> IO (IORef [(TableTag, Dyn)])
forall a. a -> IO (IORef a)
newIORef []
     Ref a -> IO (Ref a)
forall (m :: * -> *) a. Monad m => a -> m a
return (IORef [(TableTag, Dyn)] -> a -> Ref a
forall a. IORef [(TableTag, Dyn)] -> a -> Ref a
Ref IORef [(TableTag, Dyn)]
r a
a)

deref :: Ref a -> a
deref :: Ref a -> a
deref (Ref _ a :: a
a) = a
a

-----------------------------------------------------------------
-- Table IO

type TableTag
  = IORef ()

newtype TableIO a b
  = TableIO TableTag
 deriving TableIO a b -> TableIO a b -> Bool
(TableIO a b -> TableIO a b -> Bool)
-> (TableIO a b -> TableIO a b -> Bool) -> Eq (TableIO a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b. TableIO a b -> TableIO a b -> Bool
/= :: TableIO a b -> TableIO a b -> Bool
$c/= :: forall a b. TableIO a b -> TableIO a b -> Bool
== :: TableIO a b -> TableIO a b -> Bool
$c== :: forall a b. TableIO a b -> TableIO a b -> Bool
Eq

tableIO :: IO (TableIO a b)
tableIO :: IO (TableIO a b)
tableIO = TableTag -> TableIO a b
forall a b. TableTag -> TableIO a b
TableIO (TableTag -> TableIO a b) -> IO TableTag -> IO (TableIO a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` () -> IO TableTag
forall a. a -> IO (IORef a)
newIORef ()

findIO :: TableIO a b -> Ref a -> IO (Maybe b)
findIO :: TableIO a b -> Ref a -> IO (Maybe b)
findIO (TableIO t :: TableTag
t) (Ref r :: IORef [(TableTag, Dyn)]
r _) =
  do [(TableTag, Dyn)]
list <- IORef [(TableTag, Dyn)] -> IO [(TableTag, Dyn)]
forall a. IORef a -> IO a
readIORef IORef [(TableTag, Dyn)]
r
     Maybe b -> IO (Maybe b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Dyn -> b
forall a. Dyn -> a
fromDyn (Dyn -> b) -> Maybe Dyn -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` TableTag -> [(TableTag, Dyn)] -> Maybe Dyn
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup TableTag
t [(TableTag, Dyn)]
list)

extendIO :: TableIO a b -> Ref a -> b -> IO ()
extendIO :: TableIO a b -> Ref a -> b -> IO ()
extendIO (TableIO t :: TableTag
t) (Ref r :: IORef [(TableTag, Dyn)]
r _) b :: b
b =
  do [(TableTag, Dyn)]
list <- IORef [(TableTag, Dyn)] -> IO [(TableTag, Dyn)]
forall a. IORef a -> IO a
readIORef IORef [(TableTag, Dyn)]
r
     IORef [(TableTag, Dyn)] -> [(TableTag, Dyn)] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [(TableTag, Dyn)]
r ((TableTag
t,b -> Dyn
forall a. a -> Dyn
toDyn b
b) (TableTag, Dyn) -> [(TableTag, Dyn)] -> [(TableTag, Dyn)]
forall a. a -> [a] -> [a]
: ((TableTag, Dyn) -> Bool) -> [(TableTag, Dyn)] -> [(TableTag, Dyn)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((TableTag -> TableTag -> Bool
forall a. Eq a => a -> a -> Bool
/= TableTag
t) (TableTag -> Bool)
-> ((TableTag, Dyn) -> TableTag) -> (TableTag, Dyn) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TableTag, Dyn) -> TableTag
forall a b. (a, b) -> a
fst) [(TableTag, Dyn)]
list)

-----------------------------------------------------------------
-- Table ST

newtype TableST s a b
  = TableST (TableIO a b)
 deriving TableST s a b -> TableST s a b -> Bool
(TableST s a b -> TableST s a b -> Bool)
-> (TableST s a b -> TableST s a b -> Bool) -> Eq (TableST s a b)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall s a b. TableST s a b -> TableST s a b -> Bool
/= :: TableST s a b -> TableST s a b -> Bool
$c/= :: forall s a b. TableST s a b -> TableST s a b -> Bool
== :: TableST s a b -> TableST s a b -> Bool
$c== :: forall s a b. TableST s a b -> TableST s a b -> Bool
Eq

tableST :: ST s (TableST s a b)
tableST :: ST s (TableST s a b)
tableST = IO (TableST s a b) -> ST s (TableST s a b)
forall a s. IO a -> ST s a
unsafeIOtoST (TableIO a b -> TableST s a b
forall s a b. TableIO a b -> TableST s a b
TableST (TableIO a b -> TableST s a b)
-> IO (TableIO a b) -> IO (TableST s a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` IO (TableIO a b)
forall a b. IO (TableIO a b)
tableIO)

findST :: TableST s a b -> Ref a -> ST s (Maybe b)
findST :: TableST s a b -> Ref a -> ST s (Maybe b)
findST (TableST tab :: TableIO a b
tab) r :: Ref a
r = IO (Maybe b) -> ST s (Maybe b)
forall a s. IO a -> ST s a
unsafeIOtoST (TableIO a b -> Ref a -> IO (Maybe b)
forall a b. TableIO a b -> Ref a -> IO (Maybe b)
findIO TableIO a b
tab Ref a
r)

extendST :: TableST s a b -> Ref a -> b -> ST s ()
extendST :: TableST s a b -> Ref a -> b -> ST s ()
extendST (TableST tab :: TableIO a b
tab) r :: Ref a
r b :: b
b = IO () -> ST s ()
forall a s. IO a -> ST s a
unsafeIOtoST (TableIO a b -> Ref a -> b -> IO ()
forall a b. TableIO a b -> Ref a -> b -> IO ()
extendIO TableIO a b
tab Ref a
r b
b)

-----------------------------------------------------------------
-- Memo

memoRef :: (Ref a -> b) -> (Ref a -> b)
memoRef :: (Ref a -> b) -> Ref a -> b
memoRef f :: Ref a -> b
f = IO b -> b
forall a. IO a -> a
unsafePerformIO (IO b -> b) -> (Ref a -> IO b) -> Ref a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ref a -> IO b) -> Ref a -> IO b
forall a b. (Ref a -> IO b) -> Ref a -> IO b
memoRefIO (b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> IO b) -> (Ref a -> b) -> Ref a -> IO b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ref a -> b
f)

memoRefIO :: (Ref a -> IO b) -> (Ref a -> IO b)
memoRefIO :: (Ref a -> IO b) -> Ref a -> IO b
memoRefIO f :: Ref a -> IO b
f = IO (Ref a -> IO b) -> Ref a -> IO b
forall a. IO a -> a
unsafePerformIO (IO (Ref a -> IO b) -> Ref a -> IO b)
-> IO (Ref a -> IO b) -> Ref a -> IO b
forall a b. (a -> b) -> a -> b
$
  do TableIO a b
tab <- IO (TableIO a b)
forall a b. IO (TableIO a b)
tableIO
     let f' :: Ref a -> IO b
f' r :: Ref a
r = do Maybe b
mb <- TableIO a b -> Ref a -> IO (Maybe b)
forall a b. TableIO a b -> Ref a -> IO (Maybe b)
findIO TableIO a b
tab Ref a
r
                   case Maybe b
mb of
                     Just b :: b
b  -> do b -> IO b
forall (m :: * -> *) a. Monad m => a -> m a
return b
b
                     Nothing -> (b -> IO b) -> IO b
forall a. (a -> IO a) -> IO a
fixIO ((b -> IO b) -> IO b) -> (b -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \b :: b
b ->
                                  do TableIO a b -> Ref a -> b -> IO ()
forall a b. TableIO a b -> Ref a -> b -> IO ()
extendIO TableIO a b
tab Ref a
r b
b
                                     Ref a -> IO b
f Ref a
r
     (Ref a -> IO b) -> IO (Ref a -> IO b)
forall (m :: * -> *) a. Monad m => a -> m a
return Ref a -> IO b
f'

memoRefST :: (Ref a -> ST s b) -> (Ref a -> ST s b)
memoRefST :: (Ref a -> ST s b) -> Ref a -> ST s b
memoRefST f :: Ref a -> ST s b
f = ST s (Ref a -> ST s b) -> Ref a -> ST s b
forall s a. ST s a -> a
unsafePerformST (ST s (Ref a -> ST s b) -> Ref a -> ST s b)
-> ST s (Ref a -> ST s b) -> Ref a -> ST s b
forall a b. (a -> b) -> a -> b
$
  do TableST s a b
tab <- ST s (TableST s a b)
forall s a b. ST s (TableST s a b)
tableST
     let f' :: Ref a -> ST s b
f' r :: Ref a
r = do Maybe b
mb <- TableST s a b -> Ref a -> ST s (Maybe b)
forall s a b. TableST s a b -> Ref a -> ST s (Maybe b)
findST TableST s a b
tab Ref a
r
                   case Maybe b
mb of
                     Just b :: b
b  -> do b -> ST s b
forall (m :: * -> *) a. Monad m => a -> m a
return b
b
                     Nothing -> (b -> ST s b) -> ST s b
forall a s. (a -> ST s a) -> ST s a
fixST ((b -> ST s b) -> ST s b) -> (b -> ST s b) -> ST s b
forall a b. (a -> b) -> a -> b
$ \b :: b
b ->
                                  do TableST s a b -> Ref a -> b -> ST s ()
forall s a b. TableST s a b -> Ref a -> b -> ST s ()
extendST TableST s a b
tab Ref a
r b
b
                                     Ref a -> ST s b
f Ref a
r
     (Ref a -> ST s b) -> ST s (Ref a -> ST s b)
forall (m :: * -> *) a. Monad m => a -> m a
return Ref a -> ST s b
f'

-----------------------------------------------------------------
-- Dyn

data Dyn
  = Dyn

toDyn :: a -> Dyn
toDyn :: a -> Dyn
toDyn = a -> Dyn
forall a b. a -> b
unsafeCoerce

fromDyn :: Dyn -> a
fromDyn :: Dyn -> a
fromDyn = Dyn -> a
forall a b. a -> b
unsafeCoerce

-----------------------------------------------------------------
-- the end.