linear-resource.hs

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnicodeSyntax #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}

-- Linear IO, in `IO a` - `a` have to have 1 weight.
import SIO
import Data.IORef
import Linear.Common
import Linear.Std
import Linear.Unsafe
import Control.Exception as Exception
import Prelude (fromInteger, (++), show, Int, IO)
import qualified Prelude as P
import GHC.Exts
import qualified System.IO as IO
import qualified Unsafe.Coerce as Unsafe

-- Special wrapper.
data W1 a s   = W1 a

-- Key that is used for resource unregistering
data Key a s  = Key Int deriving (P.Eq)
-- Key value pair that introduces `s` that guarantees
-- that value can be used only with it's key.
data KV a = forall s . KV (W1 a s) (Key a s)
-- Information to be kept in state:
-- function to clear resource, and corresponding key
data KeyDisp  = forall a s . KeyDisp  (IO' ()) (Key a s)

-- These helper types may be complex, there another approach
-- that may allow to hide key in the wrapper for resource,
-- but that approach would require to wrap all methods working
-- with resource. However it's possible that it should be done
-- anyway as methods would have to return linear resource alongside
-- the result

-- Storage, can be part of the monad where action is running.
data S = S (IORef [KeyDisp])

-- Helper to use "ordinary" method on wrapper resource.
useW1 :: W1 a s ->. (a ->. IO' (a, r)) ->. IO' (W1 a s, r)
useW1 (W1 x) f = do
  (v, r) <- f x
  return (W1 v, r)

-- Function to free resource allocated in with a gived storage
freeRes :: S -> W1 a s ->. Key a s ->. IO' ()
freeRes (S s) w k = unsafeIOtoLIO (unsafeFree (unsafeUnrestricted w) (unsafeUnrestricted k)) where
  unsafeFree :: (Unrestricted (W1 a s)) ->. (Unrestricted (Key a s)) ->. IO ()
  unsafeFree (Unrestricted w') (Unrestricted k') =
      readIORef s P.>>= \keyvals -> go keyvals P.>>= writeIORef s P.>> P.return ()
    where
      go :: [KeyDisp] -> IO [KeyDisp]
      go ((KeyDisp f z):rest) = case (unsafeCoerce z P.== k') of
                                 P.True -> unsafeLIOtoIO f P.>> P.return rest
                                 P.False -> go rest P.>>= \rest' -> P.return ((KeyDisp f z):rest')

-- Function to allocate resouce in a gived storage
alloc :: S -> (IO a) -> (a ->. IO' ()) -> IO' (KV a)
alloc (S r) create dis = unsafeIOtoLIO go
  where
  go = readIORef r P.>>= \ls ->
         let c = case ls of
                  [] -> 0
                  (KeyDisp _ (Key x):_) -> (x P.+ 1)
             key = Key c
          in create P.>>= \value -> modifyIORef r (KeyDisp (dis value) key:) P.>> P.return (KV (W1 value) key)


-- Runner
runLIO1 :: (S -> IO' (Unrestricted r)) -> IO r
runLIO1 f =
    bracket (P.fmap S (newIORef []))
            (finalize)
            (\s -> unsafeLIOtoIO (f s) P.>>= \(Unrestricted r) -> P.return r)
  where
    finalize (S s) = readIORef s P.>>= \vals -> go vals
    go [] = P.return ()
    go (KeyDisp x _:xs) = (unsafeLIOtoIO x P.>> P.return ()) `Exception.finally` go xs

-- Test case
test = runLIO1 $ \s -> do
  (KV value key) <- alloc s (IO.openFile "foo" IO.ReadWriteMode) (hClose)
  (value', Unrestricted _) <- useW1 value hGetLine
  () <- freeRes s value' key
  return (Unrestricted ())

-- Helpers.

hOpen :: P.String -> IO' IO.Handle
hOpen str = unsafeIOtoLIO (IO.openFile str IO.ReadWriteMode)

hClose :: IO.Handle ->. IO' ()
hClose h = unsafeIOtoLIO (unsafeClose (unsafeUnrestricted h)) where
  unsafeClose :: Unrestricted IO.Handle ->. IO ()
  unsafeClose (Unrestricted p) = IO.hClose p P.>> P.return ()

hGetLine :: IO.Handle ->. IO' (IO.Handle, Unrestricted P.String)
hGetLine h = unsafeIOtoLIO (unsafeGetLine (unsafeUnrestricted h)) where
  unsafeGetLine :: Unrestricted IO.Handle ->. IO (IO.Handle, Unrestricted P.String)
  unsafeGetLine (Unrestricted p) = IO.hGetLine p P.>>= \s -> P.return (p, Unrestricted s)