{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE UndecidableInstances  #-}

module Debug.RecoverRTTI.Classifier (
    Classifier
  , PrimClassifier(..)
  , IsUserDefined(..)
    -- * Generalizations
  , Classifier_(..)
    -- * Nested classification
  , Elem(..)
  , Elems(..)
    -- * Mapping
  , mapClassifier
  ) where

import Data.Aeson (Value)
import Data.HashMap.Lazy (HashMap)
import Data.HashSet (HashSet)
import Data.Int
import Data.IntMap (IntMap)
import Data.IntSet (IntSet)
import Data.Kind
import Data.Map (Map)
import Data.Ratio
import Data.Sequence (Seq)
import Data.Set (Set)
import Data.SOP
import Data.SOP.Dict
import Data.Tree (Tree)
import Data.Void
import Data.Word

import qualified Data.ByteString             as BS.Strict
import qualified Data.ByteString.Lazy        as BS.Lazy
import qualified Data.ByteString.Short       as BS.Short
import qualified Data.HashMap.Internal.Array as HashMap (Array)
import qualified Data.Primitive.Array        as Prim (Array)
import qualified Data.Text                   as Text.Strict
import qualified Data.Text.Lazy              as Text.Lazy
import qualified Data.Vector                 as Vector.Boxed

import Debug.RecoverRTTI.Nat
import Debug.RecoverRTTI.Tuple
import Debug.RecoverRTTI.Wrappers

{-------------------------------------------------------------------------------
  Classifier
-------------------------------------------------------------------------------}

-- | Classifier
--
-- Given a value of some unknown type @a@, a @Classifier a@ will tell you what
-- the type of @a@ is. This is similar to a @TypeRep@, but since we recover
-- this information from the heap, we have less accurate type information than
-- @TypeRep@ does.
type Classifier = Classifier_ IsUserDefined

-- | User-defined types
--
-- If we classify a type as user-defined, we pair the classifier with the
-- original value. This means that a @Classifier@ is sufficient information
-- for staged inference by client code that may wish to further classify these
-- types given additional domain knowledge (see also 'reclassify_').
data IsUserDefined a where
  IsUserDefined :: UserDefined -> IsUserDefined UserDefined

instance Show (IsUserDefined a) where
  show :: IsUserDefined a -> [Char]
show (IsUserDefined UserDefined
_) = [Char]
"IsUserDefined"

{-------------------------------------------------------------------------------
  Generalizations
-------------------------------------------------------------------------------}

-- | Generalization of 'Classifier'
--
-- Type arguments:
--
-- * @o@: Classification of " other " types (not explicitly known to the lib)
--
--   Normally we instantiate this to 'IsUserDefined', classifying all unknown
--   types as 'UserDefined'.
--
-- * @a@: The type we're actually classifying
data Classifier_ (o :: Type -> Type) (a :: Type) :: Type where
  -- Primitive and user-defined types
  C_Prim  :: PrimClassifier a -> Classifier_ o a
  C_Other :: o              a -> Classifier_ o a

  -- Compound
  --
  -- NOTE: C_HashSet requires an argument; 'HashSet' and 'HashMap' cannot be
  -- distinguished from just looking at the heap ('HashSet' is a newtype
  -- around 'HashMap'), and so we classify a 'HashMap' with value type @()@
  -- as a 'HashSet'; however, we can only do this of course if we have at
  -- least one element.

  C_Maybe        :: Elems o '[a]    -> Classifier_ o (Maybe a)
  C_Either       :: Elems o '[a, b] -> Classifier_ o (Either a b)
  C_List         :: Elems o '[a]    -> Classifier_ o [a]
  C_Ratio        :: Elems o '[a]    -> Classifier_ o (Ratio a)
  C_Set          :: Elems o '[a]    -> Classifier_ o (Set a)
  C_Map          :: Elems o '[a, b] -> Classifier_ o (Map a b)
  C_IntMap       :: Elems o '[a]    -> Classifier_ o (IntMap a)
  C_Sequence     :: Elems o '[a]    -> Classifier_ o (Seq a)
  C_Tree         :: Elems o '[a]    -> Classifier_ o (Tree a)
  C_HashSet      :: Elems o '[a]    -> Classifier_ o (HashSet a)
  C_HashMap      :: Elems o '[a, b] -> Classifier_ o (HashMap a b)
  C_HM_Array     :: Elems o '[a]    -> Classifier_ o (HashMap.Array a)
  C_Prim_Array   :: Elems o '[a]    -> Classifier_ o (Prim.Array a)
  C_Vector_Boxed :: Elems o '[a]    -> Classifier_ o (Vector.Boxed.Vector a)

  C_Tuple ::
       (SListI xs, IsValidSize (Length xs))
    => Elems o xs -> Classifier_ o (WrappedTuple xs)

-- | Classifier for primitive types
data PrimClassifier (a :: Type) where
  -- Primitive types

  C_Bool     :: PrimClassifier Bool
  C_Char     :: PrimClassifier Char
  C_Double   :: PrimClassifier Double
  C_Float    :: PrimClassifier Float
  C_Int      :: PrimClassifier Int
  C_Int16    :: PrimClassifier Int16
  C_Int8     :: PrimClassifier Int8
  C_Int32    :: PrimClassifier Int32
  C_Int64    :: PrimClassifier Int64
  C_Integer  :: PrimClassifier Integer
  C_Ordering :: PrimClassifier Ordering
  C_Unit     :: PrimClassifier ()
  C_Word     :: PrimClassifier Word
  C_Word8    :: PrimClassifier Word8
  C_Word16   :: PrimClassifier Word16
  C_Word32   :: PrimClassifier Word32
  C_Word64   :: PrimClassifier Word64

  -- String types
  --
  -- We list @String@ separately, so that we show them properly (rather than
  -- as a list of characters). Of course, empty strings will be inferred as
  -- empty lists instead.

  C_String      :: PrimClassifier String
  C_BS_Strict   :: PrimClassifier BS.Strict.ByteString
  C_BS_Lazy     :: PrimClassifier BS.Lazy.ByteString
  C_BS_Short    :: PrimClassifier BS.Short.ShortByteString
  C_Text_Strict :: PrimClassifier Text.Strict.Text
  C_Text_Lazy   :: PrimClassifier Text.Lazy.Text

  -- Aeson

  C_Value :: PrimClassifier Value

  -- Reference cells

  C_STRef :: PrimClassifier SomeSTRef
  C_TVar  :: PrimClassifier SomeTVar
  C_MVar  :: PrimClassifier SomeMVar

  -- Functions

  C_Fun :: PrimClassifier SomeFun

  -- Containers with no type arguments
  --
  -- We include mutable containers here, because we currently do not attempt
  -- to peek inside them and hence cannot infer any types for their elements.

  C_IntSet            :: PrimClassifier IntSet
  C_Prim_ArrayM       :: PrimClassifier SomePrimArrayM
  C_Vector_Storable   :: PrimClassifier SomeStorableVector
  C_Vector_StorableM  :: PrimClassifier SomeStorableVectorM
  C_Vector_Primitive  :: PrimClassifier SomePrimitiveVector
  C_Vector_PrimitiveM :: PrimClassifier SomePrimitiveVectorM

{-------------------------------------------------------------------------------
  Nested classification
-------------------------------------------------------------------------------}

data Elem o a where
  Elem   :: Classifier_ o a -> Elem o a
  NoElem :: Elem o Void

newtype Elems o xs = Elems (NP (Elem o) xs)

{-------------------------------------------------------------------------------
  Show
-------------------------------------------------------------------------------}

deriving instance Show (PrimClassifier a)

deriving instance (forall x. Show (o x)) => Show (Classifier_ o a)
deriving instance (forall x. Show (o x)) => Show (Elem o a)

instance (forall a. Show (o a), SListI xs) => Show (Elems o xs) where
  showsPrec :: Int -> Elems o xs -> ShowS
showsPrec Int
p (Elems NP (Elem o) xs
xs) =
      case forall {k} (c :: k -> Constraint) (xs :: [k]).
NP (Dict c) xs -> Dict (All c) xs
all_NP NP (Dict (Compose Show (Elem o))) xs
allShow of
        Dict (All (Compose Show (Elem o))) xs
Dict -> forall a. Show a => Int -> a -> ShowS
showsPrec Int
p NP (Elem o) xs
xs
    where
      allShow :: NP (Dict (Compose Show (Elem o))) xs
      allShow :: NP (Dict (Compose Show (Elem o))) xs
allShow = forall k l (h :: (k -> *) -> l -> *) (xs :: l) (f :: k -> *).
(HPure h, SListIN h xs) =>
(forall (a :: k). f a) -> h f xs
hpure forall {k} (c :: k -> Constraint) (a :: k). c a => Dict c a
Dict

{-------------------------------------------------------------------------------
  Map over classifiers
-------------------------------------------------------------------------------}

mapClassifier :: forall m o o'.
     Applicative m
  => (forall a. o a -> m (o' a))
  -> (forall a. Classifier_ o a -> m (Classifier_ o' a))
mapClassifier :: forall (m :: * -> *) (o :: * -> *) (o' :: * -> *).
Applicative m =>
(forall a. o a -> m (o' a))
-> forall a. Classifier_ o a -> m (Classifier_ o' a)
mapClassifier forall a. o a -> m (o' a)
other = forall a. Classifier_ o a -> m (Classifier_ o' a)
go
  where
    go :: forall a. Classifier_ o a -> m (Classifier_ o' a)
    -- Primitive and user-defined types

    go :: forall a. Classifier_ o a -> m (Classifier_ o' a)
go (C_Prim  PrimClassifier a
c) = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a (o :: * -> *). PrimClassifier a -> Classifier_ o a
C_Prim PrimClassifier a
c)
    go (C_Other o a
c) = forall (o :: * -> *) a. o a -> Classifier_ o a
C_Other forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. o a -> m (o' a)
other o a
c

    -- Compound

    go (C_Maybe        Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Maybe xs)
C_Maybe        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Either       Elems o '[a, b]
c) = forall (o :: * -> *) xs b.
Elems o '[xs, b] -> Classifier_ o (Either xs b)
C_Either       forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a, b]
c
    go (C_List         Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o [xs]
C_List         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Ratio        Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Ratio xs)
C_Ratio        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Set          Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Set xs)
C_Set          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Map          Elems o '[a, b]
c) = forall (o :: * -> *) xs b.
Elems o '[xs, b] -> Classifier_ o (Map xs b)
C_Map          forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a, b]
c
    go (C_IntMap       Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (IntMap xs)
C_IntMap       forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Sequence     Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Seq xs)
C_Sequence     forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Tree         Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Tree xs)
C_Tree         forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_HashSet      Elems o '[a]
c) = forall (o :: * -> *) xs.
Elems o '[xs] -> Classifier_ o (HashSet xs)
C_HashSet      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_HashMap      Elems o '[a, b]
c) = forall (o :: * -> *) xs b.
Elems o '[xs, b] -> Classifier_ o (HashMap xs b)
C_HashMap      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a, b]
c
    go (C_HM_Array     Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Array xs)
C_HM_Array     forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Prim_Array   Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Array xs)
C_Prim_Array   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Vector_Boxed Elems o '[a]
c) = forall (o :: * -> *) xs. Elems o '[xs] -> Classifier_ o (Vector xs)
C_Vector_Boxed forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o '[a]
c
    go (C_Tuple        Elems o xs
c) = forall (xs :: [*]) (o :: * -> *).
(SListI xs, IsValidSize (Length xs)) =>
Elems o xs -> Classifier_ o (WrappedTuple xs)
C_Tuple        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems Elems o xs
c

    goElems :: SListI xs => Elems o xs -> m (Elems o' xs)
    goElems :: forall (xs :: [*]). SListI xs => Elems o xs -> m (Elems o' xs)
goElems (Elems NP (Elem o) xs
cs) = forall (o :: * -> *) (xs :: [*]). NP (Elem o) xs -> Elems o xs
Elems forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall k l (h :: (k -> *) -> l -> *) (xs :: l) (g :: * -> *)
       (f :: k -> *) (f' :: k -> *).
(HSequence h, SListIN h xs, Applicative g) =>
(forall (a :: k). f a -> g (f' a)) -> h f xs -> g (h f' xs)
htraverse' forall a. Elem o a -> m (Elem o' a)
goElem NP (Elem o) xs
cs

    goElem :: Elem o a -> m (Elem o' a)
    goElem :: forall a. Elem o a -> m (Elem o' a)
goElem (Elem Classifier_ o a
c) = forall (o :: * -> *) a. Classifier_ o a -> Elem o a
Elem forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. Classifier_ o a -> m (Classifier_ o' a)
go Classifier_ o a
c
    goElem Elem o a
NoElem   = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall (o :: * -> *). Elem o Void
NoElem