Peekers.hs

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-|
Module      : HsLua.Marshalling.Peekers
Copyright   : © 2020-2024 Albert Krewinkel
License     : MIT
Maintainer  : Albert Krewinkel <tarleb@hslua.org>
Stability   : beta
Portability : Portable

Functions which unmarshal and retrieve Haskell values from Lua's stack.
-}
module HsLua.Marshalling.Peekers
  ( -- * Peeking values from the stack
    -- ** Primitives
    peekNil
  , peekNoneOrNil
  , peekBool
  , peekIntegral
  , peekRealFloat
    -- ** Strings
  , peekByteString
  , peekLazyByteString
  , peekString
  , peekText
  , peekStringy
  , peekName
  -- ** Readable types
  , peekRead
  -- ** Collections
  , peekKeyValuePairs
  , peekList
  , peekNonEmpty
  , peekMap
  , peekSet
  -- ** Combinators
  , choice
  , peekFieldRaw
  , peekIndexRaw
  , peekNilOr
  , peekNoneOr
  , peekNoneOrNilOr
  , peekPair
  , peekTriple
  -- ** Building peek functions
  , typeChecked
  , reportValueOnFailure
  , typeMismatchMessage
  ) where

import Control.Applicative (Alternative (..))
import Control.Monad ((<$!>), (>=>), void)
import Data.ByteString (ByteString)
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Map (Map)
import Data.Set (Set)
import Data.String (IsString (fromString))
import HsLua.Core as Lua
import HsLua.Marshalling.Peek
import Text.Read (readMaybe)

import qualified Data.ByteString.Lazy as BL
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified HsLua.Core.Unsafe as Unsafe
import qualified HsLua.Core.Utf8 as Utf8

-- | Use @test@ to check whether the value at stack index @n@ has
-- the correct type and use @peekfn@ to convert it to a Haskell
-- value if possible. A successfully received value is wrapped
-- using the 'Right' constructor, while a type mismatch results
-- in @Left PeekError@ with the given error message.
typeChecked :: Name                         -- ^ expected type
            -> (StackIndex -> LuaE e Bool)  -- ^ pre-condition checker
            -> Peeker e a
            -> Peeker e a
typeChecked expectedType test peekfn idx = do
  v <- liftLua $ test idx
  if v
    then peekfn idx
    else typeMismatchMessage expectedType idx >>= failPeek

-- | Generate a type mismatch error.
typeMismatchMessage :: Name       -- ^ expected type
                    -> StackIndex -- ^ index of offending value
                    -> Peek e ByteString
typeMismatchMessage (Name expected) idx = liftLua $ do
  pushTypeMismatchError expected idx
  (tostring top <* pop 1) >>= \case
    Just !msg -> return msg
    Nothing  -> return $ mconcat
      [ "Unknown type mismatch for "
      , expected
      , " at stack index "
      , Utf8.fromString $ show (fromStackIndex idx)
      ]

-- | Report the expected and actual type of the value under the given
-- index if conversion failed.
reportValueOnFailure :: Name         -- ^ expected type
                     -> (StackIndex -> LuaE e (Maybe a))
                     -> Peeker e a
reportValueOnFailure expected peekMb idx = do
  res <- liftLua $ peekMb idx
  case res of
    Just x  -> return $! x
    Nothing -> typeMismatchMessage expected idx >>= failPeek

--
-- Primitives
--

-- | Succeeds if the value at the given index is @nil@.
peekNil :: Peeker e ()
peekNil = typeChecked "nil" Lua.isnil $ const (return ())
{-# INLINABLE peekNil #-}

-- | Succeeds if the given index is not valid or if the value at this
-- index is @nil@.
peekNoneOrNil :: Peeker e ()
peekNoneOrNil = typeChecked "none or nil" Lua.isnoneornil $ const (return ())
{-# INLINABLE peekNoneOrNil #-}

-- | Retrieves a 'Bool' as a Lua boolean.
peekBool :: Peeker e Bool
peekBool = liftLua . toboolean

--
-- Strings
--

-- | Like 'tostring', but ensures that the value at the given index is
-- not silently converted to a string, as would happen with numbers.
-- Also returns 'Nothing' if the value is a number and there is no stack
-- slot left on the Lua stack, which would be needed to convert the
-- number to a string without changing the original slot.
toByteString :: StackIndex -> LuaE e (Maybe ByteString)
toByteString idx = do
  -- Do an explicit type check, as @tostring@ converts numbers strings
  -- /in-place/, which we need to avoid.
  ltype idx >>= \case
    TypeString -> tostring idx
    _          -> checkstack 1 >>= \case
      False -> pure Nothing
      True  ->  do
        pushvalue idx
        tostring top <* pop 1
{-# INLINABLE toByteString #-}

-- | Retrieves a 'ByteString' as a raw string.
peekByteString :: Peeker e ByteString
peekByteString = reportValueOnFailure "string" toByteString
{-# INLINABLE peekByteString #-}

-- | Retrieves a lazy 'BL.ByteString' as a raw string.
peekLazyByteString :: Peeker e BL.ByteString
peekLazyByteString = (BL.fromStrict <$!>) . peekByteString
{-# INLINABLE peekLazyByteString #-}

-- | Retrieves a 'String' from an UTF-8 encoded Lua string.
peekString :: Peeker e String
peekString = peekStringy
{-# INLINABLE peekString #-}

-- | Retrieves a String-like value from an UTF-8 encoded Lua string.
--
-- This should not be used to peek 'ByteString' values or other values
-- for which construction via 'fromString' can result in loss of
-- information.
peekStringy :: forall a e. IsString a => Peeker e a
peekStringy = fmap (fromString . Utf8.toString) . peekByteString
{-# INLINABLE peekStringy #-}

-- | Retrieves a 'T.Text' value as an UTF-8 encoded string.
peekText :: Peeker e T.Text
peekText = (Utf8.toText <$!>) . peekByteString
{-# INLINABLE peekText #-}

-- | Retrieves a Lua string as 'Name'.
peekName :: Peeker e Name
peekName = (Name <$!>) . peekByteString
{-# INLINABLE peekName #-}

--
-- Arbitrary values
--

-- | Retrieves a value by getting a String from Lua, then using
-- 'readMaybe' to convert the String into a Haskell value.
peekRead :: forall a e. Read a => Peeker e a
peekRead = peekString >=> readValue
  where
    readValue s = case readMaybe s of
      Just x  -> pure x
      Nothing -> failPeek $ "Could not read: " <> Utf8.fromString s

--
-- Numbers
--

-- | Retrieves an 'Integral' value from the Lua stack.
peekIntegral :: forall a e. (Integral a, Read a) => Peeker e a
peekIntegral idx = liftLua (ltype idx) >>= \case
  TypeNumber  -> fromIntegral <$!>
                 reportValueOnFailure "Integral" tointeger idx
  TypeString  -> do
    Just str <- liftLua $ tostring idx
    case readMaybe (Utf8.toString str) of
      Nothing -> typeMismatchMessage "Integral" idx >>= failPeek
      Just x  -> return x
  _ -> typeMismatchMessage "Integral" idx >>= failPeek

-- | Retrieve a 'RealFloat' (e.g., 'Float' or 'Double') from the stack.
peekRealFloat :: forall a e. (RealFloat a, Read a) => Peeker e a
peekRealFloat idx = liftLua (ltype idx) >>= \case
  TypeString  -> do
    Just str <- liftLua $ tostring idx
    case readMaybe (Utf8.toString str) of
      Nothing -> typeMismatchMessage "RealFloat" idx >>= failPeek
      Just x  -> return x
  _ -> realToFrac <$!> reportValueOnFailure "RealFloat" tonumber idx

-- | Reads a numerically indexed table @t@ into a list, where the 'length' of
-- the list is equal to @rawlen(t)@. The operation will fail unless all
-- numerical fields between @1@ and @rawlen(t)@ can be retrieved.
peekList :: forall a e. LuaError e => Peeker e a -> Peeker e [a]
peekList peekElement = retrieving "list" . peekList' peekElement

-- | Like 'peekList', but fails if the list is empty.
peekNonEmpty :: LuaError e => Peeker e a -> Peeker e (NonEmpty a)
peekNonEmpty peekElement = retrieving "NonEmpty" .
  (peekList' peekElement >=> \case
    (x:xs) -> return (x :| xs)
    []     -> failPeek "empty list")

-- | Helper function that retrieves a list, but doesn't set a context.
peekList' :: LuaError e => Peeker e a -> Peeker e [a]
peekList' peekElement = typeChecked "table" istable $ \idx -> do
  liftLua $ checkstack' 1 "retrieving a list"
  let elementsAt [] = return []
      elementsAt (i : is) = do
        x  <- retrieving ("index " <> showInt i) $
              liftLua (rawgeti idx i) *> peekElement top `lastly` pop 1
        xs <- elementsAt is
        return (x:xs)
      showInt (Lua.Integer x) = fromString $ show x
  listLength <- liftLua (rawlen idx)
  elementsAt [1..fromIntegral listLength]

-- | Retrieves a key-value Lua table as 'Map'.
peekMap :: (LuaError e, Ord a)
        => Peeker e a -> Peeker e b -> Peeker e (Map a b)
peekMap keyPeeker valuePeeker = retrieving "Map"
  . fmap Map.fromList
  . peekKeyValuePairs keyPeeker valuePeeker

-- | Read a table into a list of pairs.
peekKeyValuePairs :: LuaError e
                  => Peeker e a -> Peeker e b -> Peeker e [(a, b)]
peekKeyValuePairs keyPeeker valuePeeker =
  typeChecked "table" istable $ \idx -> cleanup $ do
    liftLua $ checkstack' 2 "retrieving key-value pairs"
    idx' <- liftLua $ absindex idx
    let remainingPairs = nextPair keyPeeker valuePeeker idx' >>= \case
          Nothing -> return []
          Just a  -> (a:) <$!> remainingPairs
    liftLua pushnil
    remainingPairs

-- | Get the next key-value pair from a table. Assumes the last
-- key to be on the top of the stack and the table at the given
-- index @idx@. The next key, if it exists, is left at the top of
-- the stack.
--
-- The key must be either nil or must exist in the table, or this
-- function will crash with an unrecoverable error.
nextPair :: Peeker e a -> Peeker e b -> Peeker e (Maybe (a, b))
nextPair keyPeeker valuePeeker idx = retrieving "key-value pair" $ do
  hasNext <- liftLua $ Unsafe.next idx
  if not hasNext
    then return Nothing
    else do
      key   <- retrieving "key"   $ keyPeeker   (nth 2)
      value <- retrieving "value" $ valuePeeker (nth 1)
      return (Just (key, value))
        `lastly` pop 1  -- remove value, leave the key

-- | Retrieves a 'Set' from an idiomatic Lua representation. A
-- set in Lua is idiomatically represented as a table with the
-- elements as keys. Elements with falsy values are omitted.
peekSet :: (LuaError e, Ord a) => Peeker e a -> Peeker e (Set a)
peekSet elementPeeker = retrieving "Set"
  . fmap (Set.fromList . map fst . filter snd)
  . peekKeyValuePairs elementPeeker peekBool

--
-- Combinators
--

-- | Get value at key from a table.
peekFieldRaw :: LuaError e => Peeker e a -> Name -> Peeker e a
peekFieldRaw peeker name idx =
  retrieving ("raw field '" <> name <> "'") $! do
    liftLua $ do
      checkstack' 1 "peekFieldRaw"
      absidx <- Lua.absindex idx
      pushstring $ fromName name
      void (rawget absidx)
    peeker top `lastly` Lua.pop 1
{-# INLINABLE peekFieldRaw #-}

-- | Get value at integer index key from a table.
peekIndexRaw :: LuaError e => Lua.Integer -> Peeker e a -> Peeker e a
peekIndexRaw i peeker idx = do
  let showInt (Lua.Integer x) = fromString $ show x
  retrieving (fromString $ "raw index '" <> showInt i <> "'") $! do
    liftLua . void $ rawgeti idx i
    peeker top `lastly` Lua.pop 1
{-# INLINABLE peekIndexRaw #-}


-- | Returns 'empty' if the value at the given index is @nil@;
-- otherwise returns the result of peeker @p@.
peekNilOr :: Alternative m
          => Peeker e a          -- ^ p
          -> Peeker e (m a)
peekNilOr p idx = liftLua (ltype idx) >>= \case
  TypeNil  -> pure empty
  _        -> pure <$> p idx

-- | Returns 'empty' if the value at the given index is @none@;
-- otherwise returns the result of peeker @p@.
peekNoneOr :: Alternative m
           => Peeker e a          -- ^ p
           -> Peeker e (m a)
peekNoneOr p idx = liftLua (ltype idx) >>= \case
  TypeNone -> pure empty
  _        -> pure <$> p idx

-- | Returns 'empty' if the value at the given index is @none@ or
-- @nil@; otherwise returns the result of peeker @p@.
peekNoneOrNilOr :: Alternative m
                => Peeker e a          -- ^ p
                -> Peeker e (m a)
peekNoneOrNilOr p idx = liftLua (ltype idx) >>= \case
  TypeNil  -> pure empty
  TypeNone -> pure empty
  _        -> pure <$> p idx

-- | Retrieves a value pair from a table. Expects the values to be
-- stored in a numerically indexed table; does not access metamethods.
peekPair :: LuaError e
         => Peeker e a -> Peeker e b
         -> Peeker e (a, b)
peekPair peekA peekB idx = cleanup $ do
  liftLua $ checkstack' 2 "retrieving a pair"
  idx' <- liftLua $ absindex idx
  a <- liftLua (rawgeti idx' 1) *> peekA top
  b <- liftLua (rawgeti idx' 2) *> peekB top
  return (a, b)

-- | Retrieves a value triple from a table. Expects the values to be
-- stored in a numerically indexed table, with no metamethods.
peekTriple :: LuaError e
           => Peeker e a -> Peeker e b -> Peeker e c
           -> Peeker e (a, b, c)
peekTriple peekA peekB peekC idx = cleanup $ do
  liftLua $ checkstack' 3 "retrieving a triple"
  idx' <- liftLua $ absindex idx
  a <- liftLua (rawgeti idx' 1) *> peekA top
  b <- liftLua (rawgeti idx' 2) *> peekB top
  c <- liftLua (rawgeti idx' 3) *> peekC top
  return (a,b,c)

-- | Try all peekers and return the result of the first to succeed.
choice :: LuaError e
       => [Peeker e a]
       -> Peeker e a
choice peekers idx = case peekers of
  [] -> failPeek "all choices failed"
  p:ps -> p idx <|> choice ps idx
{-# INLINABLE choice #-}