module InfiniteMap
  (
    InfiniteMap,
    fromList,
    (!)
  )
  where

import System.IO.Unsafe
import Data.IORef
import qualified Data.Map as M

data InfiniteMap k v = InfiniteMap {imRef :: IORef ((M.Map k v), [(k, v)])}

fromList :: Ord k => [(k, v)] -> InfiniteMap k v
fromList l = InfiniteMap (unsafePerformIO $ newIORef (M.empty, l))

fillMapUntil :: (Ord k, Eq k) => k -> (M.Map k v, [(k, v)]) -> (M.Map k v, [(k, v)])
fillMapUntil tk (m, []) = (m, [])
fillMapUntil tk (m, ((k, v):xs))
  | tk == k = (filledMap, xs)
  | otherwise = fillMapUntil tk (filledMap, xs)
  where
    filledMap = M.insertWith' (\a _ -> a) k v m

(!) :: Ord k => InfiniteMap k v -> k -> v
m ! k = if k `M.member` cmap then (M.!) cmap k else (if k `M.member` newMap then (M.!) newMap k else error "Key not in map")
  where
    cmap = fst $ unsafePerformIO $ readIORef $ imRef m
    newMap = unsafePerformIO $ do
      (nm, nl) <- atomicModifyIORef (imRef m) (\a -> let res = fillMapUntil k a in (res, res))
      return nm

insert :: Ord k => k -> v -> InfiniteMap k v -> InfiniteMap k v
insert k v m = InfiniteMap $ unsafePerformIO $ newIORef (M.insert k v cmap, clist)
  where
    (cmap, clist) = unsafePerformIO $ readIORef $ imRef m