{-# LANGUAGE Rank2Types, ScopedTypeVariables, MagicHash, UnboxedTuples, FlexibleInstances #-}
-- | A linear type-based I/O system a la Clean.
--
--   This is an alternative to composing monads - one can decompose them into their
--   corresponding comonads, with linear operations for manipulating them.
--   (See Kieburtz, http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.46.5169&rep=rep1&type=pdf)
module Control.Linear (St, A, Blank, Pair, Fn, Open, Semiclosed, Freeable(Freeable), Foreign(Foreign), Placeholder(Placeholder), Pointer, (>>==), rtn, run, bimap, assoc1, assoc2, drop1, drop2, undrop1, undrop2, swap, apply, curry, distr, void', bimap',
open, getStdin, getStdout, getStderr, close, close1, fileSize, setFileSize, eof, seek, tell, char, line, lookahead, contents, putC, putS, random, frst, secnd, Fix(In), fixInj1, fixInj2, newForeign, new, get, weakening, contraction, free, focus, peek', poke', changeType, peek1, poke1, split) where

import Control.Arrow
import Control.Category
import Control.Monad
import GHC.Prim
import GHC.IO
import GHC.Base (realWorld#)
import System.IO
import Foreign.ForeignPtr hiding (unsafeForeignPtrToPtr)
import Foreign.ForeignPtr.Unsafe
import Foreign.Ptr
import qualified Foreign.Marshal.Alloc as A
import Foreign.Storable
import Data.Default
import Prelude hiding (id, (.), curry)
import System.Random (getStdGen)
import System.IO.Unsafe

data St = St (State# RealWorld)

-- | Linear computations are arrows over linear data, but monads over nonlinear data.
newtype A t u v = A (u -> (v, t))

type Fn t u = A () (Pair t St) (Pair u St)

data Blank = Blank

data Pair t u = Pair !t !u

newtype Open = Open Handle

newtype Semiclosed = Semiclosed Handle

data Pointer p s t = Pointer !(ForeignPtr Blank) !(Ptr t) (State# RealWorld)

-- | Pointers can be freeable or foreign. Freeable pointers are created with 'new'.
--   Linearity is enforced for them, so I can do strong update, but the need
--   to keep track of all pointers means that 'split' cannot be supported.
--   Foreign pointers on the other hand can be copied indefinitely since the
--   garbage collector is keeping track of them.
--
--   Freeable pointers can be turned into foreign pointers permanently using
--   'newForeign', or temporarily by focusing.
--
--   Placeholders classify pointers that either point to junk or to data that
--   is not allowed to be used (to maintain linearity).
data Freeable = Freeable

data Foreign = Foreign

data Placeholder = Placeholder

{-# NOINLINE dummy #-}
dummy :: ForeignPtr Blank
dummy = unsafePerformIO (A.malloc >>= newForeignPtr_)

instance (Default a) => Category (A a) where
	id = A (\x -> (x, def))
	A f . A g = A (f . fst . g)

instance (Default a) => Arrow (A a) where
	arr f = A (\x -> (f x, def))
	first (A f) = A (\(x, y) -> let (z, a) = f x in ((z, y), a))

instance (Default a) => ArrowChoice (A a) where
	A f +++ A g = A (\ei -> either (\x -> let (y, z) = f x in (Left y, z))
		(\x -> let (y, z) = g x in (Right y, z))
		ei)
	left a = a +++ id

infixl 1 >>==

-- | Monadic bind (for nonlinear data).
A f >>== g = A (\x -> let
	(y, z) = f x
	A h = g z in
	h y)

-- | Monadic return
rtn x = A (\y -> (y, x))

--- Algebraic operations
-- This setup is from http://cs.ioc.ee/~tarmo/tsem11/jeltsch1602-slides.pdf
--
-- It implements some of http://pauillac.inria.fr/~fpottier/slides/fpottier-2007-05-linear-bestiary.pdf

{-# INLINE run #-}
run :: A a St St -> IO a
run (A f) = IO $ \world -> case f (St world) of (St world', x) -> (# world', x #)

{-# INLINE bimap #-}
bimap (A f) (A g) = A (\(Pair a b) -> let
	(c, d) = f a
	(e, h) = g b in
	(Pair c e, (d, h)))

{-# INLINE assoc1 #-}
assoc1 = A (\(Pair (Pair a b) c) -> (Pair a (Pair b c), ()))

{-# INLINE assoc2 #-}
assoc2 = A (\(Pair a (Pair b c)) -> (Pair (Pair a b) c, ()))

{-# INLINE drop1 #-}
drop1 = A (\(Pair Blank x) -> (x, ()))

{-# INLINE drop2 #-}
drop2 = A (\(Pair x Blank) -> (x, ()))

{-# INLINE undrop1 #-}
undrop1 = A (\x -> (Pair Blank x, ()))

{-# INLINE undrop2 #-}
undrop2 = A (\x -> (Pair x Blank, ()))

{-# INLINE swap #-}
swap = A (\(Pair x y) -> (Pair y x, ()))

{-# INLINE apply #-}
apply = A (\(Pair (A f) x) -> f x)

{-# INLINE curry #-}
curry (A f) = A (\x -> (A (\y -> f (Pair x y)), ()))

{-# INLINE distr #-}
distr = A (\(Pair a ei) -> (either (Left . Pair a) (Right . Pair a) ei, ()))

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

{-# INLINE void' #-}
void' = (>>== const (rtn ()))

{-# INLINE bimap' #-}
bimap' :: A () t u -> A () v w -> A () (Pair t v) (Pair u w)
bimap' a a2 = void' (bimap a a2)

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

{-# INLINE lift #-}
lift f = A $ \(Pair x (St world)) -> let
	IO g = f x
	(# world', (y, z) #) = g world in (Pair y (St world'), z)

open file mode = lift (\Blank -> liftM (\hdl -> (Open hdl, ())) $ openFile file mode)

getStdin = A (\Blank -> (Open stdin, ()))

getStdout = A (\Blank -> (Open stdout, ()))

getStderr = A (\Blank -> (Open stderr, ()))

close = drop1 . lift (\(Open hdl) -> unless (hdl `elem` [stdin, stdout, stderr]) (hClose hdl) >> return (Blank, ()))

close1 = drop1 . lift (\(Semiclosed hdl) -> hClose hdl >> return (Blank, ()))

fileSize = lift (\h@(Open hdl) -> liftM ((,) h) (hFileSize hdl))

setFileSize sz = lift (\h@(Open hdl) -> hSetFileSize hdl sz >> return (h, ()))

eof = lift (\h@(Open hdl) -> liftM ((,) h) (hIsEOF hdl))

seek mode pos = lift (\h@(Open hdl) -> hSeek hdl mode pos >> return (h, ()))

tell = lift (\h@(Open hdl) -> liftM ((,) h) (hTell hdl))

char = lift (\h@(Open hdl) -> liftM ((,) h) (hGetChar hdl))

line = lift (\h@(Open hdl) -> liftM ((,) h) (hGetLine hdl))

lookahead = lift (\h@(Open hdl) -> liftM ((,) h) (hLookAhead hdl))

contents = lift (\(Open hdl) -> liftM ((,) (Semiclosed hdl)) (hGetContents hdl))

putC c = lift  (\h@(Open hdl) -> hPutChar hdl c >> return (h, ()))

putS s = lift (\h@(Open hdl) -> hPutStr hdl s >> return (h, ()))

{-# NOINLINE random #-}
-- Random numbers have no interesting dependence on the world state,
-- so it is not threaded.
random = A (\Blank -> unsafePerformIO $ liftM ((,) Blank) getStdGen)

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

instance Storable Blank where
	sizeOf _ = 0
	alignment _ = 1
	peek _ = return Blank
	poke _ Blank = return ()

align x y = ((sizeOf x - 1) `div` alignment y + 1) * alignment y

frst :: (Storable a, Storable b) => Ptr (Pair a b) -> Ptr a
frst = castPtr

secnd :: forall a b. (Storable a, Storable b) => Ptr (Pair a b) -> Ptr b
secnd = castPtr . (`plusPtr` align (undefined :: a) (undefined :: b))

instance (Storable a, Storable b) => Storable (Pair a b) where
	sizeOf _ = align (undefined :: a) (undefined :: b) + sizeOf (undefined :: b)
	alignment _ = alignment (undefined :: a)
		`lcm` alignment (undefined :: b)
	peek p = liftM2 Pair (peek (frst p)) (peek (secnd p))
	poke p (Pair x y) = do
		poke (frst p) x
		poke (secnd p) y

-- | With the Fix constructor, I can build data structures of linear data.
data Fix f = In (f (Fix f))

fixInj1 :: Pointer p s (Fix f) -> Pointer p s (f (Fix f))
fixInj1 (Pointer fp p world) = Pointer fp (castPtr p) world

fixInj2 :: Pointer p s (f (Fix f)) -> Pointer p s (Fix f)
fixInj2 (Pointer fp p world) = Pointer fp (castPtr p) world

-- | Nonlinear resources.
class Nonlinear t

instance Nonlinear Char

instance (Nonlinear t) => Nonlinear [t]

instance Nonlinear Bool

instance Nonlinear Ordering

instance Nonlinear Integer

instance Nonlinear Int

instance Nonlinear Float

instance Nonlinear Double

instance (Nonlinear t, Nonlinear u) => Nonlinear (t, u)

instance (Nonlinear t, Nonlinear u, Nonlinear v) => Nonlinear (t, u, v)

instance (Nonlinear t, Nonlinear u, Nonlinear v, Nonlinear w) => Nonlinear (t, u, v, w)

instance (Nonlinear t, Nonlinear u, Nonlinear v, Nonlinear w, Nonlinear x) => Nonlinear (t, u, v, w, x)

-- | Contraction and weakening are available for pointers created from ForeignPtrs
--   (and pointers I am focusing on).
contraction :: A () (Pointer p Foreign t) (Pair (Pointer p Foreign t) (Pointer p Foreign t))
contraction = A (\p -> (Pair p p, ()))

weakening :: A () (Pointer p Foreign t) Blank
weakening = A (\(Pointer fp _ world) -> let IO f = touchForeignPtr fp in
	case f world of (# _, () #) -> (Blank, ()))

--- Safe pointer facilities

newForeign :: A (ForeignPtr t) (Pair (Pointer p Freeable t) St) St
newForeign = lift (\(Pointer _ p _) -> liftM (\fp -> (Blank, fp)) $ newForeignPtr_ p) >>== \fp -> drop1 >>== \_ -> rtn fp

{-# NOINLINE new #-}
new :: (Storable t) => A () Blank (Pointer p Placeholder t)
new = A (\Blank -> unsafePerformIO (liftM (\p -> (Pointer dummy p realWorld#, ())) A.malloc))

-- Gets the pointer from a ForeignPtr unsafely, however linearity ensures that
-- user programs must touch the ForeignPtr in order to dispose of it
-- (see 'weakening').
get :: (Nonlinear t) => ForeignPtr t -> A () Blank (Pointer p Foreign t)
get fp = A (\Blank -> (Pointer (castForeignPtr fp) (unsafeForeignPtrToPtr fp) realWorld#, ()))

-- By using a smuggled RealWorld, I sequence freeing of a pointer without
-- requiring explicit world sequencing.
free :: A () (Pointer p2 Placeholder t) Blank
free = A (\(Pointer _ p world) -> let IO f = A.free p in
	case f world of (# _, () #) -> (Blank, ()))

-- | Focusing on a pointer.
focus :: (forall p. A a (Pair (Pointer p Foreign t) u) (Pair v St)) -> A a (Pair (Pointer p2 s t) u) (Pair (Pair (Pointer p2 s t) v) St)
focus (A f) = A (\(Pair ptr@(Pointer fp p _) x) -> first (\(Pair x st) -> Pair (Pair ptr x) st) (f (Pair (Pointer fp p realWorld#) x))) >>== \x -> updateWorld1 >>== \_ -> rtn x

--- Strong update

updateWorld :: A () (Pair (Pointer p s t) St) (Pair (Pointer p s t) St)
updateWorld = A (\(Pair (Pointer fp p _) st@(St world)) -> (Pair (Pointer fp p world) st, ()))

manipulate = assoc2 . bimap' id swap . assoc1

updateWorld1 = manipulate . bimap' updateWorld id . manipulate

peek' :: (Storable t) => Fn (Pointer p Freeable t) (Pair (Pointer p Placeholder t) t)
peek' = updateWorld1 . lift (\(Pointer fp p st) -> liftM (\x -> (Pair (Pointer fp p st) x, ())) (peek p))

poke' :: (Storable t) => Fn (Pair (Pointer p Placeholder t) t) (Pointer p Freeable t)
poke' = updateWorld . lift (\(Pair (Pointer fp p world) x) -> poke p x >> return (Pointer fp p world, ()))

changeType :: forall t u p. (Storable t, Storable u) => A () (Pointer p Placeholder t) (Pointer p Placeholder u)
changeType = if sizeOf (undefined :: u) <= sizeOf (undefined :: t) then
		A (\(Pointer fp p world) -> (Pointer (castForeignPtr fp) (castPtr p) world, ()))
	else
		error "Linear.changeType: value won't fit"

--- Weak update

peek1 :: (Nonlinear t, Storable t) => A t (Pair (Pointer p s t) St) (Pair (Pointer p s t) St)
peek1 = lift (\ptr@(Pointer _ p _) -> liftM (\x -> (ptr, x)) $ peek p) >>== \x -> updateWorld >>== \_ -> rtn x

poke1 :: (Storable t) => t -> Fn (Pointer p2 s t) (Pointer p2 s t)
poke1 x = lift (\ptr@(Pointer _ p _) -> poke p x >> return (ptr, ())) >>> updateWorld

-- | Split a pointer to a pair, into a pair of pointers.
split :: forall t u p. (Storable t, Storable u) => A () (Pointer p Foreign (Pair t u)) (Pair (Pointer p Foreign t) (Pointer p Foreign u))
split = A (\(Pointer fp p _) -> (Pair
	(Pointer fp (frst p) realWorld#)
	(Pointer fp (secnd p) realWorld#), ()))

------------------------------------------------------
--- Sample program

helloWorld = run $ undrop1
	>>> bimap' getStdout id
	>>> putS "Hello world!\n"
	>>> close