{-# LANGUAGE MagicHash #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecursiveDo #-}
module UIO where

import System.Mem
import System.Mem.Weak
import Data.IORef
import GHC.IORef
import GHC.STRef
import GHC.IO
import GHC.Weak
import GHC.Prim
import Control.Monad.Primitive
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import Data.Set (Set)
import qualified Data.Set as Set
import System.IO.Unsafe
import Control.Monad
import Control.Concurrent
import Data.Foldable
import Data.These
import Unsafe.Coerce
import Control.Monad.Fix
import Data.List.NonEmpty (NonEmpty (..))
import Data.Semigroup
import GHC.Magic

data RemainingWork = RemainingWork

instance Semigroup RemainingWork where
  {-# INLINE (<>) #-}
  (<>) = seq
  sconcat = mconcat . toList
  stimes _ d = d

instance Monoid RemainingWork where
  {-# INLINE mempty #-}
  mempty = RemainingWork
  mconcat = \case
    [] -> RemainingWork
    x : xs -> x `seq` mconcat xs

-- Unordered IO - we want to allocate things, strictly evaluate things, etc., but we don't actually care what order it is done in
newtype UIO2 a = UIO2 { unUIO2 :: State# RealWorld -> (RemainingWork, a) }

instance Functor UIO2 where
  {-# INLINE fmap #-}
  fmap f x = x >>= (pure . f)

instance Applicative UIO2 where
  {-# INLINE pure #-}
  {-# INLINE (*>) #-}
  {-# INLINE (<*>) #-}
  pure x = UIO2 (\s -> (RemainingWork, x))
  UIO2 m *> UIO2 k = UIO2 (\s ->
    let (ms, _) = m (uniqueState 1# s)
        (ks, b) = k (uniqueState 2# s)
    in (ms <> ks, b))
  UIO2 m <*> UIO2 k = UIO2 (\s ->
    let (ms, f) = m (uniqueState 1# s)
        (ks, x) = k (uniqueState 2# s)
    in (ms <> ks, f x))

instance Monad UIO2 where
  {-# INLINE (>>)   #-}
  {-# INLINE (>>=)  #-}
  (>>)      = (*>)
  UIO2 m >>= k = UIO2 (\s ->
    let (ms, a) = m (uniqueState 1# s)
        (ks, b) = unUIO2 (k a) (uniqueState 2# s)
    in (ms <> ks, b))

instance MonadFix UIO2 where
  mfix k = UIO2 (\s ->
    let (ks, result) = unUIO2 (k result) s
    in (ks, result))

runUIO2 :: UIO2 a -> IO a
runUIO2 (UIO2 m) = do
  -- We use a bang pattern here instead of "evaluate", because "evaluate" leaves a "seq#" clutting up our core, but the bang pattern does not
  (!RemainingWork, result) <- IO (\s -> (# s, m s #)) --TODO: This returns the same state we were given; should we call uniqueState 1 or something on it?
  pure result

-- The following is marked NOINLINE because unsafeDupablePerformIO is marked NOINLINE.  I don't really understand it.
{-# INLINE unordered #-}
unordered :: IO a -> UIO2 a
unordered (IO m) = UIO2 (\s -> let x = case m s of (# _, x #) -> x in (x `seq` RemainingWork, x))

-- | Perform an action only when its result is needed.  This action will be unique, but the computation will be considered finished regardless of whether this action has run.  This is appropriate for functions like `newIORef`.
{-# INLINE timeless #-}
timeless :: IO a -> UIO2 a
timeless (IO m) = UIO2 (\s -> (RemainingWork, case m s of (# _, x #) -> x))

-- Force GHC to treat each of these state tokens as unique.  This way, multiple identical calls, e.g. to newIORef are not treated as identical, because they have different state tokens.  Ideally, we would inline this after common sub-expression elimination finishes, so that it is costless.
{-# NOINLINE uniqueState #-}
uniqueState :: Int# -> State# RealWorld -> State# RealWorld
uniqueState = uniqueState'

{-# NOINLINE uniqueState' #-}
uniqueState' :: Int# -> State# RealWorld -> State# RealWorld
uniqueState' _ s = s
-- This implementation seems to work sometimes, but I don't understand why, and it seems highly dependent on other aspects of the implementation.
-- uniqueState = runRW# (\s _ _ -> s)