{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE LiberalTypeSynonyms #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE RebindableSyntax #-}
module Resource where

import qualified Prelude as P
import Prelude (Eq, Show, undefined, Integral, fromInteger, id, (.), ($), const, Monad, String, Bool(True, False))
import Control.Monad (liftM)

fail :: String -> a
fail = P.error

class IMonad m where
  return :: a -> m i i a
  (>>=)  :: m i j a -> (a -> m j k b) -> m i k b

(>>) :: IMonad m => m i j a -> m j k b -> m i k b
x >> y = x >>= const y

data Nat = Z | S Nat
type State = (Nat, [Nat]) -- why can't I promote this?

-- Other substructural flavors probably work too, but I don't need them now
class Linear (xs :: [Nat])
instance Linear '[]
instance Linear xs => Linear (S Z ': xs)

newtype Ref (n :: Nat) s a = Ref { getRef :: a }

-- Eww, when I said linear types, I didn't mean O(n)
-- This is also annoying because it doesn't reduce immediately and leads to massive types. I welcome improvements!
type family   Grow (xs :: [Nat]) :: [Nat]
type instance Grow '[] = '[Z]
type instance Grow (x ': xs) = x ': Grow xs

type family   Use (n :: Nat) (xs :: [Nat]) :: [Nat]
type instance Use Z (x ': xs) = S x ': xs
type instance Use (S i) (x ': xs) = x ': Use i xs

newtype Res s m (i :: (Nat, [Nat])) (j :: (Nat, [Nat])) a = Res { _runRes :: m a }

instance Monad m => IMonad (Res s m) where
  return  = Res . P.return
  x >>= f = Res $ _runRes x P.>>= _runRes . f

create :: Monad m => m a -> Res s m '(i, xs) '(S i, Grow xs) (Ref i s a)
create x = Res (liftM Ref x)

consume :: Monad m => Ref i s a -> Res s m '(j, xs) '(j, Use i xs) a
consume (Ref x) = Res $ P.return x

newtype Conjoined s m i a = Conjoined { runConjoined :: forall j ys. Res s m '(j, ys) '(S j, Grow (Use i ys)) (Ref j s a) }

conjoin :: Monad m => m (m a) -> Res s m '(i, xs) '(S i, Grow xs) (Conjoined s m i a)
conjoin before = Res $ before P.>>= \after -> P.return (Conjoined (Res (liftM Ref after)))

runRes :: Linear xs => (forall s. Res s m '(Z, '[]) '(j, xs) a) -> m a
runRes (Res x) = x


-- Pretend constructors etc. aren't exported. Gah I wish Haskell had real modules.


newtype Id a = Id a deriving (Eq, Show)
instance Monad Id where
  return = Id
  Id x >>= f = Id (case f x of Id x' -> x')

zomg = conjoin (Id (Id "bar"))

test = do
  x <- create (Id "foo")
  y <- create (Id True)
  Conjoined foo <- zomg
  consume y
  z <- create (Id 5)
  consume x
  a <- foo
  consume a
  consume z


run = runRes test