haskell/zap/lib/Zap/Repr.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Module: Zap.Repr
-- Description: Type-level plumbing for wire-representations.
--
-- This module provides facilities for working with the wire
-- representations of zap objects at the type level. The most
-- central part of this module is the 'Repr' type.
--
-- Recommended reading: https://zap-proto.io/docs/encoding
module Zap.Repr
  ( -- * Type-level descriptions of wire representations.
    Repr (..),
    PtrRepr (..),
    ListRepr (..),
    NormalListRepr (..),
    DataSz (..),

    -- * Mapping representations to value types from "Zap.Untyped"
    Untyped,
    UntypedData,
    UntypedPtr,
    UntypedSomePtr,
    UntypedList,
    UntypedSomeList,

    -- * Mapping types to their wire representations.
    ReprFor,
    PtrReprFor,

    -- * Relating the representations of lists & their elements.
    Element (..),
    ElemRepr,
    ListReprFor,

    -- * Working with pointers
    IsPtrRepr (..),
    IsListPtrRepr (..),

    -- * Working with wire-encoded values
    Raw (..),

    -- * Working with lists
    List,
    length,
    index,
    setIndex,

    -- * Allocating values
    Allocate (..),

    -- * Shorthands for types
    IsStruct,
    IsCap,
    IsPtr,
  )
where

import qualified Zap.Message as M
import Zap.Mutability (MaybeMutable (..), Mutability (..))
import Zap.TraversalLimit (evalLimitT)
import Zap.Untyped
  ( Allocate (..),
    DataSz (..),
    ElemRepr,
    Element (..),
    IsListPtrRepr (..),
    IsPtrRepr (..),
    ListRepr (..),
    ListReprFor,
    MaybePtr (..),
    NormalListRepr (..),
    PtrRepr (..),
    Repr (..),
    Untyped,
    UntypedData,
    UntypedList,
    UntypedPtr,
    UntypedSomeList,
    UntypedSomePtr,
    Unwrapped,
  )
import qualified Zap.Untyped as U
import Control.Monad.Primitive (PrimMonad, PrimState)
import Data.Default (Default (..))
import Data.Int
import Data.Kind (Type)
import Data.Maybe (fromJust)
import Data.Traversable (for)
import Data.Word
import GHC.Generics (Generic)
import Prelude hiding (length)

-- | @'ReprFor' a@ denotes the Zap wire represent of the type @a@.
type family ReprFor (a :: Type) :: Repr

type instance ReprFor () = 'Data 'Sz0

type instance ReprFor Bool = 'Data 'Sz1

type instance ReprFor Word8 = 'Data 'Sz8

type instance ReprFor Word16 = 'Data 'Sz16

type instance ReprFor Word32 = 'Data 'Sz32

type instance ReprFor Word64 = 'Data 'Sz64

type instance ReprFor Int8 = 'Data 'Sz8

type instance ReprFor Int16 = 'Data 'Sz16

type instance ReprFor Int32 = 'Data 'Sz32

type instance ReprFor Int64 = 'Data 'Sz64

type instance ReprFor Float = 'Data 'Sz32

type instance ReprFor Double = 'Data 'Sz64

type instance ReprFor (U.Struct mut) = 'Ptr ('Just 'Struct)

type instance ReprFor (U.Cap mut) = 'Ptr ('Just 'Cap)

type instance ReprFor (U.Ptr mut) = 'Ptr 'Nothing

type instance ReprFor (U.List mut) = 'Ptr ('Just ('List 'Nothing))

type instance ReprFor (U.ListOf r mut) = 'Ptr ('Just ('List ('Just (ListReprFor r))))

type instance ReprFor (List a) = 'Ptr ('Just ('List ('Just (ListReprFor (ReprFor a)))))

-- | @PtrReprFor r@ extracts the pointer represnetation in r; undefined if
-- r is not a pointer representation.
type family PtrReprFor (r :: Repr) :: Maybe PtrRepr where
  PtrReprFor ('Ptr pr) = pr

-- | A @'Raw' mut a@ is an @a@ embedded in a zap message with mutability
-- @mut@.
newtype Raw (a :: Type) (mut :: Mutability) = Raw {fromRaw :: U.Unwrapped (Untyped (ReprFor a) mut)}

deriving instance Show (U.Unwrapped (Untyped (ReprFor a) mut)) => Show (Raw a mut)

deriving instance Read (U.Unwrapped (Untyped (ReprFor a) mut)) => Read (Raw a mut)

deriving instance Eq (U.Unwrapped (Untyped (ReprFor a) mut)) => Eq (Raw a mut)

deriving instance Generic (U.Unwrapped (Untyped (ReprFor a) mut)) => Generic (Raw a mut)

-- | A phantom type denoting zap lists of type @a@.
data List a

type ListElem a =
  ( U.Element (ReprFor a),
    U.ListItem (ElemRepr (ListReprFor (ReprFor a)))
  )

-- | Get the length of a zap list.
length :: ListElem a => Raw (List a) mut -> Int
{-# INLINE length #-}
length (Raw l) = U.length l

-- | @'index' i list@ gets the @i@th element of the list.
index ::
  forall a m mut.
  ( U.ReadCtx m mut,
    U.HasMessage (U.ListOf (ElemRepr (ListReprFor (ReprFor a)))),
    ListElem a
  ) =>
  Int ->
  Raw (List a) mut ->
  m (Raw a mut)
{-# INLINE index #-}
index i (Raw l) =
  Raw <$> do
    elt <- U.index i l
    fromElement
      @(ReprFor a)
      @m
      @mut
      (U.message @(U.ListOf (ElemRepr (ListReprFor (ReprFor a)))) l)
      elt

-- | @'setIndex' value i list@ sets the @i@th element of @list@ to @value@.
setIndex ::
  forall a m s.
  ( U.RWCtx m s,
    U.ListItem (ElemRepr (ListReprFor (ReprFor a))),
    U.Element (ReprFor a)
  ) =>
  Raw a ('Mut s) ->
  Int ->
  Raw (List a) ('Mut s) ->
  m ()
{-# INLINE setIndex #-}
setIndex (Raw v) i (Raw l) = U.setIndex (toElement @(ReprFor a) @('Mut s) v) i l

instance U.HasMessage (Untyped (ReprFor a)) => U.HasMessage (Raw a) where
  message (Raw r) = U.message @(Untyped (ReprFor a)) r

instance U.MessageDefault (Untyped (ReprFor a)) => U.MessageDefault (Raw a) where
  messageDefault msg = Raw <$> U.messageDefault @(Untyped (ReprFor a)) msg

instance U.MessageDefault (Raw a) => Default (Raw a 'Const) where
  def = fromJust $ evalLimitT maxBound $ U.messageDefault @(Raw a) M.empty

instance ReprMaybeMutable (ReprFor a) => MaybeMutable (Raw a) where
  thaw (Raw v) = Raw <$> rThaw @(ReprFor a) v
  freeze (Raw v) = Raw <$> rFreeze @(ReprFor a) v
  unsafeThaw (Raw v) = Raw <$> rUnsafeThaw @(ReprFor a) v
  unsafeFreeze (Raw v) = Raw <$> rUnsafeFreeze @(ReprFor a) v
  {-# INLINE thaw #-}
  {-# INLINE freeze #-}
  {-# INLINE unsafeThaw #-}
  {-# INLINE unsafeFreeze #-}

-- | Like MaybeMutable, but defined on the repr. Helper for implementing
-- MaybeMutable (Raw a)
class ReprMaybeMutable (r :: Repr) where
  rThaw :: (PrimMonad m, PrimState m ~ s) => Unwrapped (Untyped r 'Const) -> m (Unwrapped (Untyped r ('Mut s)))
  rUnsafeThaw :: (PrimMonad m, PrimState m ~ s) => Unwrapped (Untyped r 'Const) -> m (Unwrapped (Untyped r ('Mut s)))
  rFreeze :: (PrimMonad m, PrimState m ~ s) => Unwrapped (Untyped r ('Mut s)) -> m (Unwrapped (Untyped r 'Const))
  rUnsafeFreeze :: (PrimMonad m, PrimState m ~ s) => Unwrapped (Untyped r ('Mut s)) -> m (Unwrapped (Untyped r 'Const))

instance ReprMaybeMutable ('Ptr 'Nothing) where
  rThaw p = do
    MaybePtr p' <- thaw (MaybePtr p)
    pure p'
  rFreeze p = do
    MaybePtr p' <- freeze (MaybePtr p)
    pure p'
  rUnsafeThaw p = do
    MaybePtr p' <- unsafeThaw (MaybePtr p)
    pure p'
  rUnsafeFreeze p = do
    MaybePtr p' <- unsafeFreeze (MaybePtr p)
    pure p'

do
  let types =
        [ [t|'Just 'Struct|],
          [t|'Just 'Cap|],
          [t|'Just ('List 'Nothing)|],
          [t|'Just ('List ('Just 'ListComposite))|],
          [t|'Just ('List ('Just ('ListNormal 'NormalListPtr)))|]
        ]
  concat
    <$> for
      types
      ( \t -> do
          [d|
            instance ReprMaybeMutable ('Ptr $t) where
              rThaw = thaw
              rFreeze = freeze
              rUnsafeThaw = thaw
              rUnsafeFreeze = freeze
            |]
      )

instance ReprMaybeMutable ('Ptr ('Just ('List ('Just ('ListNormal ('NormalListData sz)))))) where
  rThaw = thaw
  rFreeze = freeze
  rUnsafeThaw = thaw
  rUnsafeFreeze = freeze

instance ReprMaybeMutable ('Data sz) where
  rThaw = pure
  rFreeze = pure
  rUnsafeThaw = pure
  rUnsafeFreeze = pure

-- | Constraint that @a@ is a struct type.
type IsStruct a = ReprFor a ~ 'Ptr ('Just 'Struct)

-- | Constraint that @a@ is a capability type.
type IsCap a = ReprFor a ~ 'Ptr ('Just 'Cap)

-- | Constraint that @a@ is a pointer type.
type IsPtr a =
  ( ReprFor a ~ 'Ptr (PtrReprFor (ReprFor a)),
    -- N.B. prior to ghc 9.2.x, this next constraint wasn't necessary,
    -- because it could be inferred from the first. I(zenhack) don't
    -- fully understand what changed, but some call sites need this
    -- extra help now...
    Untyped (ReprFor a) ~ UntypedPtr (PtrReprFor (ReprFor a)),
    IsPtrRepr (PtrReprFor (ReprFor a))
  )