TagMap replacement (WIP)

Separate concerns for `TagMap`, add functionality, and try to
manage the inherent unsafety a bit better. I don't really see
what `MyTag` has to offer by way of optimization, so I haven't
included that business. Maybe I'm missing something there...

src/Data/TagMap.hs

@@ -0,0 +1,187 @@
+-- | This module provides 'TagMap', a version of 'IntMap' for
+-- GADT keys whose constructors can be counted by 'Int'.
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE EmptyDataDecls #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE Rank2Types #-}
+{-# LANGUAGE RecursiveDo #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+#ifdef USE_REFLEX_OPTIMIZER
+{-# OPTIONS_GHC -fplugin=Reflex.Optimizer #-}
+#endif
+module Data.TagMap
+  (
+    TagMap
+  ) where
+
+import Reflex.Class
+import Reflex.Adjustable.Class
+import Reflex.Dynamic
+import Reflex.Host.Class
+import Reflex.PerformEvent.Class
+import Reflex.PostBuild.Class
+import Reflex.Requester.Class
+import Reflex.TriggerEvent.Class
+
+import Control.Applicative (liftA2)
+import Control.Monad.Exception
+import Control.Monad.Identity
+import Control.Monad.Primitive
+import Control.Monad.Reader
+import Control.Monad.Ref
+import Control.Monad.State.Strict
+import Data.Bits
+import Data.Coerce
+import Data.Dependent.Map (DMap, DSum (..))
+import qualified Data.Dependent.Map as DMap
+import Data.Functor.Compose
+import Data.Functor.Misc
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Monoid ((<>))
+import Data.Proxy
+import qualified Data.Semigroup as S
+import Data.Some (Some(Some))
+import Data.Type.Equality
+import Data.Unique.Tag
+
+import GHC.Exts (Any, dataToTag#)
+import Unsafe.Coerce
+
+--TODO: Make this module type-safe
+
+newtype TagMap (k :: x -> *) (v :: x -> *) = TagMap (IntMap (f Any))
+type role TagMap representational representational
+
+class IsTag k where
+  -- For traversing
+  unsafeToKeyValue :: Int -> v Any -> KeyValue k v
+  -- For inspecting just keys. Do we really want to use Some here,
+  -- or are we better off (for performance) using a more legitimate
+  -- Some-like type? I don't think we actually use this yet, so it
+  -- may not matter much.
+  toKey :: Int -> Some k
+  -- For inserting and looking up
+  fromTag :: k a -> Int
+
+data MyTagType :: * -> * where
+  MyTagType_Single :: MyTagType (Single a)
+  MyTagType_Multi :: MyTagType Multi
+  MyTagType_Multi2 :: MyTagType (Multi2 k)
+  MyTagType_Multi3 :: MyTagType Multi3
+
+deriving instance Eq (MyTagType a)
+deriving instance Ord (MyTagType a)
+deriving instance Show (MyTagType a)
+
+instance IsTag MyTagType where
+  unsafeToKeyValue ki va = case ki .&. 0x3 of
+    0x0 -> KeyValue MyTagType_Single (unsafeCoerce va)
+    0x1 -> KeyValue MyTagType_Multi (unsafeCoerce va)
+    0x2 -> KeyValue MyTagType_Multi2 (unsafeCoerce va)
+    0x3 -> KeyValue MyTagType_Multi3 (unsafeCoerce va)
+    t -> error $ "Data.TagMap.unsafeToKeyValue: no such key type" <> show t
+
+  toKey ki = case ki .&. 0x3 of
+    0x0 -> Some MyTagType_Single
+    0x1 -> Some MyTagType_Multi
+    0x2 -> Some MyTagType_Multi2
+    0x3 -> Some MyTagType_Multi3
+    t -> error $ "Data.TagMap.myKeyType: no such key type" <> show t
+
+  fromTag :: MyTagType a -> Int
+  fromTag t = dataToTag# t
+
+toVany :: v a -> v Any
+toVany = unsafeCoerce
+
+empty :: TagMap f
+empty = TagMap IntMap.empty
+
+singleton :: forall f a. IsTag k => k a -> f a -> TagMap k f
+singleton k v = TagMap $ IntMap.singleton (fromTag k) $ toVany v
+
+insert :: IsTag k => k a -> v a -> TagMap k v -> TagMap k v
+insert k v (TagMap m) = TagMap $ IntMap.insert (fromTag k) (toVany v) m
+
+lookup :: IsTag k => k a -> TagMap k v -> Maybe (v a)
+lookup k (TagMap m) = fmap unsafeCoerce $ IntMap.lookup (fromTag k) m
+
+foldrWithKey :: forall k f r. IsTag k => (forall a. k a -> f a -> r -> r) -> r -> TagMap k f -> r
+foldrWithKey f b = \(TagMap m) -> IntMap.foldrWithKey go b m
+  where
+    go :: Int -> f Any -> r -> r
+    go ki fany r
+      | KeyValue k v <- unsafeToKeyValue ki fany
+      = f k v r
+
+data KeyValue k v = forall a. KeyValue !(k a) (v a)
+
+toList :: forall k f. IsTag k => TagMap k f -> [DSum k f]
+toList = foldrWithKey go []
+  where
+    go k v r = (k :=> v) : r
+
+traverseWithKey
+  :: forall k v f g. (IsTag k, Applicative f)
+  => (forall a. k a -> v a -> f (g a)) -> TagMap k v -> f (TagMap k g)
+traverseWithKey f (TagMap m) = TagMap <$> IntMap.traverseWithKey g m
+  where
+    g :: Int -> v Any -> f (g Any)
+    g ki vi
+      | KeyValue k v <- unsafeToKeyValue ki vi
+      = toVany <$> f k v
+
+data Single a
+data Multi
+data Multi2 (k :: * -> *)
+data Multi3
+
+class MyTagTypeOffset x where
+  -- | A type-directed version of `tagOffset`.
+  myTagTypeOffset :: proxy x -> Int
+
+instance MyTagTypeOffset (Single a) where
+  myTagTypeOffset _ = 0x0
+
+instance MyTagTypeOffset Multi where
+  myTagTypeOffset _ = 0x1
+
+instance MyTagTypeOffset (Multi2 k) where
+  myTagTypeOffset _ = 0x2
+
+instance MyTagTypeOffset Multi3 where
+  myTagTypeOffset _ = 0x3
+
+instance GEq MyTagType where
+  geq MyTagType_Single MyTagType_Single = Just Refl
+  geq MyTagType_Multi MyTagType_Multi = Just Refl
+  geq MyTagType_Multi2 MyTagType_Multi2 = Just Refl
+  geq MyTagType_Multi3 MyTagType_Multi3 = Just Refl
+  geq _ _ = Nothing
+
+instance GCompare MyTagType where
+  gcompare MyTagType_Single MyTagType_Single = GEQ
+  gcompare MyTagType_Single _ = GLT
+  gcompare _ MyTagType_Single = GGT
+  gcompare MyTagType_Multi MyTagType_Multi = GEQ
+  gcompare MyTagType_Multi _ = GLT
+  gcompare _ MyTagType_Multi = GGT
+  gcompare MyTagType_Multi2 MyTagType_Multi2 = GEQ
+  gcompare MyTagType_Multi2 _ = GLT
+  gcompare _ MyTagType_Multi2 = GGT
+  gcompare MyTagType_Multi3 MyTagType_Multi3 = GEQ

src/Reflex/Requester/MyTagType.hs

@@ -0,0 +1,133 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE EmptyDataDecls #-}
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE Rank2Types #-}
+{-# LANGUAGE RecursiveDo #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+#ifdef USE_REFLEX_OPTIMIZER
+{-# OPTIONS_GHC -fplugin=Reflex.Optimizer #-}
+#endif
+module Reflex.Requester.MyTagType
+  (
+    MyTagType (..)
+  , Single (..)
+  , Multi (..)
+  , Multi2 (..)
+  , MyTagTypeOffset (..)
+  ) where
+
+import Reflex.Class
+import Reflex.Adjustable.Class
+import Reflex.Dynamic
+import Reflex.Host.Class
+import Reflex.PerformEvent.Class
+import Reflex.PostBuild.Class
+import Reflex.Requester.Class
+import Reflex.TriggerEvent.Class
+
+import Control.Applicative (liftA2)
+import Control.Monad.Exception
+import Control.Monad.Identity
+import Control.Monad.Primitive
+import Control.Monad.Reader
+import Control.Monad.Ref
+import Control.Monad.State.Strict
+import Data.Bits
+import Data.Coerce
+import Data.Dependent.Map (DMap, DSum (..))
+import qualified Data.Dependent.Map as DMap
+import Data.Functor.Compose
+import Data.Functor.Misc
+import Data.IntMap.Strict (IntMap)
+import qualified Data.IntMap.Strict as IntMap
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Monoid ((<>))
+import Data.Proxy
+import qualified Data.Semigroup as S
+import Data.Some (Some(Some))
+import Data.Type.Equality
+import Data.Unique.Tag
+
+import GHC.Exts (Any, dataToTag#, I#)
+import Unsafe.Coerce
+
+data MyTagType :: * -> * where
+  MyTagType_Single :: MyTagType (Single a)
+  MyTagType_Multi :: MyTagType Multi
+  MyTagType_Multi2 :: MyTagType (Multi2 k)
+  MyTagType_Multi3 :: MyTagType Multi3
+
+deriving instance Eq (MyTagType a)
+deriving instance Ord (MyTagType a)
+deriving instance Show (MyTagType a)
+
+instance IsTag MyTagType where
+  unsafeToKeyValue ki va = case ki .&. 0x3 of
+    0x0 -> KeyValue MyTagType_Single (unsafeCoerce va)
+    0x1 -> KeyValue MyTagType_Multi (unsafeCoerce va)
+    0x2 -> KeyValue MyTagType_Multi2 (unsafeCoerce va)
+    0x3 -> KeyValue MyTagType_Multi3 (unsafeCoerce va)
+    t -> error $ "Data.TagMap.unsafeToKeyValue: no such key type" <> show t
+
+  toKey ki = case ki .&. 0x3 of
+    0x0 -> Some MyTagType_Single
+    0x1 -> Some MyTagType_Multi
+    0x2 -> Some MyTagType_Multi2
+    0x3 -> Some MyTagType_Multi3
+    t -> error $ "Data.TagMap.myKeyType: no such key type" <> show t
+
+  fromTag :: MyTagType a -> Int
+  fromTag t = I# (dataToTag# t)
+
+data Single a
+data Multi
+data Multi2 (k :: * -> *)
+data Multi3
+
+class MyTagTypeOffset x where
+  -- | A type-directed version of `tagOffset` for MyTagType
+  myTagTypeOffset :: proxy x -> Int
+
+instance MyTagTypeOffset (Single a) where
+  myTagTypeOffset _ = 0x0
+
+instance MyTagTypeOffset Multi where
+  myTagTypeOffset _ = 0x1
+
+instance MyTagTypeOffset (Multi2 k) where
+  myTagTypeOffset _ = 0x2
+
+instance MyTagTypeOffset Multi3 where
+  myTagTypeOffset _ = 0x3
+
+instance GEq MyTagType where
+  geq MyTagType_Single MyTagType_Single = Just Refl
+  geq MyTagType_Multi MyTagType_Multi = Just Refl
+  geq MyTagType_Multi2 MyTagType_Multi2 = Just Refl
+  geq MyTagType_Multi3 MyTagType_Multi3 = Just Refl
+  geq _ _ = Nothing
+
+instance GCompare MyTagType where
+  gcompare MyTagType_Single MyTagType_Single = GEQ
+  gcompare MyTagType_Single _ = GLT
+  gcompare _ MyTagType_Single = GGT
+  gcompare MyTagType_Multi MyTagType_Multi = GEQ
+  gcompare MyTagType_Multi _ = GLT
+  gcompare _ MyTagType_Multi = GGT
+  gcompare MyTagType_Multi2 MyTagType_Multi2 = GEQ
+  gcompare MyTagType_Multi2 _ = GLT
+  gcompare _ MyTagType_Multi2 = GGT
+  gcompare MyTagType_Multi3 MyTagType_Multi3 = GEQ