content: {-# language QuasiQuotes #-}
  module Render.Spec.Extends
    ( structExtends
    ) where
  
  import           Prettyprinter
  import qualified Data.Vector.Extra             as V
  import           Polysemy
  import           Polysemy.Input
  import           Relude                  hiding ( Enum
                                                  , Text
                                                  )
  import           Text.InterpolatedString.Perl6.Unindented
  
  import           Control.Exception              ( throwIO )
  import           Control.Monad.Trans.Cont       ( ContT
                                                  , evalContT
                                                  )
  import           Foreign.Ptr
  import           Foreign.Storable
  import           GHC.IO.Exception
  import           GHC.TypeLits
  import           Type.Reflection
  
  import           CType
  import           Data.Traversable
  import           Error
  import           Haskell
  import           Render.Element
  import           Render.SpecInfo
  import           Spec.Types
  
  structExtends
    :: forall r t
     . (HasErr r, HasRenderParams r, HasSpecInfo r, KnownSpecFlavor t)
    => Spec t
    -> Sem r RenderElement
  structExtends spec = genRe "Extends type family" $ do
    tellExplicitModule =<< mkModuleName ["CStruct", "Extends"]
    tellNotReexportable
    typeFamily spec
    classes spec
    case specFlavor @t of
      SpecVk -> pure ()
      SpecXr -> inheritanceClasses spec
  
  typeFamily
    :: forall r t
     . (HasErr r, HasRenderParams r, HasSpecInfo r, HasRenderElem r)
    => Spec t
    -> Sem r ()
  typeFamily Spec {..} = do
    RenderParams {..} <- input
    tellExport (EType (TyConName "Extends"))
    tellImport ''Relude.Type
    tellImport ''TypeError
    tellImport ''Constraint
    tellImportWithAll ''ErrorMessage
    cases <-
      fmap (fmap snd . sortOn fst . concat)
      $ forV (toList specStructs)
      $ \child -> do
          let cName = mkTyName (sName child)
          cTyDoc <- if V.null (sExtendedBy child)
            then renderTypeHighPrecSource (ConT (typeName cName))
            else renderTypeHighPrecSource (ConT (typeName cName) :@ PromotedNilT)
          forV (toList (sExtends child)) $ \parent -> do
            let pName = mkTyName parent
            tellSourceImport pName
            pTyDoc <- renderTypeHighPrecSource (ConT (typeName pName))
            pure (parent, "Extends" <+> pTyDoc <+> cTyDoc <+> "= ()")
    tellDoc
      $ "type family Extends (a :: [Type] -> Type) (b :: Type) :: Constraint where"
      <> line
      <> indent
           2
           (vsep
             (cases
             <> [ "Extends a b = TypeError (ShowType a :<>: Text \" is not extended by \" :<>: ShowType b)"
                ]
             )
           )
  
  classes
    :: forall r t
     . (HasErr r, HasRenderParams r, HasSpecInfo r, HasRenderElem r)
    => Spec t
    -> Sem r ()
  classes Spec {..} = do
    RenderParams {..} <- input
    tellExport (EClass (TyConName "PeekChain"))
    tellExport (EClass (TyConName "PokeChain"))
    tellExport (EType (TyConName "Chain"))
    tellExport (EType (TyConName "Extendss"))
    tellExport (EData (TyConName "SomeStruct"))
    tellExport (ETerm (TermName "extendSomeStruct"))
    tellExport (ETerm (TermName "withSomeStruct"))
    tellExport (ETerm (TermName "withSomeCStruct"))
    tellExport (ETerm (TermName "peekSomeCStruct"))
    tellExport (ETerm (TermName "pokeSomeCStruct"))
    tellExport (ETerm (TermName "forgetExtensions"))
    tellExport (EClass (TyConName "Extensible"))
    tellExport (EPat (ConName "::&"))
    tellExport (EPat (ConName ":&"))
    tellImport (TyConName "Extends")
    tellImportWithAll (TyConName "ToCStruct")
    tellImportWithAll (TyConName "FromCStruct")
    tellImportWithAll (TyConName "Zero")
    tellImport ''Relude.Type
    tellImportWith ''Proxy 'Proxy
    tellImport ''Constraint
    tellImport ''Typeable
    tellImport ''Ptr
    tellImport 'nullPtr
    tellImport 'castPtr
    tellImport 'plusPtr
    tellImport 'evalContT
    tellImport 'lift
    tellImport 'join
    tellImportWithAll ''ContT
    tellImportWith ''Storable 'peek
    tellImportWith ''Storable 'poke
    tellImport 'throwIO
    tellImportWithAll ''IOException
    tellImportWithAll ''IOErrorType
    tellImport 'typeRep
    tellImport 'fromMaybe
    estt <- maybe (throw "No extensible structure type tag type")
                  pure
                  extensibleStructTypeType
    peekChainCases <- fmap (V.mapMaybe id) . forV specStructs $ \Struct {..} ->
      if V.null sExtends
        then pure Nothing
        else Just <$> case sMembers V.!? 0 of
          Just (StructMember smName (TypeName estt') vals _ _ _)
            | Just smName == extensibleStructTypeMemberName, estt' == estt -> case
                vals
              of
              -- GHC complains if this match is inline...
                V.Singleton typeEnum -> do
                  -- If this type contains a union then it doesn't have a PeekCStruct
                  -- instance
                  --
                  -- TODO: Actually check here for the instance, and don't repeat this
                  -- union checking logic.
                  let isDiscriminated u =
                        Spec.Types.sName u
                          `elem` (udUnionType <$> toList unionDiscriminators)
                  unions <- filter (not . isDiscriminated) <$> containsUnion sName
  
                  let tagCon = pretty (mkConName estt (CName typeEnum))
                      match  = tagCon <+> "->"
                  case unions of
                    [] -> do
                      let n = mkTyName sName
                      tDoc <- renderTypeHighPrecSource $ if V.null sExtendedBy
                        then ConT (typeName n)
                        else ConT (typeName n) :@ PromotedNilT
                      tellImportWithAll (mkTyName estt)
                      tellSourceImport n
                      pure $ match <+> "go @" <> tDoc
                    u : _ ->
                      pure
                        $   match
                        <+> "throwIO $ IOError Nothing InvalidArgument \"peekChainHead\" (\"struct type"
                        <+> tagCon
                        <+> "contains an undiscriminated union ("
                        <>  pretty (mkTyName (Spec.Types.sName u))
                        <>  ") and can't be safely peeked\") Nothing Nothing"
                _ -> throw "Multiple values for sType"
          _ -> throw $ "Unable to find sType member in " <> show sName
    completeStructTailPragmas <-
      fmap (V.mapMaybe id) . forV specStructs $ \Struct {..} -> if V.null sExtends
        then pure Nothing
        else Just <$> do
          let n = mkTyName sName
          tellSourceImport n
          pure $ "{-# complete (::&) ::" <+> pretty n <+> "#-}"
    tellBoot $ do
      tellExport (EType (TyConName "Extendss"))
      tellExport (EType (TyConName "PeekChain"))
      tellExport (EType (TyConName "PokeChain"))
      tellExport (EType (TyConName "Chain"))
      tellImport ''Relude.Type
      tellImport ''Constraint
      tellDoc [qqi|
        class PeekChain (xs :: [Type])
        class PokeChain (xs :: [Type])
        type family Extends (p :: [Type] -> Type) (x :: Type) :: Constraint where ..
        type family Extendss (p :: [Type] -> Type) (xs :: [Type]) :: Constraint where
          Extendss p '[]      = ()
          Extendss p (x : xs) = (Extends p x, Extendss p xs)
        type family Chain (xs :: [a]) = (r :: a) | r -> xs where
          Chain '[]    = ()
          Chain (x:xs) = (x, Chain xs)
      |]
    estmn <- maybe (throw "No extensible struct member name")
                   (pure . pretty . mkMemberName "XrBaseOutStructure")
                   extensibleStructTypeMemberName
    tellDoc [qqi|
      data SomeStruct (a :: [Type] -> Type) where
        SomeStruct
          :: forall a es
           . (Extendss a es, PokeChain es, Show (Chain es))
          => a es
          -> SomeStruct a
  
      deriving instance (forall es. Show (Chain es) => Show (a es)) => Show (SomeStruct a)
  
      -- | The constraint is so on this instance to encourage type inference
      instance Zero (a '[]) => Zero (SomeStruct a) where
        zero = SomeStruct (zero :: a '[])
  
      -- | Forget which extensions a pointed-to struct has by casting the pointer
      forgetExtensions :: Ptr (a es) -> Ptr (SomeStruct a)
      forgetExtensions = castPtr
  
      -- | Add an extension to the beginning of the struct chain
      --
      -- This can be used to optionally extend structs based on some condition (for
      -- example, an extension or layer being available)
      extendSomeStruct
        :: (Extensible a, Extends a e, ToCStruct e, Show e)
        => e
        -> SomeStruct a
        -> SomeStruct a
      extendSomeStruct e (SomeStruct a) = SomeStruct (setNext a (e, getNext a))
  
      -- | Consume a 'SomeStruct' value
      withSomeStruct
        :: forall a b
         . SomeStruct a
        -> (forall es . (Extendss a es, PokeChain es, Show (Chain es)) => a es -> b)
        -> b
      withSomeStruct (SomeStruct s) f = f s
  
      -- | Write the C representation of some extended @a@ and use the pointer,
      -- the pointer must not be returned from the continuation.
      withSomeCStruct
        :: forall a b
         . (forall es . (Extendss a es, PokeChain es) => ToCStruct (a es))
        => SomeStruct a
        -> (forall es . (Extendss a es, PokeChain es) => Ptr (a es) -> IO b)
        -> IO b
      withSomeCStruct s f = withSomeStruct s (`withCStruct` f)
  
      -- | Given some memory for the head of the chain, allocate and poke the
      -- tail and run an action.
      pokeSomeCStruct
        :: (forall es . (Extendss a es, PokeChain es) => ToCStruct (a es))
        => Ptr (SomeStruct a)
        -- ^ Pointer to some memory at least the size of the head of the struct
        -- chain.
        -> SomeStruct a
        -- ^ The struct to poke
        -> IO b
        -- ^ Computation to run while the poked tail is valid
        -> IO b
      pokeSomeCStruct p (SomeStruct s) = pokeCStruct (castPtr p) s
  
      -- | Given a pointer to a struct with an unknown chain, peek the struct and
      -- its chain.
      peekSomeCStruct
        :: forall a
         . (Extensible a, forall es . (Extendss a es, PeekChain es) => FromCStruct (a es))
        => Ptr (SomeStruct a)
        -> IO (SomeStruct a)
      peekSomeCStruct p = do
        head'  <- peekCStruct (castPtr @_ @(a '[]) p)
        pNext <- peek @(Ptr BaseOutStructure) (p `plusPtr` 8)
        peekSomeChain @a pNext $ \\tail' -> SomeStruct (setNext head' tail')
  
      peekSomeChain
        :: forall a b
         . (Extensible a)
        => Ptr BaseOutStructure
        -> (  forall es
            . (Extendss a es, PokeChain es, Show (Chain es))
           => Chain es
           -> b
           )
        -> IO b
      peekSomeChain p c = if p == nullPtr
        then pure (c ())
        else do
          baseOut <- peek p
          join
            $ peekChainHead @a (case baseOut of BaseOutStructure\{{estmn}} -> {estmn})
                               (castPtr @BaseOutStructure @() p)
            $ \\head' -> peekSomeChain @a (case baseOut of BaseOutStructure\{next} -> next)
                                        (\\tail' -> c (head', tail'))
  
      peekChainHead
        :: forall a b
         . Extensible a
        => StructureType
        -> Ptr ()
        -> (forall e . (Extends a e, ToCStruct e, Show e) => e -> b)
        -> IO b
      peekChainHead ty p c = case ty of
      {indent 2 (vsep (toList peekChainCases))}
        t -> throwIO $ IOError Nothing InvalidArgument "peekChainHead" ("Unrecognized struct type: " <> show t) Nothing Nothing
       where
        go :: forall e . (Typeable e, FromCStruct e, ToCStruct e, Show e) => IO b
        go =
          let r = extends @a @e Proxy $ do
                head' <- peekCStruct @e (castPtr p)
                pure $ c head'
          in  fromMaybe
                (throwIO $ IOError
                  Nothing
                  InvalidArgument
                  "peekChainHead"
                  (  "Illegal struct extension of "
                  <> extensibleTypeName @a
                  <> " with "
                  <> show ty
                  )
                  Nothing
                  Nothing
                )
                r
  
      class Extensible (a :: [Type] -> Type) where
        extensibleTypeName :: String
        -- ^ For error reporting an invalid extension
        getNext :: a es -> Chain es
        setNext :: a ds -> Chain es -> a es
        extends :: forall e b proxy. Typeable e => proxy e -> (Extends a e => b) -> Maybe b
  
      type family Chain (xs :: [a]) = (r :: a) | r -> xs where
        Chain '[]    = ()
        Chain (x:xs) = (x, Chain xs)
  
      -- | A pattern synonym to separate the head of a struct chain from the
      -- tail, use in conjunction with ':&' to extract several members.
      --
      -- @
      -- Head\{..} ::& () <- returningNoTail a b c
      -- -- Equivalent to
      -- Head\{..} <- returningNoTail @'[] a b c
      -- @
      --
      -- @
      -- Head\{..} ::& Foo\{..} :& Bar\{..} :& () <- returningWithTail a b c
      -- @
      --
      -- @
      -- myFun (Head\{..} :&& Foo\{..} :& ())
      -- @
      pattern (::&) :: Extensible a => a es' -> Chain es -> a es
      pattern a ::& es <- (\\a -> (a, getNext a) -> (a, es))
        where a ::& es = setNext a es
      infix 6 ::&
      {vsep (toList completeStructTailPragmas)}
  
      -- | View the head and tail of a 'Chain', see '::&'
      --
      -- Equivalent to @(,)@
      pattern (:&) :: e -> Chain es -> Chain (e:es)
      pattern e :& es = (e, es)
      infixr 7 :&
      \{-# complete (:&) #-}
  
      type family Extendss (p :: [Type] -> Type) (xs :: [Type]) :: Constraint where
        Extendss p '[]      = ()
        Extendss p (x : xs) = (Extends p x, Extendss p xs)
  
      class PokeChain es where
        withChain :: Chain es -> (Ptr (Chain es) -> IO a) -> IO a
        withZeroChain :: (Ptr (Chain es) -> IO a) -> IO a
  
      instance PokeChain '[] where
        withChain () f = f nullPtr
        withZeroChain f = f nullPtr
  
      instance (ToCStruct e, PokeChain es) => PokeChain (e:es) where
        withChain (e, es) f = evalContT $ do
          t <- ContT $ withChain es
          h <- ContT $ withCStruct e
          lift $ linkChain h t
          lift $ f (castPtr h)
        withZeroChain f = evalContT $ do
          t <- ContT $ withZeroChain @es
          h <- ContT $ withZeroCStruct @e
          lift $ linkChain h t
          lift $ f (castPtr h)
  
      class PeekChain es where
        peekChain :: Ptr (Chain es) -> IO (Chain es)
  
      instance PeekChain '[] where
        peekChain _ = pure ()
  
      instance (FromCStruct e, PeekChain es) => PeekChain (e:es) where
        peekChain p = do
          h <- peekCStruct @e (castPtr p)
          tPtr <- peek (p `plusPtr` 8)
          t <- peekChain tPtr
          pure (h, t)
  
      linkChain :: Ptr a -> Ptr b -> IO ()
      linkChain head' tail' = poke (head' `plusPtr` 8) tail'
    |]
  
  inheritanceClasses
    :: forall r t
     . (HasErr r, HasRenderParams r, HasSpecInfo r, HasRenderElem r)
    => Spec t
    -> Sem r ()
  inheritanceClasses Spec {..} = do
    RenderParams {..} <- input
    tellExport (EData (TyConName "SomeChild"))
    tellExport (ETerm (TermName "withSomeChild"))
    tellExport (ETerm (TermName "lowerChildPointer"))
    tellExport (EType (TyConName "Inherits"))
    tellExport (EClass (TyConName "Inheritable"))
    tellImport (TyConName "ToCStruct")
    tellImport ''Relude.Type
    tellImport ''Ptr
    tellImport 'castPtr
    tellImport ''TypeError
    tellImportWithAll ''ErrorMessage
    let parentsAndChildren =
          toList . flip V.concatMap specStructs $ \Struct {..} ->
            (sName, ) <$> sInheritedBy
    inheritsCases <- for parentsAndChildren $ \(p, c) -> do
      let pTyName = mkTyName p
          pTy     = case p of
            "XrCompositionLayerBaseHeader" -> parens (pretty pTyName <+> "'[]")
            _                              -> pretty pTyName
          cTy = mkTyName c
      tellSourceImport pTyName
      tellSourceImport cTy
      pure [qqi|Inherits {pTy} {cTy} = ()|]
    tellDoc [qqi|
      data SomeChild (a :: Type) where
        SomeChild :: forall a b . (Inherits a b, Typeable b, ToCStruct b, Show b) => b -> SomeChild a
      deriving instance Show (SomeChild a)
  
      type family Inherits (a :: Type) (b :: Type) :: Constraint where
      {indent 2 (vsep inheritsCases)}
        Inherits parent child =
          TypeError (ShowType parent :<>: Text " is not inherited by " :<>: ShowType child)
  
      class Inheritable (a :: Type) where
        peekSomeCChild :: Ptr (SomeChild a) -> IO (SomeChild a)
  
      withSomeChild :: SomeChild a -> (Ptr (SomeChild a) -> IO b) -> IO b
      withSomeChild (SomeChild c) f = withCStruct c (f . lowerChildPointer)
  
      lowerChildPointer :: Inherits a b => Ptr b -> Ptr (SomeChild a)
      lowerChildPointer = castPtr
    |]
  