{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
module Data.Parameterized.FinMap.Safe
( FinMap
, null
, lookup
, size
, incMax
, embed
, empty
, singleton
, insert
, buildFinMap
, append
, fromVector
, delete
, decMax
, mapWithKey
, unionWithKey
, unionWith
, union
) where
import Prelude hiding (lookup, null)
import Data.Foldable.WithIndex (FoldableWithIndex(ifoldMap))
import Data.Functor.WithIndex (FunctorWithIndex(imap))
import Data.Maybe (isJust)
import Data.Proxy (Proxy(Proxy))
import Data.Map (Map)
import qualified Data.Map as Map
import GHC.TypeLits (KnownNat, Nat)
import Data.Parameterized.Fin (Fin)
import qualified Data.Parameterized.Fin as Fin
import Data.Parameterized.NatRepr (NatRepr, type (+), type (<=))
import qualified Data.Parameterized.NatRepr as NatRepr
import Data.Parameterized.Vector (Vector)
import qualified Data.Parameterized.Vector as Vec
data FinMap (n :: Nat) a =
FinMap
{ forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap :: Map (Fin n) a
, forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize :: NatRepr n
}
instance Eq a => Eq (FinMap n a) where
FinMap n a
fm1 == :: FinMap n a -> FinMap n a -> Bool
== FinMap n a
fm2 = FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm1 Map (Fin n) a -> Map (Fin n) a -> Bool
forall a. Eq a => a -> a -> Bool
== FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm2
{-# INLINABLE (==) #-}
instance Semigroup (FinMap n a) where
<> :: FinMap n a -> FinMap n a -> FinMap n a
(<>) = FinMap n a -> FinMap n a -> FinMap n a
forall (n :: Nat) a. FinMap n a -> FinMap n a -> FinMap n a
union
{-# INLINE (<>) #-}
instance KnownNat n => Monoid (FinMap n a) where
mempty :: FinMap n a
mempty = FinMap n a
forall (n :: Nat) a. KnownNat n => FinMap n a
empty
{-# INLINE mempty #-}
instance Functor (FinMap n) where
fmap :: forall a b. (a -> b) -> FinMap n a -> FinMap n b
fmap a -> b
f FinMap n a
fm = FinMap n a
fm { getFinMap = fmap f (getFinMap fm) }
{-# INLINABLE fmap #-}
instance Foldable (FinMap n) where
foldMap :: forall m a. Monoid m => (a -> m) -> FinMap n a -> m
foldMap a -> m
f = (a -> m) -> Map (Fin n) a -> m
forall m a. Monoid m => (a -> m) -> Map (Fin n) a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
f (Map (Fin n) a -> m)
-> (FinMap n a -> Map (Fin n) a) -> FinMap n a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap
{-# INLINABLE foldMap #-}
instance Traversable (FinMap n) where
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> FinMap n a -> f (FinMap n b)
traverse a -> f b
f FinMap n a
fm = Map (Fin n) b -> NatRepr n -> FinMap n b
forall (n :: Nat) a. Map (Fin n) a -> NatRepr n -> FinMap n a
FinMap (Map (Fin n) b -> NatRepr n -> FinMap n b)
-> f (Map (Fin n) b) -> f (NatRepr n -> FinMap n b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Map (Fin n) a -> f (Map (Fin n) b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map (Fin n) a -> f (Map (Fin n) b)
traverse a -> f b
f (FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm) f (NatRepr n -> FinMap n b) -> f (NatRepr n) -> f (FinMap n b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> NatRepr n -> f (NatRepr n)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FinMap n a -> NatRepr n
forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize FinMap n a
fm)
instance FunctorWithIndex (Fin n) (FinMap n) where
imap :: forall a b. (Fin n -> a -> b) -> FinMap n a -> FinMap n b
imap Fin n -> a -> b
f FinMap n a
fm = FinMap n a
fm { getFinMap = Map.mapWithKey f (getFinMap fm) }
{-# INLINE imap #-}
instance FoldableWithIndex (Fin n) (FinMap n) where
ifoldMap :: forall m a. Monoid m => (Fin n -> a -> m) -> FinMap n a -> m
ifoldMap Fin n -> a -> m
f = (Fin n -> a -> m) -> Map (Fin n) a -> m
forall m k a. Monoid m => (k -> a -> m) -> Map k a -> m
Map.foldMapWithKey Fin n -> a -> m
f (Map (Fin n) a -> m)
-> (FinMap n a -> Map (Fin n) a) -> FinMap n a -> m
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap
{-# INLINABLE ifoldMap #-}
instance Show a => Show (FinMap n a) where
show :: FinMap n a -> String
show FinMap n a
fm = Map (Fin n) a -> String
forall a. Show a => a -> String
show (FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm)
{-# INLINABLE show #-}
null :: FinMap n a -> Bool
null :: forall (n :: Nat) a. FinMap n a -> Bool
null = Map (Fin n) a -> Bool
forall k a. Map k a -> Bool
Map.null (Map (Fin n) a -> Bool)
-> (FinMap n a -> Map (Fin n) a) -> FinMap n a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap
{-# INLINABLE null #-}
lookup :: Fin n -> FinMap n a -> Maybe a
lookup :: forall (n :: Nat) a. Fin n -> FinMap n a -> Maybe a
lookup Fin n
k = Fin n -> Map (Fin n) a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Fin n
k (Map (Fin n) a -> Maybe a)
-> (FinMap n a -> Map (Fin n) a) -> FinMap n a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap
{-# INLINABLE lookup #-}
size :: forall n a. FinMap n a -> Fin (n + 1)
size :: forall (n :: Nat) a. FinMap n a -> Fin (n + 1)
size FinMap n a
fm =
NatRepr n
-> (forall (n :: Nat).
((n + 1) <= n) =>
NatRepr n -> Fin (n + 1) -> Bool)
-> Fin (n + 1)
forall (m :: Nat).
NatRepr m
-> (forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> Fin (n + 1) -> Bool)
-> Fin (m + 1)
Fin.countFin (FinMap n a -> NatRepr n
forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize FinMap n a
fm) (\NatRepr n
k Fin (n + 1)
_count -> Maybe a -> Bool
forall a. Maybe a -> Bool
isJust (Fin n -> FinMap n a -> Maybe a
forall (n :: Nat) a. Fin n -> FinMap n a -> Maybe a
lookup (NatRepr n -> Fin n
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
Fin.mkFin NatRepr n
k) FinMap n a
fm))
incMax :: forall n a. FinMap n a -> FinMap (n + 1) a
incMax :: forall (n :: Nat) a. FinMap n a -> FinMap (n + 1) a
incMax FinMap n a
fm =
case Proxy n -> LeqProof n (n + 1)
forall (f :: Nat -> *) (z :: Nat). f z -> LeqProof z (z + 1)
NatRepr.leqSucc (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n) of
LeqProof n (n + 1)
NatRepr.LeqProof -> NatRepr (n + 1) -> FinMap n a -> FinMap (n + 1) a
forall (n :: Nat) (m :: Nat) a.
(n <= m) =>
NatRepr m -> FinMap n a -> FinMap m a
embed (NatRepr n -> NatRepr (n + 1)
forall (n :: Nat). NatRepr n -> NatRepr (n + 1)
NatRepr.incNat (FinMap n a -> NatRepr n
forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize FinMap n a
fm)) FinMap n a
fm
embed :: (n <= m) => NatRepr m -> FinMap n a -> FinMap m a
embed :: forall (n :: Nat) (m :: Nat) a.
(n <= m) =>
NatRepr m -> FinMap n a -> FinMap m a
embed NatRepr m
m FinMap n a
fm =
FinMap
{ getFinMap :: Map (Fin m) a
getFinMap = (Fin n -> Fin m) -> Map (Fin n) a -> Map (Fin m) a
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys Fin n -> Fin m
forall (n :: Nat) (m :: Nat). (n <= m) => Fin n -> Fin m
Fin.embed (FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm)
, maxSize :: NatRepr m
maxSize = NatRepr m
m
}
empty :: KnownNat n => FinMap n a
empty :: forall (n :: Nat) a. KnownNat n => FinMap n a
empty = Map (Fin n) a -> NatRepr n -> FinMap n a
forall (n :: Nat) a. Map (Fin n) a -> NatRepr n -> FinMap n a
FinMap Map (Fin n) a
forall k a. Map k a
Map.empty NatRepr n
forall (n :: Nat). KnownNat n => NatRepr n
NatRepr.knownNat
{-# INLINABLE empty #-}
singleton :: a -> FinMap 1 a
singleton :: forall a. a -> FinMap 1 a
singleton a
item =
FinMap
{ getFinMap :: Map (Fin 1) a
getFinMap = Fin 1 -> a -> Map (Fin 1) a
forall k a. k -> a -> Map k a
Map.singleton (NatRepr 0 -> Fin 1
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
Fin.mkFin (NatRepr 0
forall (n :: Nat). KnownNat n => NatRepr n
NatRepr.knownNat :: NatRepr 0)) a
item
, maxSize :: NatRepr 1
maxSize = NatRepr 1
forall (n :: Nat). KnownNat n => NatRepr n
NatRepr.knownNat :: NatRepr 1
}
insert :: Fin n -> a -> FinMap n a -> FinMap n a
insert :: forall (n :: Nat) a. Fin n -> a -> FinMap n a -> FinMap n a
insert Fin n
k a
v FinMap n a
fm = FinMap n a
fm { getFinMap = Map.insert k v (getFinMap fm) }
{-# INLINABLE insert #-}
newtype FinMap' a (n :: Nat) = FinMap' { forall a (n :: Nat). FinMap' a n -> FinMap n a
unFinMap' :: FinMap n a }
buildFinMap ::
forall m a.
NatRepr m ->
(forall n. (n + 1 <= m) => NatRepr n -> FinMap n a -> FinMap (n + 1) a) ->
FinMap m a
buildFinMap :: forall (m :: Nat) a.
NatRepr m
-> (forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> FinMap n a -> FinMap (n + 1) a)
-> FinMap m a
buildFinMap NatRepr m
m forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> FinMap n a -> FinMap (n + 1) a
f =
let f' :: forall k. (k + 1 <= m) => NatRepr k -> FinMap' a k -> FinMap' a (k + 1)
f' :: forall (k :: Nat).
((k + 1) <= m) =>
NatRepr k -> FinMap' a k -> FinMap' a (k + 1)
f' = (\NatRepr k
n (FinMap' FinMap k a
fin) -> FinMap (k + 1) a -> FinMap' a (k + 1)
forall a (n :: Nat). FinMap n a -> FinMap' a n
FinMap' (NatRepr k -> FinMap k a -> FinMap (k + 1) a
forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> FinMap n a -> FinMap (n + 1) a
f NatRepr k
n FinMap k a
fin))
in FinMap' a m -> FinMap m a
forall a (n :: Nat). FinMap' a n -> FinMap n a
unFinMap' (NatRepr m
-> FinMap' a 0
-> (forall (k :: Nat).
((k + 1) <= m) =>
NatRepr k -> FinMap' a k -> FinMap' a (k + 1))
-> FinMap' a m
forall (m :: Nat) (f :: Nat -> *).
NatRepr m
-> f 0
-> (forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> f n -> f (n + 1))
-> f m
NatRepr.natRecStrictlyBounded NatRepr m
m (FinMap 0 a -> FinMap' a 0
forall a (n :: Nat). FinMap n a -> FinMap' a n
FinMap' FinMap 0 a
forall (n :: Nat) a. KnownNat n => FinMap n a
empty) NatRepr n -> FinMap' a n -> FinMap' a (n + 1)
forall (k :: Nat).
((k + 1) <= m) =>
NatRepr k -> FinMap' a k -> FinMap' a (k + 1)
f')
append :: NatRepr n -> a -> FinMap n a -> FinMap (n + 1) a
append :: forall (n :: Nat) a.
NatRepr n -> a -> FinMap n a -> FinMap (n + 1) a
append NatRepr n
k a
v FinMap n a
fm =
case NatRepr n -> LeqProof n (n + 1)
forall (f :: Nat -> *) (z :: Nat). f z -> LeqProof z (z + 1)
NatRepr.leqSucc NatRepr n
k of
LeqProof n (n + 1)
NatRepr.LeqProof -> Fin (n + 1) -> a -> FinMap (n + 1) a -> FinMap (n + 1) a
forall (n :: Nat) a. Fin n -> a -> FinMap n a -> FinMap n a
insert (NatRepr n -> Fin (n + 1)
forall (i :: Nat) (n :: Nat). ((i + 1) <= n) => NatRepr i -> Fin n
Fin.mkFin NatRepr n
k) a
v (FinMap n a -> FinMap (n + 1) a
forall (n :: Nat) a. FinMap n a -> FinMap (n + 1) a
incMax FinMap n a
fm)
fromVector :: forall n a. Vector n (Maybe a) -> FinMap n a
fromVector :: forall (n :: Nat) a. Vector n (Maybe a) -> FinMap n a
fromVector Vector n (Maybe a)
v =
NatRepr n
-> (forall (n :: Nat).
((n + 1) <= n) =>
NatRepr n -> FinMap n a -> FinMap (n + 1) a)
-> FinMap n a
forall (m :: Nat) a.
NatRepr m
-> (forall (n :: Nat).
((n + 1) <= m) =>
NatRepr n -> FinMap n a -> FinMap (n + 1) a)
-> FinMap m a
buildFinMap
(Vector n (Maybe a) -> NatRepr n
forall (n :: Nat) a. Vector n a -> NatRepr n
Vec.length Vector n (Maybe a)
v)
(\NatRepr n
k FinMap n a
m ->
case NatRepr n -> Vector n (Maybe a) -> Maybe a
forall (i :: Nat) (n :: Nat) a.
((i + 1) <= n) =>
NatRepr i -> Vector n a -> a
Vec.elemAt NatRepr n
k Vector n (Maybe a)
v of
Just a
e -> NatRepr n -> a -> FinMap n a -> FinMap (n + 1) a
forall (n :: Nat) a.
NatRepr n -> a -> FinMap n a -> FinMap (n + 1) a
append NatRepr n
k a
e FinMap n a
m
Maybe a
Nothing -> FinMap n a -> FinMap (n + 1) a
forall (n :: Nat) a. FinMap n a -> FinMap (n + 1) a
incMax FinMap n a
m)
delete :: Fin n -> FinMap n a -> FinMap n a
delete :: forall (n :: Nat) a. Fin n -> FinMap n a -> FinMap n a
delete Fin n
k FinMap n a
fm = FinMap n a
fm { getFinMap = Map.delete k (getFinMap fm) }
{-# INLINABLE delete #-}
decMax :: NatRepr n -> FinMap (n + 1) a -> FinMap n a
decMax :: forall (n :: Nat) a. NatRepr n -> FinMap (n + 1) a -> FinMap n a
decMax NatRepr n
n FinMap (n + 1) a
fm =
FinMap
{ getFinMap :: Map (Fin n) a
getFinMap = (Fin (n + 1) -> Maybe (Fin n))
-> Map (Fin (n + 1)) a -> Map (Fin n) a
forall k2 k1 a. Ord k2 => (k1 -> Maybe k2) -> Map k1 a -> Map k2 a
maybeMapKeys (NatRepr (n + 1) -> NatRepr n -> Fin (n + 1) -> Maybe (Fin n)
forall (n :: Nat) (m :: Nat).
NatRepr n -> NatRepr m -> Fin n -> Maybe (Fin m)
Fin.tryEmbed NatRepr (n + 1)
sz NatRepr n
n) (FinMap (n + 1) a -> Map (Fin (n + 1)) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap (n + 1) a
fm)
, maxSize :: NatRepr n
maxSize = NatRepr n
n
}
where
sz :: NatRepr (n + 1)
sz = FinMap (n + 1) a -> NatRepr (n + 1)
forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize FinMap (n + 1) a
fm
maybeMapKeys :: Ord k2 => (k1 -> Maybe k2) -> Map k1 a -> Map k2 a
maybeMapKeys :: forall k2 k1 a. Ord k2 => (k1 -> Maybe k2) -> Map k1 a -> Map k2 a
maybeMapKeys k1 -> Maybe k2
f Map k1 a
m =
(k1 -> a -> Map k2 a -> Map k2 a)
-> Map k2 a -> Map k1 a -> Map k2 a
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
Map.foldrWithKey
(\k1
k a
v Map k2 a
accum ->
case k1 -> Maybe k2
f k1
k of
Just k2
k' -> k2 -> a -> Map k2 a -> Map k2 a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k2
k' a
v Map k2 a
accum
Maybe k2
Nothing -> Map k2 a
accum)
Map k2 a
forall k a. Map k a
Map.empty
Map k1 a
m
mapWithKey :: (Fin n -> a -> b) -> FinMap n a -> FinMap n b
mapWithKey :: forall (n :: Nat) a b.
(Fin n -> a -> b) -> FinMap n a -> FinMap n b
mapWithKey Fin n -> a -> b
f FinMap n a
fm = FinMap n a
fm { getFinMap = Map.mapWithKey f (getFinMap fm) }
{-# INLINE mapWithKey #-}
unionWithKey :: (Fin n -> a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWithKey :: forall (n :: Nat) a.
(Fin n -> a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWithKey Fin n -> a -> a -> a
f FinMap n a
fm1 FinMap n a
fm2 =
FinMap
{ getFinMap :: Map (Fin n) a
getFinMap = (Fin n -> a -> a -> a)
-> Map (Fin n) a -> Map (Fin n) a -> Map (Fin n) a
forall k a.
Ord k =>
(k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWithKey Fin n -> a -> a -> a
f (FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm1) (FinMap n a -> Map (Fin n) a
forall (n :: Nat) a. FinMap n a -> Map (Fin n) a
getFinMap FinMap n a
fm2)
, maxSize :: NatRepr n
maxSize = FinMap n a -> NatRepr n
forall (n :: Nat) a. FinMap n a -> NatRepr n
maxSize FinMap n a
fm1
}
unionWith :: (a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWith :: forall a (n :: Nat).
(a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWith a -> a -> a
f = (Fin n -> a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
forall (n :: Nat) a.
(Fin n -> a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWithKey (\Fin n
_ a
v1 a
v2 -> a -> a -> a
f a
v1 a
v2)
union :: FinMap n a -> FinMap n a -> FinMap n a
union :: forall (n :: Nat) a. FinMap n a -> FinMap n a -> FinMap n a
union = (a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
forall a (n :: Nat).
(a -> a -> a) -> FinMap n a -> FinMap n a -> FinMap n a
unionWith a -> a -> a
forall a b. a -> b -> a
const