Appendix: Functional programmers enter a diner…

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}

module Scratch.Locks where

import Control.Concurrent
import Control.Monad.Cont
import Data.Coerce (coerce)
import Data.Functor.Compose
import Data.IORef
import Data.List (sortOn)

data Locks = Locks
  { aLock :: MVar Int,
    bLock :: MVar String,
    cLock :: MVar ()
  }

mkLocks :: IO Locks
mkLocks =
  Locks
    <$> newMVar 0
    <*> newMVar ""
    <*> newMVar ()

type family Result r where
  Result (a -> r) = Result r
  Result r = r

type family First r where
  First (a -> r) = a
  First r = r

data ALock = ALock

data BLock = BLock

data CLock = CLock

data AllLocks
  = ALock_
  | BLock_
  | CLock_
  deriving (Eq, Ord)

data LockRequest l where
  ALockRequest :: LockRequest ALock
  BLockRequest :: LockRequest BLock
  CLockRequest :: LockRequest CLock

lockToValue :: Some LockRequest -> AllLocks
lockToValue s = case s of
  Some ALockRequest -> ALock_
  Some BLockRequest -> BLock_
  Some CLockRequest -> CLock_

valueToLock :: AllLocks -> Some LockRequest
valueToLock s = case s of
  ALock_ -> Some ALockRequest
  BLock_ -> Some BLockRequest
  CLock_ -> Some CLockRequest

data ContentsOf l where
  LockOfA :: Int -> ContentsOf ALock
  LockOfB :: String -> ContentsOf BLock
  LockOfC :: () -> ContentsOf CLock

data Some f = forall l. Some {getSome :: f l}

data LockProvider provider continuation where
  LockProviderCons ::
    forall l provider continuation r.
    LockRequest l ->
    (provider (ContentsOf l)) ->
    LockProvider provider (continuation -> r) ->
    LockProvider provider ((ContentsOf l -> continuation) -> r)
  LockProviderNil ::
    LockProvider provider (r -> r)

newtype Acquire l = Acquire {acquire :: forall r. ContT r IO l}
  deriving (Functor)

getA :: Locks -> Acquire (ContentsOf ALock)
getA locks = Acquire $ ContT $ \action' -> do
  withMVar (aLock locks) $
    \contents -> do
      print "acuire A"
      result <- action' (LockOfA contents)
      print "release A"
      pure result

getB :: Locks -> Acquire (ContentsOf BLock)
getB locks = Acquire $ ContT $ \action' -> do
  withMVar (bLock locks) $
    \contents -> do
      print "acuire B"
      result <- action' (LockOfB contents)
      print "release B"
      pure result

getC :: Locks -> Acquire (ContentsOf CLock)
getC locks = Acquire $ ContT $ \action' -> do
  withMVar (cLock locks) $
    \contents -> do
      print "acuire C"
      result <- action' (LockOfC contents)
      print "release C"
      pure result

data Defer l
  = Defer {unDefer :: Acquire l}
  | Evaluated l

data BiSome f g = forall l. BiSome {getBiSome :: (f l, g l)}

type f ∘ g = Compose f g

infixr 8 ∘

prepareLockProvider ::
  LockProvider Acquire continuation ->
  IO (LockProvider (IORef ∘ Defer) continuation)
prepareLockProvider lp = do
  case lp of
    LockProviderCons lockRequest action continuation -> do
      ref <- Compose <$> newIORef (Defer action)
      rest <- prepareLockProvider continuation
      pure $ LockProviderCons lockRequest ref rest
    LockProviderNil -> pure LockProviderNil

gatherLocks ::
  LockProvider (IORef ∘ Defer) continuation ->
  IO [BiSome LockRequest (IORef ∘ Defer ∘ ContentsOf)]
gatherLocks lp = case lp of
  LockProviderCons lockRequest action continuation -> do
    let action' = action
    rest <- gatherLocks continuation
    pure $ BiSome (lockRequest, coerce action') : rest
  LockProviderNil -> pure []

giveLockContents :: LockProvider (IORef ∘ Defer) continuation -> IO continuation
giveLockContents lp = case lp of
  LockProviderCons _ ioRef continuation -> do
    ref <- readIORef (getCompose ioRef)
    case ref of
      Evaluated l -> do
        continue <- giveLockContents continuation
        pure $ \f -> continue (f l)
      _ -> error "not evaluated yet"
  LockProviderNil -> pure id

requestLocksSorted ::
  forall r.
  [BiSome LockRequest (IORef ∘ Defer ∘ ContentsOf)] ->
  IO r ->
  IO r
requestLocksSorted ls action = do
  let eval :: [BiSome LockRequest (IORef ∘ Defer ∘ ContentsOf)] -> IO r
      eval [] = action
      eval ((BiSome (_, Compose ioRef)) : rest) = do
        ref' <- getCompose <$> readIORef ioRef
        case ref' of
          Defer (Acquire l) -> do
            runContT l $ \l' -> do
              writeIORef ioRef (Compose (Evaluated l'))
              eval rest
          Evaluated _ ->
            eval rest
      lockValue (BiSome (lr, _ioRef)) = lockToValue (Some lr)
  eval (sortOn lockValue ls)

provideAllLocks :: LockProvider Acquire continuation -> (continuation -> IO r) -> IO r
provideAllLocks lp action = do
  lp' <- prepareLockProvider lp
  gathered <- gatherLocks lp'
  let actWithLocks = giveLockContents lp' >>= action
  requestLocksSorted gathered actWithLocks

example :: Locks -> (ContentsOf ALock -> ContentsOf BLock -> IO b) -> IO b
example locks action =
  provideAllLocks
    (LockProviderCons ALockRequest (getA locks) (LockProviderCons BLockRequest (getB locks) LockProviderNil))
    $ \with -> with action

class AccessLock l where
  accessLock :: Locks -> Acquire (ContentsOf l)
  requestsLock :: LockRequest l

instance AccessLock ALock where
  accessLock = getA
  requestsLock = ALockRequest

instance AccessLock BLock where
  accessLock = getB
  requestsLock = BLockRequest

instance AccessLock CLock where
  accessLock = getC
  requestsLock = CLockRequest

class ProvidesLock r where
  provideLock :: Locks -> LockProvider Acquire r

instance {-# OVERLAPPABLE #-} ProvidesLock (IO f -> IO f) where
  provideLock _ = LockProviderNil

instance
  {-# OVERLAPPING #-}
  (ProvidesLock (continuation -> r), AccessLock l) =>
  ProvidesLock ((ContentsOf l -> continuation) -> r)
  where
  provideLock locks = LockProviderCons requestsLock (accessLock locks) (provideLock locks)

provideLocks ::
  (ProvidesLock continuation) => Locks -> (continuation -> IO r) -> IO r
provideLocks state = provideAllLocks (provideLock state)

example' :: Locks -> (ContentsOf BLock -> IO Int) -> IO Int
example' locks action = do
  provideLocks locks $ \with -> with action

example'' :: Locks -> (ContentsOf BLock -> ContentsOf ALock -> IO Int) -> IO Int
example'' locks action = do
  provideLocks locks $ \with -> with action

withLocks :: (ProvidesLock (f -> IO b), Result f ~ IO b) => Locks -> f -> IO b
withLocks locks action =
  provideLocks locks $ \with ->
    with action

example''' :: IO ()
example''' = do
  locks <- mkLocks
  withLocks @(ContentsOf BLock -> ContentsOf ALock -> IO ()) locks $
    \_b _a -> print "contents"