sbv-14.4: Data/SBV/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, enumFromThenToH)
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) -> do ea <- parseHaskellExpr loc a
eb <- parseHaskellExpr loc b
ec <- parseHaskellExpr loc c
-- Pass the from/then step as a hint when it's a statically-known integer
-- (e.g. @[m, m-1 .. n]@ => @-1@). Exact-arithmetic instances fold it; the
-- rest ignore it. See 'constStep'.
varE 'enumFromThenToH `appE` pure ea `appE` pure eb `appE` pure ec `appE` liftMStep (constStep ea eb)
_ -> errorWithLoc loc $ unlines [ "Data.SBV.Enum: Invalid format. Use one of:"
, ""
, " [sEnum| a .. |]"
, " [sEnum| a, b .. |]"
, " [sEnum| a .. c |]"
, " [sEnum| a, b .. c |]"
]
-- | Read a parsed expression as @base + offset@: a single opaque atom plus an integer constant.
-- @Nothing@ base means the whole thing is a pure integer constant. We only look through @+@ and @-@
-- of integer literals; anything else is treated as an atom. This is intentionally a single-base
-- peel, not a general linear normalizer -- it's exactly enough to recognize @m@, @m-1@, @m+k@, etc.
peel :: Exp -> (Maybe Exp, Integer)
peel (LitE (IntegerL n)) = (Nothing, n)
peel (ParensE e) = peel e
peel (SigE e _) = peel e
peel (InfixE (Just l) (VarE op) (Just (LitE (IntegerL n)))) = shift op l n
peel (UInfixE l (VarE op) (LitE (IntegerL n))) = shift op l n
peel (InfixE (Just (LitE (IntegerL n))) (VarE op) (Just r)) | base op == "+" = add (peel r) n
peel (UInfixE (LitE (IntegerL n)) (VarE op) r) | base op == "+" = add (peel r) n
peel e = (Just e, 0)
-- | Helper for 'peel': fold a @base <op> lit@ where @op@ is @+@ or @-@.
shift :: Name -> Exp -> Integer -> (Maybe Exp, Integer)
shift op l n = case base op of
"+" -> add (peel l) n
"-" -> add (peel l) (negate n)
_ -> (Just (UInfixE l (VarE op) (LitE (IntegerL n))), 0)
-- | Add a constant to a peeled @(base, offset)@.
add :: (Maybe Exp, Integer) -> Integer -> (Maybe Exp, Integer)
add (b, k) n = (b, k + n)
-- | Unqualified name of an operator (haskell-src-meta may leave it unqualified, so compare by base).
base :: Name -> String
base = nameBase
-- | The from->then step @then - from@, when it's a statically-known integer (same atom on both
-- sides, so the atoms cancel). Returns @Nothing@ for genuinely-symbolic steps (distinct atoms),
-- in which case the quasiquoter falls back to the ordinary, hint-free behavior.
constStep :: Exp -> Exp -> Maybe Integer
constStep from thn
| bf == bt = Just (kt - kf)
| True = Nothing
where (bf, kf) = peel from
(bt, kt) = peel thn
-- | Splice a @`Maybe` `Integer`@ step hint into the generated call.
liftMStep :: Maybe Integer -> Q Exp
liftMStep Nothing = conE 'Nothing
liftMStep (Just n) = conE 'Just `appE` litE (integerL n)
-- | 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