STMUtil.hs

module STMUtil where

import Prelude

import Control.Concurrent.STM (STM)
import Data.HashMap.Strict (HashMap)
import Data.Hashable (Hashable)

import qualified Data.HashMap.Strict as HashMap
import qualified Debug.Trace
import qualified ListT as LT
import qualified StmContainers.Map as StmMap
import qualified Text.Show

{-# WARNING traceWithM "traceWithM" #-}
traceWithM :: (Monad m, Show a) => String -> m a -> m ()
traceWithM prefix ma = do
  a <- ma
  Debug.Trace.traceM $ prefix <> show a

stmToHashMap
  :: (Eq k, Hashable k)
  => StmMap.Map k v
  -> STM (HashMap k v)
stmToHashMap = stmToHashMapWith pure

stmToHashMap2
  :: (Eq k0, Hashable k0, Eq k1, Hashable k1)
  => StmMap.Map k0 (StmMap.Map k1 v)
  -> STM (HashMap k0 (HashMap k1 v))
stmToHashMap2 = stmToHashMapWith stmToHashMap

stmToHashMap3
  :: (Eq k0, Hashable k0, Eq k1, Hashable k1, Eq k2, Hashable k2)
  => StmMap.Map k0 (StmMap.Map k1 (StmMap.Map k2 v))
  -> STM (HashMap k0 (HashMap k1 (HashMap k2 v)))
stmToHashMap3 = stmToHashMapWith stmToHashMap2

stmToHashMap4
  :: (Eq k0, Hashable k0, Eq k1, Hashable k1, Eq k2, Hashable k2, Eq k3, Hashable k3)
  => StmMap.Map k0 (StmMap.Map k1 (StmMap.Map k2 (StmMap.Map k3 v)))
  -> STM (HashMap k0 (HashMap k1 (HashMap k2 (HashMap k3 v))))
stmToHashMap4 = stmToHashMapWith stmToHashMap3

stmToHashMapWith
  :: (Eq k, Hashable k)
  => (v0 -> STM v1)
  -> StmMap.Map k v0
  -> STM (HashMap k v1)
stmToHashMapWith f s =
  LT.fold go HashMap.empty $ StmMap.listT s
  where
  go m (k, v0) = do
    v1 <- f v0
    pure $ HashMap.insert k v1 m