Comparison of multiplicity polymorphic vs linear monad

LMonadMonomorphic.hs

{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilyDependencies #-}

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE TupleSections #-}
module LMonad where
import GHC.Prim
import qualified Unsafe.Coerce as NonLinear

import qualified GHC.Types as T
import Prelude hiding (return, (>>=), (>>), IO)
import qualified Prelude as P
import qualified System.IO as IO

class LMonad m where
    return :: a ⊸ m a
    (>>=) :: m a ⊸ (a ⊸ m b) ⊸ m b
    (>>) :: Comonoid a => m a ⊸ m b ⊸ m b
    ma >> mb = ma >>= \a -> next (del a) mb
      where
        next :: () ⊸ m b ⊸ m b
        next () mb' = mb'

main :: P.IO ()
main = toIO test

test :: IO ()
test = do
    h <- openFile "LMonadMonomorphic.hs" IO.ReadMode
    Unrestricted content <- readLines h
    fromIO (mapM_ putStrLn content)

readLines :: IO.Handle ⊸ IO [String]
readLines h = go h [] where
  go :: IO.Handle ⊸ [String] -> IO [String]
  go h' ls = do
    (h'', Unrestricted eof) <- hIsEOF h'
    choose eof h''
      (\h''' -> do
        () <- hClose h'''
        return (Unrestricted $ reverse ls))
      (\h''' -> do
        (h'''', Unrestricted ln) <- hReadLine h'''
        go h'''' (ln : ls))

type IO a = IOL (Unrestricted a)
newtype IOL a = IOL (State# RealWorld -> (# State# RealWorld, a #))
runIOL :: IOL a ⊸ State# RealWorld -> (# State# RealWorld, a #)
runIOL (IOL f) w = f w

instance LMonad IOL where
    return a = IOL (\w -> (# w, a #))
    m >>= f = IOL (\w -> matchUnboxedLin w (runIOL m) (\w' a -> runIOL (f a) w'))


 
openFile :: FilePath -> IO.IOMode -> IOL IO.Handle
openFile path mode = fromIOLin (IO.openFile path mode)

hReadLine :: IO.Handle ⊸ IOL (IO.Handle, Unrestricted String)
hReadLine handle = fromIOLin (unsafeCastLinear almost handle)
  where almost h = (h,) . Unrestricted <$> IO.hGetLine h

hIsEOF :: IO.Handle ⊸ IOL (IO.Handle, Unrestricted Bool)
hIsEOF handle = fromIOLin (unsafeCastLinear almost handle)
  where almost h = (h,) . Unrestricted  <$> IO.hIsEOF h

hClose :: IO.Handle ⊸ IOL ()
hClose handle = fromIOLin (unsafeCastLinear IO.hClose handle)

data Unrestricted a where
    Unrestricted :: a -> Unrestricted a


class Comonoid a where
    del :: a ⊸ ()
    dup :: a ⊸ (a, a)
instance Comonoid  () where
    del () = ()
    dup () = ((), ())
instance Comonoid  (Unrestricted a) where
    del (Unrestricted _) = ()
    dup (Unrestricted a) = (Unrestricted a, Unrestricted a)
choose :: Bool -> a ⊸ (a ⊸ b) -> (a ⊸ b) -> b
choose True v f _ = f v
choose False v _ f = f v

fromIOLin :: P.IO a ⊸ IOL (a)
fromIOLin (T.IO f) = IOL f

toIO :: IO a -> P.IO a
toIO (IOL f) = T.IO (\w -> case f w of (# w', Unrestricted a #) -> (# w', a #))

fromIO :: P.IO a -> IOL (Unrestricted a)
fromIO (T.IO f) = IOL (\w ->  case  f w of (# w', a #) -> (# w', Unrestricted a #))

unsafeCoerce :: a ⊸ b
unsafeCoerce = NonLinear.unsafeCoerce NonLinear.unsafeCoerce

unsafeCastLinear :: (a -> b) ⊸ (a ⊸ b)
unsafeCastLinear = unsafeCoerce


matchUnboxedLin :: State# RealWorld -> (State# RealWorld -> (# State# RealWorld, c #)) ⊸ (State# RealWorld -> c ⊸ (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
matchUnboxedLin = almost
  where
    almost :: State# RealWorld -> (State# RealWorld ⊸ (# State# RealWorld, c #)) -> (State# RealWorld ⊸ c ⊸ (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
    almost a f c =  case f a of
        (# l, r #) -> c l r

LMonadPolymorphic.hs

{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilyDependencies #-}

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE TupleSections #-}
module LMonad where

import qualified Unsafe.Coerce as NonLinear
import GHC.Prim

import qualified GHC.Types as T
import Prelude hiding (return, (>>=), (>>), IO)
import qualified Prelude as P
import qualified System.IO as IO

class LMonad m where
    return :: (SingI p) => Arr p a (m p a)
    (>>=) :: (SingI p) => m p a ⊸ Arr p a (m q b) ⊸ m q b
    (>>) :: m 'Omega a ⊸ m q b ⊸ m q b
    a >> b = a >>= \_ -> b

main :: P.IO ()
main = toIO test

test :: IO ()
test = do
    h <- openFile "LMonadPolymorphic.hs" IO.ReadMode
    content <- readLines h
    fromIO (mapM_ putStrLn content)

readLines :: IO.Handle ⊸ IOL 'Omega [String]
readLines h = go h [] where
  go :: IO.Handle ⊸ [String] -> IOL 'Omega [String]
  go h' ls = do
    (h'', eof) <- hIsEOF h'
    choose eof h''
      (\h''' -> do
        () <- hClose h'''
        return (reverse ls))
      (\h''' -> do
        (h'''', ln) <- hReadLine h'''
        go h'''' (ln : ls))

type IO = IOL 'Omega
newtype IOL (p :: Multiplicity) a = IOL (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf p a #))
runIOL :: IOL p a ⊸ State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf p a #)
runIOL (IOL f) w = f w

data Multiplicity = One | Omega
type family Arr (p :: Multiplicity) a b = result | result -> a b where
    Arr 'One a b = a ⊸ b
    Arr 'Omega a b = a -> b

instance LMonad IOL where
    return :: forall p a. SingI p => Arr p a (IOL p a)
    return = caseF sing where
        caseF :: Sing p -> Arr p a (IOL p a)
        caseF SOne = \a -> IOL (\w -> (# w, Linear a #))
        caseF SOmega = \a -> IOL (\w -> (# w, NotLinear a #))

    (>>=) :: forall p q a b. (SingI p) => IOL p a ⊸ Arr p a (IOL q b) ⊸ IOL q b
    (>>=) = caseF sing where
        caseF :: Sing p -> IOL p a ⊸ Arr p a (IOL q b) ⊸ IOL q b
        caseF SOne = \m (f :: a ⊸ IOL q b) -> IOL (\w -> matchUnboxedLin w (runIOL m) (\w' a -> runIOL (f a) w'))
        caseF SOmega = \m f -> IOL (\w -> matchUnboxed w (runIOL m) (\w' a -> runIOL (f a) w'))



openFile :: FilePath -> IO.IOMode -> IOL 'One IO.Handle
openFile path mode = fromIOLin (IO.openFile path mode)

hReadLine :: IO.Handle ⊸ IOL 'One (IO.Handle, String)
hReadLine handle = fromIOLin (unsafeCastLinear almost handle)
  where almost h = (h,) <$> IO.hGetLine h

hIsEOF :: IO.Handle ⊸ IOL 'Omega (IO.Handle, Bool)
hIsEOF handle = fromIO (unsafeCastLinear almost handle)
  where almost h = (h,) <$> IO.hIsEOF h

hClose :: IO.Handle ⊸ IOL 'Omega ()
hClose handle = fromIO (unsafeCastLinear IO.hClose handle)

data Unrestricted a where
    Unrestricted :: a -> Unrestricted a
data UnrestrictedIf p a where
    Linear :: a ⊸ UnrestrictedIf 'One a
    NotLinear :: a -> UnrestrictedIf 'Omega a

choose :: Bool -> a ⊸ (a ⊸ b) -> (a ⊸ b) -> b
choose True v f _ = f v
choose False v _ f = f v

toIO :: IO a -> P.IO a
toIO (IOL f) = T.IO (removeUnrestricted f)

fromIOLin :: P.IO a ⊸ IOL 'One a
fromIOLin (T.IO f) = IOL (addRestricted f)
fromIO :: P.IO a ⊸ IO a
fromIO (T.IO f) = IOL (addUnrestricted f)

-- Guess I will just let the bug handle this for now
unsafeCoerce :: a ⊸ b
unsafeCoerce = NonLinear.unsafeCoerce NonLinear.unsafeCoerce

unsafeCastLinear :: (a -> b) ⊸ (a ⊸ b)
unsafeCastLinear = unsafeCoerce


addRestricted :: (State# RealWorld ⊸ (# State# RealWorld, a #)) ⊸  State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'One a #)
addRestricted = almost
  where
    almost :: (State# RealWorld ⊸ (# State# RealWorld, a #)) -> State# RealWorld -> (# State# RealWorld, UnrestrictedIf 'One a #)
    almost f w = case f w of
      (# w', a #) -> (# w', Linear a #)
addUnrestricted :: (State# RealWorld ⊸ (# State# RealWorld, a #)) ⊸  State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'Omega a #)
addUnrestricted = almost
  where
    almost :: (State# RealWorld ⊸ (# State# RealWorld, a #)) -> State# RealWorld -> (# State# RealWorld, UnrestrictedIf 'Omega a #)
    almost f w = case f w of
      (# w', a #) -> (# w', NotLinear a #)

removeUnrestricted :: (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'Omega a #)) ⊸  State# RealWorld ⊸ (# State# RealWorld, a #)
removeUnrestricted = almost
  where
    almost :: (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'Omega a #)) -> State# RealWorld -> (# State# RealWorld, a #)
    almost f w = case f w of
      (# w', NotLinear a #) -> (# w', a #)

-- Unboxed tuples have to be case matched and case doesn't work for linear types so we
-- cast the function as linear and it seems to work
matchUnboxedLin :: State# RealWorld ⊸ (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'One c #)) ⊸ (State# RealWorld ⊸ c ⊸ (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
matchUnboxedLin = almost
  where
    -- can't be polymorphic over State# Realworld
    almost :: State# RealWorld -> (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'One c #)) -> (State# RealWorld ⊸ c ⊸ (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
    almost a f c =  case f a of
        (# l, Linear r #) -> c l r

matchUnboxed :: State# RealWorld ⊸ (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'Omega c #)) ⊸ (State# RealWorld ⊸ c -> (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
matchUnboxed = almost
  where
    almost :: State# RealWorld -> (State# RealWorld ⊸ (# State# RealWorld, UnrestrictedIf 'Omega c #)) -> (State# RealWorld ⊸ c -> (# State# RealWorld, d #)) ⊸ (# State# RealWorld, d #)
    almost a f c =  case f a of
        (# l, NotLinear r #) -> c l r

data family Sing (a :: k)
data instance Sing (a :: Multiplicity) where
    SOne :: Sing 'One
    SOmega :: Sing 'Omega
class SingI (p :: k) where
    sing :: Sing p
instance SingI 'One where
    sing = SOne
instance SingI 'Omega where
    sing = SOmega