{-# 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 {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 :: OpenMap f
empty = HashMap (StableName Any) (SomeExpr f) -> OpenMap f
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap HashMap (StableName Any) (SomeExpr f)
forall k v. HashMap k v
HashMap.empty

map :: forall f g. (forall x. f x -> g x) -> OpenMap f -> OpenMap g
map :: (forall x. f x -> g x) -> OpenMap f -> OpenMap g
map forall x. f x -> g x
f = HashMap (StableName Any) (SomeExpr g) -> OpenMap g
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (HashMap (StableName Any) (SomeExpr g) -> OpenMap g)
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr g))
-> OpenMap f
-> OpenMap g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SomeExpr f -> SomeExpr g)
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SomeExpr f -> SomeExpr g
go (HashMap (StableName Any) (SomeExpr f)
 -> HashMap (StableName Any) (SomeExpr g))
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr f))
-> OpenMap f
-> HashMap (StableName Any) (SomeExpr g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenMap f -> HashMap (StableName Any) (SomeExpr f)
forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    go :: SomeExpr f -> SomeExpr g
go (SomeExpr f v
y) = g v -> SomeExpr g
forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (f v -> g v
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 x. f x -> Maybe (g x)) -> OpenMap f -> OpenMap g
mapMaybe forall x. f x -> Maybe (g x)
f = HashMap (StableName Any) (SomeExpr g) -> OpenMap g
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (HashMap (StableName Any) (SomeExpr g) -> OpenMap g)
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr g))
-> OpenMap f
-> OpenMap g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SomeExpr f -> Maybe (SomeExpr g))
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr g)
forall v1 v2 k. (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapMaybe SomeExpr f -> Maybe (SomeExpr g)
go (HashMap (StableName Any) (SomeExpr f)
 -> HashMap (StableName Any) (SomeExpr g))
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr f))
-> OpenMap f
-> HashMap (StableName Any) (SomeExpr g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenMap f -> HashMap (StableName Any) (SomeExpr f)
forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    go :: SomeExpr f -> Maybe (SomeExpr g)
go (SomeExpr f v
y) = case f v -> Maybe (g v)
forall x. f x -> Maybe (g x)
f f v
y of
      Just g v
fy -> SomeExpr g -> Maybe (SomeExpr g)
forall a. a -> Maybe a
Just (g v -> SomeExpr g
forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr g v
fy)
      Maybe (g v)
Nothing -> Maybe (SomeExpr g)
forall a. Maybe a
Nothing

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) -> OpenMap f -> OpenMap g
mapWithKey forall d. OpenKey d -> f d -> g d
f = HashMap (StableName Any) (SomeExpr g) -> OpenMap g
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (HashMap (StableName Any) (SomeExpr g) -> OpenMap g)
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr g))
-> OpenMap f
-> OpenMap g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StableName Any -> SomeExpr f -> SomeExpr g)
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr g)
forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
HashMap.mapWithKey StableName Any -> SomeExpr f -> SomeExpr g
go (HashMap (StableName Any) (SomeExpr f)
 -> HashMap (StableName Any) (SomeExpr g))
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr f))
-> OpenMap f
-> HashMap (StableName Any) (SomeExpr g)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenMap f -> HashMap (StableName Any) (SomeExpr f)
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) = g v -> SomeExpr g
forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (OpenKey v -> f v -> g v
forall d. OpenKey d -> f d -> g d
f (StableName Any -> OpenKey v
forall x. StableName Any -> OpenKey x
OpenKey StableName Any
key) f v
y)

lookup :: OpenMap f -> OpenKey x -> Maybe (f x)
lookup :: OpenMap f -> OpenKey x -> Maybe (f x)
lookup (OpenMap HashMap (StableName Any) (SomeExpr f)
m) (OpenKey StableName Any
k) = SomeExpr f -> f x
forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr (SomeExpr f -> f x) -> Maybe (SomeExpr f) -> Maybe (f x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StableName Any
-> HashMap (StableName Any) (SomeExpr f) -> Maybe (SomeExpr f)
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 :: OpenMap f -> [SomeOpenItem f]
toList = ((StableName Any, SomeExpr f) -> SomeOpenItem f)
-> [(StableName Any, SomeExpr f)] -> [SomeOpenItem f]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (StableName Any, SomeExpr f) -> SomeOpenItem f
forall (f :: * -> *).
(StableName Any, SomeExpr f) -> SomeOpenItem f
wrap ([(StableName Any, SomeExpr f)] -> [SomeOpenItem f])
-> (OpenMap f -> [(StableName Any, SomeExpr f)])
-> OpenMap f
-> [SomeOpenItem f]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap (StableName Any) (SomeExpr f)
-> [(StableName Any, SomeExpr f)]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList (HashMap (StableName Any) (SomeExpr f)
 -> [(StableName Any, SomeExpr f)])
-> (OpenMap f -> HashMap (StableName Any) (SomeExpr f))
-> OpenMap f
-> [(StableName Any, SomeExpr f)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OpenMap f -> HashMap (StableName Any) (SomeExpr f)
forall (f :: * -> *).
OpenMap f -> HashMap (StableName Any) (SomeExpr f)
unOpenMap
  where
    wrap :: (StableName Any, SomeExpr f) -> SomeOpenItem f
    wrap :: (StableName Any, SomeExpr f) -> SomeOpenItem f
wrap (StableName Any
key, SomeExpr f
x) = case SomeExpr f
x of
      SomeExpr f v
x' -> OpenKey v -> f v -> SomeOpenItem f
forall (f :: * -> *) x. OpenKey x -> f x -> SomeOpenItem f
SomeOpenItem (StableName Any -> OpenKey v
forall x. StableName Any -> OpenKey x
OpenKey StableName Any
key) f v
x'

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

unsafeCastTypeSomeExpr :: SomeExpr f -> f v
unsafeCastTypeSomeExpr :: SomeExpr f -> f v
unsafeCastTypeSomeExpr = \case
  SomeExpr f v
x -> f v -> f v
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 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) = HashMap (StableName Any) (SomeExpr h) -> OpenMap h
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap ((SomeExpr f -> SomeExpr g -> SomeExpr h)
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr g)
-> HashMap (StableName Any) (SomeExpr h)
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 = h v -> SomeExpr h
forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (f v -> g v -> h v
forall x. f x -> g x -> h x
f f v
x' g v
forall v. g v
y')
      where
        y' :: g v
y' = SomeExpr g -> g v
forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr SomeExpr g
sy

insert :: forall f dx. OpenKey dx -> f dx -> OpenMap f -> OpenMap f
insert :: OpenKey dx -> f dx -> OpenMap f -> OpenMap f
insert (OpenKey StableName Any
k) f dx
x (OpenMap HashMap (StableName Any) (SomeExpr f)
m) = HashMap (StableName Any) (SomeExpr f) -> OpenMap f
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap (StableName Any
-> SomeExpr f
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr f)
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert StableName Any
k (f dx -> SomeExpr f
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 :: (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) = HashMap (StableName Any) (SomeExpr f) -> OpenMap f
forall (f :: * -> *).
HashMap (StableName Any) (SomeExpr f) -> OpenMap f
OpenMap HashMap (StableName Any) (SomeExpr f)
m'
  where
    m' :: HashMap (StableName Any) (SomeExpr f)
m' = (SomeExpr f -> SomeExpr f)
-> StableName Any
-> HashMap (StableName Any) (SomeExpr f)
-> HashMap (StableName Any) (SomeExpr f)
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 = f x -> SomeExpr f
forall (f :: * -> *) v. f v -> SomeExpr f
SomeExpr (f x -> f x
f (SomeExpr f -> f x
forall (f :: * -> *) v. SomeExpr f -> f v
unsafeCastTypeSomeExpr SomeExpr f
x))

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