Fast Map Union and Local Instances Through Instance Types

insttypes.hs

{-# LANGUAGE GHC2021, FunctionalDependencies, AllowAmbiguousTypes, OverloadedRecordDot, BlockArguments #-}
module Lib where

-- Code accompanying https://prophetlabs.de/posts/insttypes.html

import qualified Data.Map as Map

import Unsafe.Coerce

import Data.Proxy
import Data.Kind

class OrdI (inst :: Type) a | inst -> a where
    compareI :: a -> a -> Ordering

newtype OrdFor inst a = MkOrdFor {unOrdFor :: a}

instance (OrdI inst a) => Eq (OrdFor inst a) where
    (MkOrdFor x) == (MkOrdFor y) =
        compareI @inst x y == EQ

instance (OrdI inst a) => Ord (OrdFor inst a) where
    compare (MkOrdFor x) (MkOrdFor y) = compareI @inst x y


data IMap inst k v = MkIMap {
    underlying :: Map.Map (OrdFor inst k) v
}

empty :: IMap inst k v
empty = MkIMap Map.empty

insert :: forall inst k v. OrdI inst k => k -> v -> IMap inst k v -> IMap inst k v
insert k v (MkIMap map) = MkIMap { underlying = Map.insert (MkOrdFor @inst k) v map }

union :: (Eq k, OrdI inst k) => IMap inst k v -> IMap inst k v -> IMap inst k v
union map1 map2 = MkIMap { underlying = Map.union map1.underlying map2.underlying }

data RegularIntOrd
instance OrdI RegularIntOrd Int where
    compareI = compare


data ReverseIntOrd
instance OrdI ReverseIntOrd Int where
    compareI x y = case compareI @RegularIntOrd x y of
        LT -> GT
        EQ -> EQ
        GT -> LT

data OrdIDict a where
    OrdIDict :: OrdI inst a => OrdIDict a

data FakeDict a = FakeDict a


withOrdI :: forall a b. (a -> a -> Ordering) -> (forall inst. OrdI inst a => b) -> b
withOrdI dict body = 
    case unsafeCoerce (FakeDict dict) :: OrdIDict a of
        (OrdIDict @inst) -> body @inst

withOrdIProxy :: forall a b. (a -> a -> Ordering) -> (forall inst. OrdI inst a => Proxy inst -> b) -> b
withOrdIProxy dict body =
    case unsafeCoerce (FakeDict dict) :: OrdIDict a of
        (OrdIDict @inst) -> body (Proxy @inst)


map1 = insert 1 1 (empty @RegularIntOrd)
map2 = insert 2 2 (empty @RegularIntOrd)

map3 = insert 3 3 (empty @ReverseIntOrd)
map4 = insert 4 4 (empty @ReverseIntOrd)

fine1 = union map1 map2
fine2 = union map3 map4

-- error = union map1 map3

local = withOrdIProxy (\_ _ -> EQ) \(Proxy @inst) -> do
    let localMap1 :: IMap inst Int Int = insert @inst 1 1 (empty @inst)
    -- let localMap2 :: IMap RegularIntOrd Int Int = insert @inst 2 2 (empty @inst) -- fails!
    -- let _ = union fine1 localMap1 -- also fails!
    -- localMap1 -- fails in return position!
    ()