{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Data.TypeMap.Internal.Map where

import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap

import Data.TypeMap.Internal.Unsafe

-- | IntMap-backed type-map.
data TypeMap d = TypeMap !Int !(TypeMap' d) !Int

newtype TypeMap' d = TypeMap' (IntMap Any)

type role TypeMap' nominal

-- | Empty vector.
empty :: TypeMap '[]
empty :: TypeMap '[]
empty = Int -> TypeMap' '[] -> Int -> TypeMap '[]
forall k (d :: k). Int -> TypeMap' d -> Int -> TypeMap d
TypeMap Int
0 (IntMap Any -> TypeMap' '[]
forall k (d :: k). IntMap Any -> TypeMap' d
TypeMap' IntMap Any
forall a. IntMap a
IntMap.empty) Int
0


-- | Access an element indexed by type @a@. /O(log n)/
--
-- If @a@ is associated to @b@ in the type list @d@:
--
-- @
-- 'index' @a (v :: 'TypeMap' d) :: b
-- @
--
-- >>> let v = (0 :: Int) <| True <| "Hello" <| empty :: TypeMap '[ '("a", Int), '("b", Bool), '("c", String)]
-- >>> index @"c" v
-- "Hello"
index
  :: forall a d
  .  KnownNat (Index a d)
  => TypeMap d -> Lookup a d
index :: TypeMap d -> Lookup a d
index (TypeMap Int
inf TypeMap' d
m1 Int
_) = (forall c. IntMap c -> Int -> c) -> TypeMap' d -> Lookup a d
forall k (a :: k) (d :: [(k, *)]) (f :: * -> *)
       (m :: [(k, *)] -> *).
(KnownNat (Index a d), Coercible (f Any) (m d)) =>
(forall c. f c -> Int -> c) -> m d -> Lookup a d
unsafeIndex @a @d forall c. IntMap c -> Int -> c
index' TypeMap' d
m1
  where
    index' :: forall c. IntMap c -> Int -> c
    index' :: IntMap c -> Int -> c
index' IntMap c
m Int
n = IntMap c
m IntMap c -> Int -> c
forall c. IntMap c -> Int -> c
IntMap.! (Int
inf Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)

-- | Add an element to the beginning of a map. /O(log n)/
cons
  :: forall a d b
  .  b -> TypeMap d -> TypeMap ('(a, b) ': d)
cons :: b -> TypeMap d -> TypeMap ('(a, b) : d)
cons b
b (TypeMap Int
inf TypeMap' d
m1 Int
sup) =
  Int -> TypeMap' ('(a, b) : d) -> Int -> TypeMap ('(a, b) : d)
forall k (d :: k). Int -> TypeMap' d -> Int -> TypeMap d
TypeMap (Int
inf Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ((forall c. c -> IntMap c -> IntMap c)
-> b -> TypeMap' d -> TypeMap' ('(a, b) : d)
forall k (a :: k) (d :: [(k, *)]) b (f :: * -> *)
       (m :: [(k, *)] -> *).
(Coercible (f Any) (m d), Coercible (f Any) (m ('(a, b) : d))) =>
(forall c. c -> f c -> f c) -> b -> m d -> m ('(a, b) : d)
unsafeCons forall c. c -> IntMap c -> IntMap c
cons' b
b TypeMap' d
m1) Int
sup
  where
    cons' :: forall c. c -> IntMap c -> IntMap c
    cons' :: c -> IntMap c -> IntMap c
cons' = Int -> c -> IntMap c -> IntMap c
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert (Int
inf Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

-- | Synonym of 'cons'.
(<|)
  :: forall a d b
  .  b -> TypeMap d -> TypeMap ('(a, b) ': d)
<| :: b -> TypeMap d -> TypeMap ('(a, b) : d)
(<|) = b -> TypeMap d -> TypeMap ('(a, b) : d)
forall k (a :: k) (d :: [(k, *)]) b.
b -> TypeMap d -> TypeMap ('(a, b) : d)
cons

infixr 5 <|, `cons`

-- | Add an element to the end of a list. /O(log n)/
snoc
  :: forall a d b
  .  (Last d ~ '(a, b))
  => TypeMap (Init d) -> b -> TypeMap d
snoc :: TypeMap (Init d) -> b -> TypeMap d
snoc (TypeMap Int
inf TypeMap' (Init d)
m1 Int
sup) b
b1 = Int -> TypeMap' d -> Int -> TypeMap d
forall k (d :: k). Int -> TypeMap' d -> Int -> TypeMap d
TypeMap Int
inf ((forall c. IntMap c -> c -> IntMap c)
-> TypeMap' (Init d) -> b -> TypeMap' d
forall k (a :: k) (d :: [(k, *)]) b (f :: * -> *)
       (m :: [(k, *)] -> *).
(Last d ~ '(a, b), Coercible (f Any) (m (Init d)),
 Coercible (f Any) (m d)) =>
(forall c. f c -> c -> f c) -> m (Init d) -> b -> m d
unsafeSnoc @a @d @b forall c. IntMap c -> c -> IntMap c
snoc' TypeMap' (Init d)
m1 b
b1) (Int
sup Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  where
    snoc' :: forall c. IntMap c -> c -> IntMap c
    snoc' :: IntMap c -> c -> IntMap c
snoc' IntMap c
m c
b = Int -> c -> IntMap c -> IntMap c
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
sup c
b IntMap c
m

-- | Synonym of 'snoc'.
(|>)
  :: forall a d b
  .  (Last d ~ '(a, b))
  => TypeMap (Init d) -> b -> TypeMap d
|> :: TypeMap (Init d) -> b -> TypeMap d
(|>) = TypeMap (Init d) -> b -> TypeMap d
forall k (a :: k) (d :: [(k, *)]) b.
(Last d ~ '(a, b)) =>
TypeMap (Init d) -> b -> TypeMap d
snoc

infixl 5 |>, `snoc`