{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE UndecidableInstances #-}

module Downhill.Internal.Graph.NodeMap
  ( -- * NodeMap
    NodeMap,
    NodeKey,
    -- * Construction
    fromOpenMap,
    generate,
    -- * Query
    lookup,
    tryLookup,
    toList,
    toListWith,
    elems,
    -- * Modify
    map,
    mapWithKey,
    adjust,
    zipWith,
    -- * Node Set
    IsNodeSet,
    SomeNodeMap (..),
    KeyAndValue (..),
  )
where

import Control.Applicative (Const)
import Data.Data (Proxy (Proxy))
import Data.Reflection (Reifies (reflect), reify)
import Downhill.Internal.Graph.OpenMap (OpenKey, OpenMap, SomeOpenItem (SomeOpenItem))
import qualified Downhill.Internal.Graph.OpenMap as OpenMap
import Prelude (Maybe (Just, Nothing), const, error, (.), (<$>))

type role NodeKey nominal nominal

-- | Valid key, guaranteed to be a member of @s@
newtype NodeKey s x = NodeKey (OpenKey x)

-- | @NodeMap s f@ is a map where value of type @f x@ is associated with key @NodeKey s x@.
-- Type variable `s` tracks the set of nodes. Lookups never fail. Maps can
-- be zipped without losing any nodes.
newtype NodeMap s f = NodeMap {NodeMap s f -> OpenMap f
unNodeMap :: OpenMap f}

data KeyAndValue s f = forall x. KeyAndValue (NodeKey s x) (f x)

class IsNodeSet s where
  allNodes :: OpenMap Proxy

map :: forall s f g. (forall v. f v -> g v) -> NodeMap s f -> NodeMap s g
map :: (forall v. f v -> g v) -> NodeMap s f -> NodeMap s g
map forall v. f v -> g v
f = OpenMap g -> NodeMap s g
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap (OpenMap g -> NodeMap s g)
-> (NodeMap s f -> OpenMap g) -> NodeMap s f -> NodeMap s g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall v. f v -> g v) -> OpenMap f -> OpenMap g
forall (f :: * -> *) (g :: * -> *).
(forall x. f x -> g x) -> OpenMap f -> OpenMap g
OpenMap.map forall v. f v -> g v
f (OpenMap f -> OpenMap g)
-> (NodeMap s f -> OpenMap f) -> NodeMap s f -> OpenMap g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NodeMap s f -> OpenMap f
forall s (f :: * -> *). NodeMap s f -> OpenMap f
unNodeMap

mapWithKey :: forall s f g. (forall x. NodeKey s x -> f x -> g x) -> NodeMap s f -> NodeMap s g
mapWithKey :: (forall x. NodeKey s x -> f x -> g x) -> NodeMap s f -> NodeMap s g
mapWithKey forall x. NodeKey s x -> f x -> g x
f (NodeMap OpenMap f
x) = OpenMap g -> NodeMap s g
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap ((forall d. OpenKey d -> f d -> g d) -> OpenMap f -> OpenMap g
forall (f :: * -> *) (g :: * -> *).
(forall d. OpenKey d -> f d -> g d) -> OpenMap f -> OpenMap g
OpenMap.mapWithKey forall d. OpenKey d -> f d -> g d
f' OpenMap f
x)
  where
    f' :: OpenKey dx -> f dx -> g dx
    f' :: OpenKey dx -> f dx -> g dx
f' OpenKey dx
key' = NodeKey s dx -> f dx -> g dx
forall x. NodeKey s x -> f x -> g x
f (OpenKey dx -> NodeKey s dx
forall s x. OpenKey x -> NodeKey s x
NodeKey OpenKey dx
key')

toList :: NodeMap s f -> [KeyAndValue s f]
toList :: NodeMap s f -> [KeyAndValue s f]
toList = (forall x. NodeKey s x -> f x -> KeyAndValue s f)
-> NodeMap s f -> [KeyAndValue s f]
forall s (f :: * -> *) r.
(forall x. NodeKey s x -> f x -> r) -> NodeMap s f -> [r]
toListWith forall x. NodeKey s x -> f x -> KeyAndValue s f
forall s (f :: * -> *) x. NodeKey s x -> f x -> KeyAndValue s f
KeyAndValue

toListWith :: forall s f r. (forall x. NodeKey s x -> f x -> r) -> NodeMap s f -> [r]
toListWith :: (forall x. NodeKey s x -> f x -> r) -> NodeMap s f -> [r]
toListWith forall x. NodeKey s x -> f x -> r
f (NodeMap OpenMap f
m) = SomeOpenItem f -> r
wrap (SomeOpenItem f -> r) -> [SomeOpenItem f] -> [r]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> OpenMap f -> [SomeOpenItem f]
forall (f :: * -> *). OpenMap f -> [SomeOpenItem f]
OpenMap.toList OpenMap f
m
  where
    wrap :: SomeOpenItem f -> r
    wrap :: SomeOpenItem f -> r
wrap (SomeOpenItem OpenKey x
key f x
value) = NodeKey s x -> f x -> r
forall x. NodeKey s x -> f x -> r
f (OpenKey x -> NodeKey s x
forall s x. OpenKey x -> NodeKey s x
NodeKey OpenKey x
key) f x
value

elems :: NodeMap s (Const b) -> [b]
elems :: NodeMap s (Const b) -> [b]
elems (NodeMap OpenMap (Const b)
m) = OpenMap (Const b) -> [b]
forall b. OpenMap (Const b) -> [b]
OpenMap.elems OpenMap (Const b)
m

lookup :: NodeMap s f -> NodeKey s v -> f v
lookup :: NodeMap s f -> NodeKey s v -> f v
lookup (NodeMap OpenMap f
m) (NodeKey OpenKey v
key) =
  case OpenMap f -> OpenKey v -> Maybe (f v)
forall (f :: * -> *) x. OpenMap f -> OpenKey x -> Maybe (f x)
OpenMap.lookup OpenMap f
m OpenKey v
key of
    Just f v
x -> f v
x
    Maybe (f v)
Nothing -> [Char] -> f v
forall a. HasCallStack => [Char] -> a
error [Char]
"oh fuck"

-- | If key belongs to @s@, @tryLookup@ will return a proof of this fact
-- and a corresponding value from the map. Otherwise returns @Nothing@.
tryLookup :: NodeMap s f -> OpenKey x -> Maybe (NodeKey s x, f x)
tryLookup :: NodeMap s f -> OpenKey x -> Maybe (NodeKey s x, f x)
tryLookup (NodeMap OpenMap f
m) OpenKey x
key =
  case OpenMap f -> OpenKey x -> Maybe (f x)
forall (f :: * -> *) x. OpenMap f -> OpenKey x -> Maybe (f x)
OpenMap.lookup OpenMap f
m OpenKey x
key of
    Just f x
x -> (NodeKey s x, f x) -> Maybe (NodeKey s x, f x)
forall a. a -> Maybe a
Just (OpenKey x -> NodeKey s x
forall s x. OpenKey x -> NodeKey s x
NodeKey OpenKey x
key, f x
x)
    Maybe (f x)
Nothing -> Maybe (NodeKey s x, f x)
forall a. Maybe a
Nothing

generate :: forall s f. IsNodeSet s => (forall x. NodeKey s x -> f x) -> NodeMap s f
generate :: (forall x. NodeKey s x -> f x) -> NodeMap s f
generate forall x. NodeKey s x -> f x
f = case IsNodeSet s => OpenMap Proxy
forall s. IsNodeSet s => OpenMap Proxy
allNodes @s of
  OpenMap Proxy
m -> (forall x. NodeKey s x -> Proxy x -> f x)
-> NodeMap s Proxy -> NodeMap s f
forall s (f :: * -> *) (g :: * -> *).
(forall x. NodeKey s x -> f x -> g x) -> NodeMap s f -> NodeMap s g
mapWithKey (\NodeKey s x
key Proxy x
_ -> NodeKey s x -> f x
forall x. NodeKey s x -> f x
f NodeKey s x
key) (OpenMap Proxy -> NodeMap s Proxy
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap OpenMap Proxy
m)

zipWith :: forall s f g h. (forall x. f x -> g x -> h x) -> NodeMap s f -> NodeMap s g -> NodeMap s h
zipWith :: (forall x. f x -> g x -> h x)
-> NodeMap s f -> NodeMap s g -> NodeMap s h
zipWith forall x. f x -> g x -> h x
f (NodeMap OpenMap f
x) (NodeMap OpenMap g
y) = OpenMap h -> NodeMap s h
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap ((forall x. f x -> g x -> h x)
-> OpenMap f -> OpenMap g -> OpenMap h
forall (f :: * -> *) (g :: * -> *) (h :: * -> *).
(forall x. f x -> g x -> h x)
-> OpenMap f -> OpenMap g -> OpenMap h
OpenMap.intersectionWith forall x. f x -> g x -> h x
f OpenMap f
x OpenMap g
y)

adjust :: forall s f x. (f x -> f x) -> NodeKey s x -> NodeMap s f -> NodeMap s f
adjust :: (f x -> f x) -> NodeKey s x -> NodeMap s f -> NodeMap s f
adjust f x -> f x
f (NodeKey OpenKey x
key) (NodeMap OpenMap f
m) = OpenMap f -> NodeMap s f
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap ((f x -> f x) -> OpenKey x -> OpenMap f -> OpenMap f
forall (f :: * -> *) x.
(f x -> f x) -> OpenKey x -> OpenMap f -> OpenMap f
OpenMap.adjust f x -> f x
f OpenKey x
key OpenMap f
m)

data NodeSetWrapper s

instance Reifies s (OpenMap Proxy) => IsNodeSet (NodeSetWrapper s) where
  allNodes :: OpenMap Proxy
allNodes = Proxy s -> OpenMap Proxy
forall k (s :: k) a (proxy :: k -> *). Reifies s a => proxy s -> a
reflect @s Proxy s
forall k (t :: k). Proxy t
Proxy

-- | 'NodeMap' with existential set of nodes.
data SomeNodeMap f where
  SomeNodeMap :: IsNodeSet s => NodeMap s f -> SomeNodeMap f

fromOpenMap :: forall f. OpenMap f -> SomeNodeMap f
fromOpenMap :: OpenMap f -> SomeNodeMap f
fromOpenMap OpenMap f
x = OpenMap Proxy
-> (forall s.
    Reifies s (OpenMap Proxy) =>
    Proxy s -> SomeNodeMap f)
-> SomeNodeMap f
forall a r. a -> (forall s. Reifies s a => Proxy s -> r) -> r
reify OpenMap Proxy
nodes forall s. Reifies s (OpenMap Proxy) => Proxy s -> SomeNodeMap f
go
  where
    nodes :: OpenMap Proxy
    nodes :: OpenMap Proxy
nodes = (forall x. f x -> Proxy x) -> OpenMap f -> OpenMap Proxy
forall (f :: * -> *) (g :: * -> *).
(forall x. f x -> g x) -> OpenMap f -> OpenMap g
OpenMap.map (Proxy x -> f x -> Proxy x
forall a b. a -> b -> a
const Proxy x
forall k (t :: k). Proxy t
Proxy) OpenMap f
x
    go :: forall s. Reifies s (OpenMap Proxy) => Proxy s -> SomeNodeMap f
    go :: Proxy s -> SomeNodeMap f
go Proxy s
_proxy = NodeMap (NodeSetWrapper s) f -> SomeNodeMap f
forall s (f :: * -> *). IsNodeSet s => NodeMap s f -> SomeNodeMap f
SomeNodeMap @(NodeSetWrapper s) (OpenMap f -> NodeMap (NodeSetWrapper s) f
forall s (f :: * -> *). OpenMap f -> NodeMap s f
NodeMap OpenMap f
x)