content: -----------------------------------------------------------------------------
  -- |
  -- 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 { quoteExp  = parseSEnumExpr
                      , quotePat  = err "patterns"
                      , quoteType = err "types"
                      , quoteDec  = err "declarations"
                      }
    where err ctx = error $ "Data.SBV.sEnum does not support " ++ 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 input = do
    loc <- location
  
    -- Make sure there's a .. somewhere
    unless (".." `isInfixOf` input) $ errorWithLoc loc "There must be exactly one occurrence of '..'"
  
    -- Find that occurrence of ..
    (prefix, mEnd) <- do
          let walk ('.':'.':cs) sofar
               | ".." `isInfixOf` cs = errorWithLoc loc "Unexpected multiple occurrences of '..'"
               | True                = pure (reverse sofar, cs)
              walk (c:cs)         sofar = walk cs (c : sofar)
              walk ""             sofar = pure (reverse sofar, "")
  
          (pre, post) <- walk (trim input) ""
          pure (trim pre, case trim post of
                            "" -> Nothing
                            s  -> Just 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.
    prefixParts <- do
         let walk (',':cs) sofar
              | ',' `elem` cs = errorWithLoc loc "Unexpected multiple commas."
              | True          = pure (reverse sofar, cs)
             walk (c:cs) sofar = walk cs (c : sofar)
             walk ""     sofar = pure (reverse sofar, "")
  
             hasComma = ',' `elem` prefix
  
         (pre, post) <- walk prefix ""
  
         -- post can be empty but pre can't
         case (trim pre, trim post) of
           ("", _)  | hasComma -> errorWithLoc loc "parse error on input ','"
                    | True     -> errorWithLoc loc "parse error on input '..'"
           (a,  "") | hasComma -> errorWithLoc loc "parse error on input '..'"
                    | True     -> pure [a]
           (a,  b)             -> pure [a, b]
  
    case (prefixParts, mEnd) of
      ([a],    Nothing) -> varE 'enumFrom       `appE` parseHaskellExpr loc a
      ([a, b], Nothing) -> varE 'enumFromThen   `appE` parseHaskellExpr loc a `appE` parseHaskellExpr loc b
      ([a],    Just c)  -> varE 'enumFromTo     `appE` parseHaskellExpr loc a `appE`                               parseHaskellExpr loc c
      ([a, b], Just c)  -> varE 'enumFromThenTo `appE` parseHaskellExpr loc a `appE` parseHaskellExpr loc b `appE` parseHaskellExpr loc c
  
      _ -> errorWithLoc loc $ unlines [ "Data.SBV.Enum: Invalid format. Use one of:"
                                      , ""
                                      , "  [sEnum| a    ..   |]"
                                      , "  [sEnum| a, b ..   |]"
                                      , "  [sEnum| a    .. c |]"
                                      , "  [sEnum| a, b .. c |]"
                                      ]
  
  -- | Parses a string into a Haskell TH Exp using haskell-src-meta
  parseHaskellExpr :: Loc -> String -> Q Exp
  parseHaskellExpr loc s = case parse (trim s) of
                             Left err -> errorWithLoc loc $ intercalate "\n"
                                                               [ "*** Could not parse expression:"
                                                               , "***"
                                                               , "***   " ++ s ++ if all isSpace s then "<empty>" else ""
                                                               , "***"
                                                               , "*** Error: " ++ err
                                                               ]
                             Right e  -> return e
    where parse = fmap Meta.toExp . Meta.parseResultToEither . Exts.parseExpWithMode mode
          mode = 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 loc msg = fail $ intercalate "\n" $ ("Data.SBV.sEnum: error at " ++ formatLoc loc)
                                                 : map ("        " ++) (lines msg)
  
  -- | Show `file.hs:line:col`
  formatLoc :: Loc -> String
  formatLoc loc = loc_filename loc ++ ":" ++ show line ++ ":" ++ show col
    where (line, col) = loc_start loc
  
  -- | Trim whitespace from both ends
  trim :: String -> String
  trim = f . f
    where f = reverse . dropWhile isSpace
  