{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module Downhill.Internal.Graph.OpenMap
  ( -- * OpenMap
    OpenMap,
    OpenKey,
    SomeOpenItem (SomeOpenItem),
    -- * Construction
    makeOpenKey,
    empty,
    insert,
    -- * Query
    lookup,
    toList,
    elems,
    -- * Modify
    map,
    mapWithKey,
    mapMaybe,
    adjust,
    intersectionWith,
  )
where

import Control.Applicative (Const (Const))
import Control.Exception (evaluate)
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import Data.Kind (Type)
import GHC.Base (Any, Maybe (Just, Nothing), coerce)
import GHC.StableName (StableName)
import System.Mem.StableName (makeStableName)
import Unsafe.Coerce (unsafeCoerce)
import Prelude (Functor (fmap), IO, Monad (return), (.), (<$>))

data SomeExpr f = forall v. SomeExpr (f v)

-- | A key of @OpenMap@.
newtype OpenKey x = OpenKey (StableName Any)

-- | Heterogeneous map with 'StableName' as a key.
newtype OpenMap (f :: Type -> Type) = OpenMap {forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap :: HashMap (StableName Any) (SomeExpr f)}

-- | Key and value.
data SomeOpenItem f = forall x. SomeOpenItem (OpenKey x) (f x)

empty :: OpenMap f
empty :: forall (f :: * -> *). OpenMap f
empty = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap forall k v. HashMap k v
HashMap.empty

map :: forall f g. (forall x. f x -> g x) -> OpenMap f -> OpenMap g
map :: forall (f :: * -> *) (g :: * -> *).
(forall x. f x -> g x) -> OpenMap f -> OpenMap g
map forall x. f x -> g x
f = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SomeExpr f -> SomeExpr g
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    go :: SomeExpr f -> SomeExpr g
go (SomeExpr f v
y) = forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (forall x. f x -> g x
f f v
y)

mapMaybe :: forall f g. (forall x. f x -> Maybe (g x)) -> OpenMap f -> OpenMap g
mapMaybe :: forall (f :: * -> *) (g :: * -> *).
(forall x. f x -> Maybe (g x)) -> OpenMap f -> OpenMap g
mapMaybe forall x. f x -> Maybe (g x)
f = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall v1 v2 k. (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapMaybe SomeExpr f -> Maybe (SomeExpr g)
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    go :: SomeExpr f -> Maybe (SomeExpr g)
go (SomeExpr f v
y) = case forall x. f x -> Maybe (g x)
f f v
y of
      Just g v
fy -> forall a. a -> Maybe a
Just (forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr g v
fy)
      Maybe (g v)
Nothing -> forall a. Maybe a
Nothing

mapWithKey :: forall f g. (forall d. OpenKey d -> f d -> g d) -> OpenMap f -> OpenMap g
mapWithKey :: forall (f :: * -> *) (g :: * -> *).
(forall d. OpenKey d -> f d -> g d) -> OpenMap f -> OpenMap g
mapWithKey forall d. OpenKey d -> f d -> g d
f = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapWithKey StableName Any -> SomeExpr f -> SomeExpr g
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    go :: StableName Any -> SomeExpr f -> SomeExpr g
go StableName Any
key (SomeExpr f v
y) = forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (forall d. OpenKey d -> f d -> g d
f (forall x. StableName Any -> OpenKey x
OpenKey StableName Any
key) f v
y)

lookup :: OpenMap f -> OpenKey x -> Maybe (f x)
lookup :: forall (f :: * -> *) x. OpenMap f -> OpenKey x -> Maybe (f x)
lookup (OpenMap HashMap (StableName Any) (SomeExpr f)
m) (OpenKey StableName Any
k) = forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup StableName Any
k HashMap (StableName Any) (SomeExpr f)
m

toList :: OpenMap f -> [SomeOpenItem f]
toList :: forall (f :: * -> *). OpenMap f -> [SomeOpenItem f]
toList = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (f :: * -> *).
(StableName Any, SomeExpr f) -> SomeOpenItem f
wrap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. HashMap k v -> [(k, v)]
HashMap.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    wrap :: (StableName Any, SomeExpr f) -> SomeOpenItem f
    wrap :: forall (f :: * -> *).
(StableName Any, SomeExpr f) -> SomeOpenItem f
wrap (StableName Any
key, SomeExpr f
x) = case SomeExpr f
x of
      SomeExpr f v
x' -> forall (f :: * -> *) x. OpenKey x -> f x -> SomeOpenItem f
SomeOpenItem (forall x. StableName Any -> OpenKey x
OpenKey StableName Any
key) f v
x'

elems :: OpenMap (Const b) -> [b]
elems :: forall b. OpenMap (Const b) -> [b]
elems = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall r. SomeExpr (Const r) -> r
unSomeExpr forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. HashMap k v -> [v]
HashMap.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    unSomeExpr :: SomeExpr (Const r) -> r
    unSomeExpr :: forall r. SomeExpr (Const r) -> r
unSomeExpr (SomeExpr (Const r
x)) = r
x

unsafeCastTypeSomeExpr :: SomeExpr f -> f v
unsafeCastTypeSomeExpr :: forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr = \case
  SomeExpr f v
x -> forall a b. a -> b
unsafeCoerce f v
x

intersectionWith :: forall f g h. (forall x. f x -> g x -> h x) -> OpenMap f -> OpenMap g -> OpenMap h
intersectionWith :: forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
(forall x. f x -> g x -> h x)
-> OpenMap f -> OpenMap g -> OpenMap h
intersectionWith forall x. f x -> g x -> h x
f (OpenMap HashMap (StableName Any) (SomeExpr f)
x) (OpenMap HashMap (StableName Any) (SomeExpr g)
y) = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (forall k v1 v2 v3.
(Eq k, Hashable k) =>
(v1 -> v2 -> v3) -> HashMap k v1 -> HashMap k v2 -> HashMap k v3
HashMap.intersectionWith SomeExpr f -> SomeExpr g -> SomeExpr h
f' HashMap (StableName Any) (SomeExpr f)
x HashMap (StableName Any) (SomeExpr g)
y)
  where
    f' :: SomeExpr f -> SomeExpr g -> SomeExpr h
f' (SomeExpr f v
x') SomeExpr g
sy = forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (forall x. f x -> g x -> h x
f f v
x' forall {v}. g v
y')
      where
        y' :: g v
y' = forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr SomeExpr g
sy

insert :: forall f dx. OpenKey dx -> f dx -> OpenMap f -> OpenMap f
insert :: forall (f :: * -> *) dx.
OpenKey dx -> f dx -> OpenMap f -> OpenMap f
insert (OpenKey StableName Any
k) f dx
x (OpenMap HashMap (StableName Any) (SomeExpr f)
m) = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert StableName Any
k (forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr f dx
x) HashMap (StableName Any) (SomeExpr f)
m)

adjust :: forall f x. (f x -> f x) -> OpenKey x -> OpenMap f -> OpenMap f
adjust :: forall (f :: * -> *) x.
(f x -> f x) -> OpenKey x -> OpenMap f -> OpenMap f
adjust f x -> f x
f (OpenKey StableName Any
key) (OpenMap HashMap (StableName Any) (SomeExpr f)
m) = forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap HashMap (StableName Any) (SomeExpr f)
m'
  where
    m' :: HashMap (StableName Any) (SomeExpr f)
m' = forall k v.
(Eq k, Hashable k) =>
(v -> v) -> k -> HashMap k v -> HashMap k v
HashMap.adjust SomeExpr f -> SomeExpr f
f' StableName Any
key HashMap (StableName Any) (SomeExpr f)
m
    f' :: SomeExpr f -> SomeExpr f
f' SomeExpr f
x = forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (f x -> f x
f (forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr SomeExpr f
x))

makeOpenKey :: f v -> IO (OpenKey v)
makeOpenKey :: forall (f :: * -> *) v. f v -> IO (OpenKey v)
makeOpenKey f v
x = do
  f v
x' <- forall a. a -> IO a
evaluate f v
x
  StableName (f v)
z <- forall a. a -> IO (StableName a)
makeStableName f v
x'
  forall (m :: * -> *) a. Monad m => a -> m a
return (forall x. StableName Any -> OpenKey x
OpenKey (coerce :: forall a b. Coercible a b => a -> b
coerce StableName (f v)
z))