Ids.hs

{-# LANGUAGE BangPatterns    #-}
{-# LANGUAGE GADTs           #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE UnboxedTuples   #-}

-- Note that argument order is more like IORef than Map, because its mutable
module Development.IDE.Graph.Internal.Ids(
    Ids, Id,
    empty, insert, lookup, fromList,
    null, size, sizeUpperBound,
    forWithKeyM_, forCopy, forMutate,
    toList, elems, toMap
    ) where

import           Control.Exception
import           Control.Monad.Extra
import           Data.Functor
import qualified Data.HashMap.Strict                   as Map
import           Data.IORef.Extra
import           Data.List.Extra                       (zipFrom)
import           Data.Maybe
import           Data.Primitive.Array                  hiding (fromList)
import           Development.IDE.Graph.Internal.Intern (Id)
import           GHC.Exts                              (RealWorld)
import           GHC.IO                                (IO (..))
import           Prelude                               hiding (lookup, null)


newtype Ids a = Ids (IORef (S a))

data S a = S
    {capacity :: {-# UNPACK #-} !Int -- ^ Number of entries in values, initially 0
    ,used     :: {-# UNPACK #-} !Int -- ^ Capacity that has been used, assuming no gaps from index 0, initially 0
    ,values   :: {-# UNPACK #-} !(MutableArray RealWorld (Maybe a))
    }


empty :: IO (Ids a)
empty = do
    let capacity = 0
    let used = 0
    values <- newArray capacity Nothing
    Ids <$> newIORef S{..}

fromList :: [a] -> IO (Ids a)
fromList xs = do
    let capacity = length xs
    let used = capacity
    values <- newArray capacity Nothing
    forM_ (zipFrom 0 xs) $ \(i, x) ->
        writeArray values i $ Just x
    Ids <$> newIORef S{..}

sizeUpperBound :: Ids a -> IO Int
sizeUpperBound (Ids ref) = do
    S{..} <- readIORef ref
    pure used


size :: Ids a -> IO Int
size (Ids ref) = do
    S{..} <- readIORef ref
    let go !acc i
            | i < 0 = pure acc
            | otherwise = do
                v <- readArray values i
                if isJust v then go (acc+1) (i-1) else go acc (i-1)
    go 0 (used-1)


toMap :: Ids a -> IO (Map.HashMap Id a)
toMap ids = do
    mp <- Map.fromList <$> toListUnsafe ids
    pure $! mp

forWithKeyM_ :: Ids a -> (Id -> a -> IO ()) -> IO ()
forWithKeyM_ (Ids ref) f = do
    S{..} <- readIORef ref
    let go !i | i >= used = pure ()
              | otherwise = do
                v <- readArray values i
                whenJust v $ f $ fromIntegral i
                go $ i+1
    go 0

forCopy :: Ids a -> (a -> b) -> IO (Ids b)
forCopy (Ids ref) f = do
    S{..} <- readIORef ref
    values2 <- newArray capacity Nothing
    let go !i | i >= used = pure ()
              | otherwise = do
                v <- readArray values i
                whenJust v $ \v -> writeArray values2 i $ Just $ f v
                go $ i+1
    go 0
    Ids <$> newIORef (S capacity used values2)


forMutate :: Ids a -> (Id -> a -> a) -> IO ()
forMutate (Ids ref) f = do
    S{..} <- readIORef ref
    let go !i | i >= used = pure ()
              | otherwise = do
                v <- readArray values i
                whenJust v $ \v -> writeArray values i $! Just $! f i v
                go $ i+1
    go 0


toListUnsafe :: Ids a -> IO [(Id, a)]
toListUnsafe (Ids ref) = do
    S{..} <- readIORef ref

    -- execute in O(1) stack
    -- see https://neilmitchell.blogspot.co.uk/2015/09/making-sequencemapm-for-io-take-o1-stack.html
    let index _ i | i >= used = []
        index r i | IO io <- readArray values i = case io r of
            (# r, Nothing #) -> index r (i+1)
            (# r, Just v  #) -> (fromIntegral i, v) : index r (i+1)

    IO $ \r -> (# r, index r 0 #)


toList :: Ids a -> IO [(Id, a)]
toList ids = do
    xs <- toListUnsafe ids
    let demand (_:xs) = demand xs
        demand []     = ()
    evaluate $ demand xs
    pure xs

elems :: Ids a -> IO [a]
elems ids = map snd <$> toList ids

null :: Ids a -> IO Bool
null ids = (== 0) <$> sizeUpperBound ids


insert :: Ids a -> Id -> a -> IO ()
insert (Ids ref) (i) v = do
    S{..} <- readIORef ref
    let ii = fromIntegral i
    if ii < capacity then do
        writeArray values ii $ Just v
        when (ii >= used) $ writeIORef' ref S{used=ii+1,..}
     else do
        c2<- pure $ max (capacity * 2) (ii + 10000)
        v2 <- newArray c2 Nothing
        copyMutableArray v2 0 values 0 capacity
        writeArray v2 ii $ Just v
        writeIORef' ref $ S c2 (ii+1) v2

lookup :: Ids a -> Id -> IO (Maybe a)
lookup (Ids ref) (i) = do
    S{..} <- readIORef ref
    let ii = fromIntegral i
    if ii < used then
        readArray values ii
     else
        pure Nothing