Singleton maps

SingMap.hs

{-# language TypeInType, RankNTypes, ScopedTypeVariables, TypeFamilies, TypeOperators,
  GADTs, UndecidableInstances, NoStarIsType, TemplateHaskell, InstanceSigs,
  TypeSynonymInstances, FlexibleInstances, BangPatterns #-}

module SingMap (module SingMap, module P, module F) where
import GHC.TypeLits hiding (type (<=))
import Data.Type.Bool
import Data.Type.Equality
import Data.Singletons
import Data.Traversable
import Data.Singletons.Prelude hiding (type Map, type Lookup, sLookup, type InsertSym0, type InsertSym1)
import Data.Singletons.TH
import qualified Data.Map.Lazy as M
import qualified Data.Map.Internal as MI
import Numeric.Natural
import qualified Data.Singletons.Prelude as P
import qualified Data.Singletons.Prelude.Foldable as F


$(promote [d|
 data Map k a
   = Bin Nat k a (Map k a) (Map k a)
   | Tip
 |])

data instance Sing (x :: Map k a) where
  SBin :: Sing n -> Sing k -> Sing a -> Sing l -> Sing r -> Sing ('Bin n k a l r)
  STip :: Sing 'Tip
  
instance SingI 'Tip where
  sing = STip
instance (SingI s, SingI k, SingI a, SingI l, SingI r) => SingI ('Bin s k a l r) where
  sing = SBin sing sing sing sing sing


instance (SingKind k, SingKind a) => SingKind (Map k a) where
  type Demote (Map k a) = M.Map (Demote k) (Demote a)
  fromSing (SBin n k a l r) = MI.Bin (fromIntegral $ fromSing n) (fromSing k) (fromSing a) (fromSing l) (fromSing r)
  fromSing STip = MI.Tip

  toSing (MI.Bin s k a l r)
    | SomeSing s' <- toSing (fromIntegral s :: Natural)
    , SomeSing k' <- toSing k
    , SomeSing a' <- toSing a
    , SomeSing l' <- toSing l
    , SomeSing r' <- toSing r
    = SomeSing (SBin s' k' a' l' r')
  toSing MI.Tip = SomeSing STip

$(singletons [d|
 delta, ratio :: Nat
 delta = 3
 ratio = 2

 empty :: Map k v
 empty = Tip

 size :: Map k v -> Nat
 size Tip = 0
 size (Bin s _ _ _ _) = s

 singleton :: k -> v -> Map k v
 singleton k v = Bin 1 k v Tip Tip

 lookup :: Ord k => k -> Map k v -> Maybe v
 lookup _ Tip = Nothing
 lookup kx (Bin s ky y l r) =
   case compare kx ky of
     LT -> lookup kx l
     GT -> lookup kx r
     EQ -> Just y

 instance Foldable (Map k) where
   foldMap _ Tip = mempty
   foldMap f (Bin _ _k v l r) = foldMap f l <> f v <> foldMap f r

 instance Functor (Map k) where
   fmap _ Tip = Tip
   fmap f (Bin s k v l r) = Bin s k (f v) (fmap f l) (fmap f r)

 balanceL :: k -> a -> Map k a -> Map k a -> Map k a
 balanceL k x l r = case r of
   Tip -> case l of
            Tip -> Bin 1 k x Tip Tip
            (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip
            (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip)
            (Bin _ lk lx ll@(Bin _ _ _ _ _) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip)
            (Bin ls lk lx ll@(Bin lls _ _ _ _) lr@(Bin lrs lrk lrx lrl lrr))
              -> if lrs < ratio*lls
                 then Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip)
                 else Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip)

   (Bin rs _ _ _ _) -> case l of
            Tip -> Bin (1+rs) k x Tip r

            (Bin ls lk lx ll lr)
               -> if ls > delta*rs
                  then case (ll, lr) of
                    (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr)
                      -> if lrs < ratio*lls
                         then Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r)
                         else Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r)
                    (Bin _ _ _ _ _, Tip) -> error "oops"
                    (Tip, Bin _ _ _ _ _) -> error "oops"
                    (Tip, Tip) -> error "oops"
                  else Bin (1+ls+rs) k x l r
                  
 balanceR :: k -> a -> Map k a -> Map k a -> Map k a
 balanceR k x l r = case l of
   Tip -> case r of
            Tip -> Bin 1 k x Tip Tip
            (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r
            (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr
            (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip)
            (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _))
              -> if rls < ratio*rrs
                 then Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr
                 else Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr)

   (Bin ls _ _ _ _) -> case r of
            Tip -> Bin (1+ls) k x l Tip

            (Bin rs rk rx rl rr)
               -> if rs > delta*ls
                  then case (rl, rr) of
                    (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _)
                      -> if rls < ratio*rrs
                         then Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr
                         else Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr)
                    (Bin _ _ _ _ _, Tip) -> error "oops"
                    (Tip, Bin _ _ _ _ _) -> error "oops"
                    (Tip, Tip) -> error "oops"
                  else Bin (1+ls+rs) k x l r
   
 insert :: Ord k => k -> a -> Map k a -> Map k a
 insert = go
   where
     go :: Ord k => k -> a -> Map k a -> Map k a
     go !kx x Tip = singleton kx x
     go !kx x t@(Bin sz ky y l r) =
         case compare kx ky of
             LT -> balanceL ky y l' r
                where !l' = go kx x l
             GT -> balanceR ky y l r'
                where !r' = go kx x r
             EQ -> Bin sz kx x l r

 fromList :: Ord k => [(k, v)] -> Map k v
 fromList [] = empty
 fromList ((k,v) : kvs) = insert k v (fromList kvs)

 assocs :: Map k v -> [(k, v)]
 assocs xs0 = go xs0 []
   where
     go Tip rest = rest
     go (Bin _ k v l r) rest = go l ((k,v) : go r rest)

 smallest :: Map k v -> Maybe (k, v)
 smallest Tip = Nothing
 smallest (Bin _ k v Tip _) = Just (k, v)
 smallest (Bin _ _ _ (Bin s k v l r) _) = smallest (Bin s k v l r)
 |])