Supply.hs

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | This module turns every monoid into a (linear) comonoid.
--
-- It generalises the Writer monad, since `tell` is `consume . spend`.
--
-- It looks a lot like the following twisted (applicativg) functor (see also
-- https://dl.acm.org/doi/abs/10.1145/2976002.2976004 /
-- http://ozark.hendrix.edu/~yorgey/pub/twisted.pdf)
--
-- type F b a = (b , b -> a)
-- pure a = (mempty, pure a)
-- (bf, ff) <*> (ba, fa) = (bf <> ba, \b -> ff b (fa (b<>bf)))
--
-- But `Reader (Supply a)` is a monad, not just an applicative.
--
-- On the other hand it seems weaker than the State monad. I thought that maybe
-- `Reader (Supply (Endo s))` could be equivalent to State, but when you call
-- `Spend`, you only get a function `s -> s`, not quite the current state.
--
-- Anyway, I don't think that this comonoid is actually that useful. At least,
-- used as a monad, despite being ostensibly a Reader monad, this creates all
-- the memory leaks of the Writer monad, as we need to find all the occurrences
-- of `spend` immediately upon leaving `withSupply`. So it's pretty strict, and
-- we create this tree of references mirroring the monadic structure. Just like
-- Writer. And there's nothing you can do with this that you can't do with a
-- State monad. I guess the evaluation order is less sequential. I find it
-- satisfying, but I'm not sure that it actually has value. Besides, maybe it's
-- not very sequential, but it's still pretty strict.
--
-- Maybe if you use the `Supply a` as an explicit argument rather than in a
-- monad, then you actually have fewer synchronisation points and it may have
-- performance implication or something. Me, I use it to generate fresh names
-- and I kind of like it. But I'm not sure I'm being reasonable about all this.

module Supply
  ( Supply,
    withSupply,
    spend,
  )
where

import Control.Exception
import Control.Functor.Linear qualified as Control
import Control.Functor.Linear.Internal.Reader qualified as Linear
import Control.Monad
import Data.Functor.Identity
import Data.IORef
import Data.Unrestricted.Linear
import Prelude.Linear
import System.IO.Unsafe (unsafePerformIO)
import Unsafe.Linear qualified as Unsafe
import Prelude qualified

data Supply a where
  MkSupply :: (IO a -> IO ()) -> IO a -> IO a -> Supply

instance Consumable (Supply a) where
  consume (MkSupply _ _ _) = ()

instance Prelude.Monoid a => Dupable (Supply a) where
  dup2 (MkSupply setCount readCount readMyLeft) = unsafePerformIO $ do
    rl <- newIORef (Prelude.return mempty)
    rr <- newIORef (Prelude.return mempty)
    let setCountl c = writeIORef rl c
        setCountr c = writeIORef rr c
        readCountl = join $ readIORef rl
        readCountr = join $ readIORef rr
        readMyLeftl = readMyLeft -- The left of the left duplicate is the left of the original copy.
        readMyLeftr = do
          cll <- readMyLeft
          cl <- readCountl
          Prelude.return $ cll <> cl
    setCount (do l <- readCountl; r <- readCountr; Prelude.return (l <> r))
    Prelude.return (MkSupply setCountl readCountl readMyLeftl, MkSupply setCountr readCountr readMyLeftr)
  {-# NOINLINE dup2 #-}

spend :: Prelude.Monoid a => Supply a %1 -> a -> Ur a
spend (MkSupply setCount _ readMyLeft) = unsafePerformIO $ do
  setCount (Prelude.return a) -- Note: because of linearity, the count is always set from 0 to a.
  Prelude.return $ Ur $ unsafePerformIO $ do
    -- /!\ Important! Only reads under `Ur`, because we have no control over the
    -- evaluation order anymore. In fact, also only writes before so that all
    -- the writes precede the reads and reads are then deterministic.
    readMyLeft
{-# NOINLINE fresh #-}

withSupply :: (Movable a) => (Supply %1 -> a) %1 -> a
withSupply = Unsafe.toLinear $ \f -> unsafePerformIO $ do
  r <- newIORef (Prelude.return 0)
  let setCount c = writeIORef r c
      readCount = join $ readIORef r
      readMyLeft = Prelude.return 0
  u <- evaluate $ move $ f (MkSupply setCount readCount readMyLeft)
  return $ unur u
{-# NOINLINE withSupply #-}