-----------------------------------------------------------------------------
-- |
-- Module    : Data.SBV.SEnum
-- Copyright : (c) Levent Erkok
-- License   : BSD3
-- Maintainer: erkokl@gmail.com
-- Stability : experimental
--
-- Add support for symbolic enumerations via a quasi-quoter. The code in this
-- file was initially generated by ChatGPT, which didn't quite work but was
-- close enough to let me finish it off.
--
-- Provides a quasiquoter `[sEnum| ... |]` for enumerations, like:
--
-- > [sEnum| a .. |]       ==> enumFrom a
-- > [sEnum| a, b .. |]    ==> enumFromThen a b
-- > [sEnum| a .. c |]     ==> enumFromTo a c
-- > [sEnum| a, b .. c |]  ==> enumFromThenTo a b c
--
-- All of @a@, @b@, @c@ can be arbitrary expressions.
--
-- If you pass invalid Haskell expressions or incorrect format, a detailed
-- error is raised with source location.
-----------------------------------------------------------------------------

{-# LANGUAGE TemplateHaskellQuotes #-}

{-# OPTIONS_GHC -Wall -Werror #-}

module Data.SBV.SEnum (sEnum) where

import Language.Haskell.TH
import Language.Haskell.TH.Quote

import qualified Language.Haskell.Exts                  as Exts
import qualified Language.Haskell.Meta.Parse            as Meta
import qualified Language.Haskell.Meta.Syntax.Translate as Meta

import Data.Char (isSpace)

import Prelude hiding (enumFrom, enumFromThen, enumFromTo, enumFromThenTo)
import Data.SBV.List  (enumFrom, enumFromThen, enumFromTo, enumFromThenTo)

import Control.Monad (unless)
import Data.List (isInfixOf, intercalate)

-- | The `sEnum` quasiquoter.
--
-- Supports formats:
--
--   * [sEnum| a    ..   |]
--   * [sEnum| a, b ..   |]
--   * [sEnum| a    .. c |]
--   * [sEnum| a, b .. c |]
--
-- All expressions may be arbitrary Haskell expressions, including floating point.
sEnum :: QuasiQuoter
sEnum :: QuasiQuoter
sEnum = QuasiQuoter { quoteExp :: String -> Q Exp
quoteExp  = String -> Q Exp
parseSEnumExpr
                    , quotePat :: String -> Q Pat
quotePat  = String -> String -> Q Pat
forall {a}. String -> a
err String
"patterns"
                    , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall {a}. String -> a
err String
"types"
                    , quoteDec :: String -> Q [Dec]
quoteDec  = String -> String -> Q [Dec]
forall {a}. String -> a
err String
"declarations"
                    }
  where err :: String -> a
err String
ctx = String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$ String
"Data.SBV.sEnum does not support " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
ctx

-- | Parse the sequence syntax into a TH Exp. This isn't the most robust parser, but it gets the job done.
parseSEnumExpr :: String -> Q Exp
parseSEnumExpr :: String -> Q Exp
parseSEnumExpr String
input = do
  Loc
loc <- Q Loc
location

  -- Make sure there's a .. somewhere
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (String
".." String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isInfixOf` String
input) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ Loc -> String -> Q ()
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"There must be exactly one occurrence of '..'"

  -- Find that occurrence of ..
  (String
prefix, Maybe String
mEnd) <- do
        let walk :: String -> String -> Q (String, String)
walk (Char
'.':Char
'.':String
cs) String
sofar
             | String
".." String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isInfixOf` String
cs = Loc -> String -> Q (String, String)
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"Unexpected multiple occurrences of '..'"
             | Bool
True                = (String, String) -> Q (String, String)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
sofar, String
cs)
            walk (Char
c:String
cs)         String
sofar = String -> String -> Q (String, String)
walk String
cs (Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
sofar)
            walk String
""             String
sofar = (String, String) -> Q (String, String)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
sofar, String
"")

        (String
pre, String
post) <- String -> String -> Q (String, String)
walk (String -> String
trim String
input) String
""
        (String, Maybe String) -> Q (String, Maybe String)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
trim String
pre, case String -> String
trim String
post of
                          String
"" -> Maybe String
forall a. Maybe a
Nothing
                          String
s  -> String -> Maybe String
forall a. a -> Maybe a
Just String
s)

  -- Now find the comma in the prefix. We only expect one comma here; though I suspect there might be more
  -- in complicated expressions. Let's ignore that for now.
  [String]
prefixParts <- do
       let walk :: String -> String -> Q (String, String)
walk (Char
',':String
cs) String
sofar
            | Char
',' Char -> String -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
cs = Loc -> String -> Q (String, String)
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"Unexpected multiple commas."
            | Bool
True          = (String, String) -> Q (String, String)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
sofar, String
cs)
           walk (Char
c:String
cs) String
sofar = String -> String -> Q (String, String)
walk String
cs (Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
sofar)
           walk String
""     String
sofar = (String, String) -> Q (String, String)
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
sofar, String
"")

           hasComma :: Bool
hasComma = Char
',' Char -> String -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
prefix

       (String
pre, String
post) <- String -> String -> Q (String, String)
walk String
prefix String
""

       -- post can be empty but pre can't
       case (String -> String
trim String
pre, String -> String
trim String
post) of
         (String
"", String
_)  | Bool
hasComma -> Loc -> String -> Q [String]
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"parse error on input ','"
                  | Bool
True     -> Loc -> String -> Q [String]
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"parse error on input '..'"
         (String
a,  String
"") | Bool
hasComma -> Loc -> String -> Q [String]
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
"parse error on input '..'"
                  | Bool
True     -> [String] -> Q [String]
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [String
a]
         (String
a,  String
b)             -> [String] -> Q [String]
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [String
a, String
b]

  case ([String]
prefixParts, Maybe String
mEnd) of
    ([String
a],    Maybe String
Nothing) -> Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'enumFrom       Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
a
    ([String
a, String
b], Maybe String
Nothing) -> Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'enumFromThen   Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
a Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
b
    ([String
a],    Just String
c)  -> Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'enumFromTo     Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
a Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE`                               Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
c
    ([String
a, String
b], Just String
c)  -> Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE 'enumFromThenTo Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
a Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
b Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
c

    ([String], Maybe String)
_ -> Loc -> String -> Q Exp
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc (String -> Q Exp) -> String -> Q Exp
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines [ String
"Data.SBV.Enum: Invalid format. Use one of:"
                                    , String
""
                                    , String
"  [sEnum| a    ..   |]"
                                    , String
"  [sEnum| a, b ..   |]"
                                    , String
"  [sEnum| a    .. c |]"
                                    , String
"  [sEnum| a, b .. c |]"
                                    ]

-- | Parses a string into a Haskell TH Exp using haskell-src-meta
parseHaskellExpr :: Loc -> String -> Q Exp
parseHaskellExpr :: Loc -> String -> Q Exp
parseHaskellExpr Loc
loc String
s = case String -> Either String Exp
parse (String -> String
trim String
s) of
                           Left String
err -> Loc -> String -> Q Exp
forall a. Loc -> String -> Q a
errorWithLoc Loc
loc (String -> Q Exp) -> String -> Q Exp
forall a b. (a -> b) -> a -> b
$ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n"
                                                             [ String
"*** Could not parse expression:"
                                                             , String
"***"
                                                             , String
"***   " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ if (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isSpace String
s then String
"<empty>" else String
""
                                                             , String
"***"
                                                             , String
"*** Error: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err
                                                             ]
                           Right Exp
e  -> Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
  where parse :: String -> Either String Exp
parse = (Exp SrcSpanInfo -> Exp)
-> Either String (Exp SrcSpanInfo) -> Either String Exp
forall a b. (a -> b) -> Either String a -> Either String b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Exp SrcSpanInfo -> Exp
forall a. ToExp a => a -> Exp
Meta.toExp (Either String (Exp SrcSpanInfo) -> Either String Exp)
-> (String -> Either String (Exp SrcSpanInfo))
-> String
-> Either String Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseResult (Exp SrcSpanInfo) -> Either String (Exp SrcSpanInfo)
forall a. ParseResult a -> Either String a
Meta.parseResultToEither (ParseResult (Exp SrcSpanInfo) -> Either String (Exp SrcSpanInfo))
-> (String -> ParseResult (Exp SrcSpanInfo))
-> String
-> Either String (Exp SrcSpanInfo)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseMode -> String -> ParseResult (Exp SrcSpanInfo)
Exts.parseExpWithMode ParseMode
mode
        mode :: ParseMode
mode = ParseMode
Exts.defaultParseMode {
                  Exts.extensions = Exts.extensions Exts.defaultParseMode
                                        ++ [ Exts.EnableExtension Exts.TypeApplications
                                           , Exts.EnableExtension Exts.DataKinds
                                           ]
              }

-- | Utility: add filename and line number to an error
errorWithLoc :: Loc -> String -> Q a
errorWithLoc :: forall a. Loc -> String -> Q a
errorWithLoc Loc
loc String
msg = String -> Q a
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q a) -> String -> Q a
forall a b. (a -> b) -> a -> b
$ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
"\n" ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ (String
"Data.SBV.sEnum: error at " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Loc -> String
formatLoc Loc
loc)
                                               String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (String -> String) -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String
"        " String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> [String]
lines String
msg)

-- | Show `file.hs:line:col`
formatLoc :: Loc -> String
formatLoc :: Loc -> String
formatLoc Loc
loc = Loc -> String
loc_filename Loc
loc String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
":" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
line String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
":" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
col
  where (Int
line, Int
col) = Loc -> (Int, Int)
loc_start Loc
loc

-- | Trim whitespace from both ends
trim :: String -> String
trim :: String -> String
trim = String -> String
f (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
f
  where f :: String -> String
f = String -> String
forall a. [a] -> [a]
reverse (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile Char -> Bool
isSpace