SEnum.hs

-----------------------------------------------------------------------------
-- |
-- 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