Unique.hs

{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE UndecidableInstances #-}

module BrushStroking.Unique
  ( MonadUnique(freshUnique)
  , Unique, unsafeUnique
  , uniqueText
  , UniqueSupply, newUniqueSupply
  , uniqueMapFromList
  )
  where

-- base
import Control.Arrow
  ( (&&&) )
import Data.Int
  ( Int64 )
import Data.Word
  ( Word32 )
import Foreign.Storable
  ( Storable )

-- containers
import Data.Map.Strict
  ( Map )
import qualified Data.Map.Strict as Map
  ( fromList )

-- deepseq
import Control.DeepSeq
  ( NFData )

-- generic-lens
import Data.Generics.Product.Typed
  ( HasType(typed) )

-- lens
import Control.Lens
  ( view )

-- mtl
import Control.Monad.Reader
  ( MonadReader(..) )

-- stm
import Control.Concurrent.STM
  ( STM )
import qualified Control.Concurrent.STM as STM

-- text
import Data.Text
  ( Text )
import qualified Data.Text as Text
  ( pack )

-- transformers
import Control.Monad.IO.Class
  ( MonadIO(..) )
import Control.Monad.Trans.Class
  ( lift )
import Control.Monad.Trans.Reader
  ( ReaderT )

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

newtype Unique = Unique { unique :: Int64 }
  deriving stock   Show
  deriving newtype ( Eq, Ord, Enum, Storable, NFData )

unsafeUnique :: Word32 -> Unique
unsafeUnique i = Unique ( -(fromIntegral i) - 1 )

uniqueText :: Unique -> Text
uniqueText ( Unique i )
  | i >= 0
  = "%" <> Text.pack ( show i )
  | otherwise
  = "§" <> Text.pack ( show $ -i - 1 )

newtype UniqueSupply = UniqueSupply { uniqueSupplyTVar :: STM.TVar Unique }

instance Show UniqueSupply where { show _ = "Unique supply" }

newUniqueSupply :: IO UniqueSupply
newUniqueSupply = UniqueSupply <$> STM.newTVarIO ( Unique 1 )

uniqueMapFromList :: HasType Unique a => [ a ] -> Map Unique a
uniqueMapFromList = Map.fromList . map ( view typed &&& id )

class Monad m => MonadUnique m where
  freshUnique :: m Unique

instance {-# OVERLAPPABLE #-} ( Monad m, MonadReader r m, HasType UniqueSupply r, MonadIO m ) => MonadUnique m where
  freshUnique = do
    UniqueSupply { uniqueSupplyTVar } <- view ( typed @UniqueSupply )
    liftIO $ STM.atomically $ STM.stateTVar uniqueSupplyTVar doSucc

instance MonadUnique ( ReaderT UniqueSupply STM ) where
  freshUnique = do
    UniqueSupply { uniqueSupplyTVar } <- ask
    lift $ STM.stateTVar uniqueSupplyTVar doSucc

doSucc :: Unique -> ( Unique, Unique )
doSucc uniq@( Unique !i ) = ( uniq, Unique ( succ i ) )