{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}

-- | Template Haskell utilities for working with units of measure in a
-- nice syntax.
module Data.UnitsOfMeasure.TH
    ( u
    , declareBaseUnit
    , declareDerivedUnit
    , declareConvertibleUnit
    ) where

import Data.Char
import Numeric
import Text.Parse.Units

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

import Data.UnitsOfMeasure.Internal
import Data.UnitsOfMeasure.Convert

-- | The 'u' quasiquoter may be used to create units or quantities;
-- its meaning depends on the context:
--
-- * in a declaration context, it creates new base and derived units
--   from a comma-separated list of names with optional definitions,
--   for example @['u'|kg, m, s, N = kg * m/s^2|]@;
--
-- * in a type context, it parses a single unit and converts it into
--   the corresponding type, so @['u'|m/s|]@ becomes the type
--   @'Base' "m" /: 'Base' "s"@ of kind 'Unit';
--
-- * in an expression context, it can be used to create a 'Quantity'
--   corresponding to a numeric literal, for example @['u'|42 m|]@ is
--   an expression of type @'Quantity' 'Integer' ('Base' "m")@,
--   @['u'|-2.2 m|]@ is an expression of type @'Quantity' 'Double' ('Base' "m")@,
--   and @['u'|m|]@ alone is a function of type @a -> 'Quantity' a ('Base' "m")@;
--
-- * in a pattern context, it can be used to match on a particular
--   value of a quantity with an 'Integer' or 'Rational'
--   representation type, for example @f ['u'| 42 m |] = 'True'@ is a
--   (partial) function of type @'Quantity' 'Integer' [u|m|] -> Bool@.
--
u :: QuasiQuoter
u :: QuasiQuoter
u = QuasiQuoter
      { quoteExp :: String -> Q Exp
quoteExp  = String -> Q Exp
uExp
      , quotePat :: String -> Q Pat
quotePat  = String -> Q Pat
uPat
      , quoteType :: String -> Q Type
quoteType = String -> Q Type
uType
      , quoteDec :: String -> Q [Dec]
quoteDec  = String -> Q [Dec]
uDec
      }

-- | Parse a unit expression optionally preceded by a literal, and
-- create a constructor for 'Quantity' with the given units (applied
-- to the literal if one is present).
uExp :: String -> Q Exp
uExp :: String -> Q Exp
uExp String
s
  | Just (Either Integer Rational
ei, String
s') <- String -> Maybe (Either Integer Rational, String)
readNumber String
s = Either Integer Rational -> UnitExp () String -> Q Exp
mkLiteral Either Integer Rational
ei (UnitExp () String -> Q Exp) -> Q (UnitExp () String) -> Q Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Q (UnitExp () String)
parseUnitQ String
s'
  | Bool
otherwise                     = UnitExp () String -> Q Exp
mkConversion (UnitExp () String -> Q Exp) -> Q (UnitExp () String) -> Q Exp
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Q (UnitExp () String)
parseUnitQ String
s
  where
    mkLiteral :: Either Integer Rational -> UnitExp () String -> Q Exp
mkLiteral (Left  Integer
0) UnitExp () String
Unity = [| zero |]
    mkLiteral (Right Rational
0) UnitExp () String
Unity = [| MkQuantity 0.0 |]
    mkLiteral Either Integer Rational
ei        UnitExp () String
expr  = [| (MkQuantity :: a -> Quantity a $(reifyUnit expr))
                                                                  $(litE (either integerL rationalL ei)) |]
    mkConversion :: UnitExp () String -> Q Exp
mkConversion UnitExp () String
expr = [|  MkQuantity :: a -> Quantity a $(reifyUnit expr) |]

-- | Parse an integer or rational literal followed by a unit
-- expression, and create a pattern match on @'Quantity' 'Integer' u@
-- or @'Quantity' 'Rational' u@.  Unfortunately we cannot easily
-- support arbitrary representation types.
uPat :: String -> Q Pat
uPat :: String -> Q Pat
uPat String
s
  | Just (Left  Integer
i, String
s') <- String -> Maybe (Either Integer Rational, String)
readNumber String
s  = Lit -> Q Type -> String -> Q Pat
mkPat (Integer -> Lit
integerL  Integer
i) [t|Integer |] String
s'
  | Just (Right Rational
r, String
s') <- String -> Maybe (Either Integer Rational, String)
readNumber String
s  = Lit -> Q Type -> String -> Q Pat
mkPat (Rational -> Lit
rationalL Rational
r) [t|Rational|] String
s'
  | Bool
otherwise                           = String -> Q Pat
forall a. HasCallStack => String -> a
error String
"unable to parse literal"
  where
    mkPat :: Lit -> Q Type -> String -> Q Pat
mkPat Lit
l Q Type
t String
s' = [p| MkQuantity $(litP l) |] Q Pat -> Q Type -> Q Pat
forall (m :: * -> *). Quote m => m Pat -> m Type -> m Pat
`sigP` [t| Quantity $t $(uType s') |]

-- | Parse a unit expression and convert it into the corresponding type.
uType :: String -> Q Type
uType :: String -> Q Type
uType String
s = UnitExp () String -> Q Type
reifyUnit (UnitExp () String -> Q Type) -> Q (UnitExp () String) -> Q Type
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> Q (UnitExp () String)
parseUnitQ String
s

parseUnitQ :: String -> Q (UnitExp () String)
parseUnitQ :: String -> Q (UnitExp () String)
parseUnitQ String
s = case SymbolTable () String
-> String -> Either String (UnitExp () String)
forall pre u.
(Show pre, Show u) =>
SymbolTable pre u -> String -> Either String (UnitExp pre u)
parseUnit SymbolTable () String
forall a. SymbolTable a String
universalSymbolTable String
s of
                 Right UnitExp () String
expr -> UnitExp () String -> Q (UnitExp () String)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return UnitExp () String
expr
                 Left  String
err  -> String -> Q (UnitExp () String)
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"unable to parse unit expression \"" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"\": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err)

-- | Convert a unit expression into the corresponding type.
reifyUnit :: UnitExp () String -> Q Type
reifyUnit :: UnitExp () String -> Q Type
reifyUnit UnitExp () String
Unity        = [t| One |]
reifyUnit (Unit Maybe ()
_ String
s)   = [t| MkUnit $(litT (strTyLit s))            |]
reifyUnit (UnitExp () String
u `Mult` UnitExp () String
v) = [t| $(reifyUnit u) *: $(reifyUnit v)       |]
reifyUnit (UnitExp () String
u `Div`  UnitExp () String
v) = [t| $(reifyUnit u) /: $(reifyUnit v)       |]
reifyUnit (UnitExp () String
u `Pow`  Integer
n) | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0    = [t| $(reifyUnit u) ^: $(litT (numTyLit n)) |]
                       | Bool
otherwise = [t| One /: $(reifyUnit u) ^: $(litT (numTyLit (- n))) |]


-- | Parse the string as a mixture of base units and derived units,
-- and create corresponding 'MkUnit' type instance declarations.
uDec :: String -> Q [Dec]
uDec :: String -> Q [Dec]
uDec String
s = case String -> Maybe [(String, UnitDecl)]
parseUnitDecs String
s of
           Just [(String, UnitDecl)]
xs -> [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((String, UnitDecl) -> Q [Dec])
-> [(String, UnitDecl)] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((String -> UnitDecl -> Q [Dec]) -> (String, UnitDecl) -> Q [Dec]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry String -> UnitDecl -> Q [Dec]
declareUnit) [(String, UnitDecl)]
xs
           Maybe [(String, UnitDecl)]
Nothing -> String -> Q ()
reportError (String
"unable to parse unit declarations: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s) Q () -> Q [Dec] -> Q [Dec]
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return []

data UnitDecl = BaseUnit
              | DefinedUnit    (UnitExp () String)
              | ConversionUnit Rational (UnitExp () String)

-- | Parse a comma-separated list of unit declarations, for example:
--
-- > kg, m, s, N = kg * m/s^2
parseUnitDecs :: String -> Maybe [(String, UnitDecl)]
parseUnitDecs :: String -> Maybe [(String, UnitDecl)]
parseUnitDecs = String -> Maybe [(String, UnitDecl)]
go
  where
    go :: String -> Maybe [(String, UnitDecl)]
go [] = [(String, UnitDecl)] -> Maybe [(String, UnitDecl)]
forall a. a -> Maybe a
Just []
    go (Char
c:String
xs) | Char -> Bool
isSpace Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
',' = String -> Maybe [(String, UnitDecl)]
go String
xs
    go String
xs = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
span Char -> Bool
isAlpha String
xs of
              ([], String
_) -> Maybe [(String, UnitDecl)]
forall a. Maybe a
Nothing
              (String
u, String
ys) -> String -> String -> Maybe [(String, UnitDecl)]
go' String
u String
ys

    go' :: String -> String -> Maybe [(String, UnitDecl)]
go' String
u [] = [(String, UnitDecl)] -> Maybe [(String, UnitDecl)]
forall a. a -> Maybe a
Just [(String
u, UnitDecl
BaseUnit)]
    go' String
u (Char
c:String
xs) | Char -> Bool
isSpace Char
c = String -> String -> Maybe [(String, UnitDecl)]
go' String
u String
xs
    go' String
u (Char
',':String
xs) = ((String
u, UnitDecl
BaseUnit) (String, UnitDecl) -> [(String, UnitDecl)] -> [(String, UnitDecl)]
forall a. a -> [a] -> [a]
:) ([(String, UnitDecl)] -> [(String, UnitDecl)])
-> Maybe [(String, UnitDecl)] -> Maybe [(String, UnitDecl)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Maybe [(String, UnitDecl)]
go String
xs
    go' String
u (Char
'=':String
xs) = let (String
d, String
ys) = (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
',') String
xs
                     in case String -> Maybe (Either Integer Rational, String)
readNumber String
d of
                          Just (Either Integer Rational
ei, String
s)
                            | Bool -> Bool
not ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Char -> Bool
isSpace String
s) -- parse "x = 1" as DefinedUnit, not ConversionUnit
                              -> case SymbolTable () String
-> String -> Either String (UnitExp () String)
forall pre u.
(Show pre, Show u) =>
SymbolTable pre u -> String -> Either String (UnitExp pre u)
parseUnit SymbolTable () String
forall a. SymbolTable a String
universalSymbolTable String
s of
                                   Right UnitExp () String
e -> ((String
u, Rational -> UnitExp () String -> UnitDecl
ConversionUnit ((Integer -> Rational)
-> (Rational -> Rational) -> Either Integer Rational -> Rational
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either Integer -> Rational
forall a. Num a => Integer -> a
fromInteger Rational -> Rational
forall a. a -> a
id Either Integer Rational
ei) UnitExp () String
e) (String, UnitDecl) -> [(String, UnitDecl)] -> [(String, UnitDecl)]
forall a. a -> [a] -> [a]
:) ([(String, UnitDecl)] -> [(String, UnitDecl)])
-> Maybe [(String, UnitDecl)] -> Maybe [(String, UnitDecl)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Maybe [(String, UnitDecl)]
go String
ys
                                   Either String (UnitExp () String)
_                -> Maybe [(String, UnitDecl)]
forall a. Maybe a
Nothing
                          Maybe (Either Integer Rational, String)
_   -> case SymbolTable () String
-> String -> Either String (UnitExp () String)
forall pre u.
(Show pre, Show u) =>
SymbolTable pre u -> String -> Either String (UnitExp pre u)
parseUnit SymbolTable () String
forall a. SymbolTable a String
universalSymbolTable String
d of
                                   Right UnitExp () String
e -> ((String
u, UnitExp () String -> UnitDecl
DefinedUnit UnitExp () String
e) (String, UnitDecl) -> [(String, UnitDecl)] -> [(String, UnitDecl)]
forall a. a -> [a] -> [a]
:) ([(String, UnitDecl)] -> [(String, UnitDecl)])
-> Maybe [(String, UnitDecl)] -> Maybe [(String, UnitDecl)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Maybe [(String, UnitDecl)]
go String
ys
                                   Left  String
_ -> Maybe [(String, UnitDecl)]
forall a. Maybe a
Nothing
    go' String
_ String
_        = Maybe [(String, UnitDecl)]
forall a. Maybe a
Nothing

-- | Given a unit name and an optional definition, create an
-- appropriate instance of the 'MkUnit' type family.
declareUnit :: String -> UnitDecl -> Q [Dec]
declareUnit :: String -> UnitDecl -> Q [Dec]
declareUnit String
s UnitDecl
ud = case UnitDecl
ud of
  UnitDecl
BaseUnit           -> [d| type instance MkUnit $(litT (strTyLit s)) = Base $(litT (strTyLit s))
                            instance HasCanonicalBaseUnit $(litT (strTyLit s))
                          |]
  DefinedUnit UnitExp () String
u      -> [d| type instance MkUnit $(litT (strTyLit s)) = $(reifyUnit u) |]
  ConversionUnit Rational
_ (Unit Maybe ()
Nothing String
s') | String
s String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
s'
                     -> String -> Q ()
reportError (String
"cannot define cyclic convertible unit: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
s) Q () -> Q [Dec] -> Q [Dec]
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return []
  ConversionUnit Rational
r UnitExp () String
u -> [d| type instance MkUnit $(litT (strTyLit s)) = Base $(litT (strTyLit s))
                            instance HasCanonicalBaseUnit $(litT (strTyLit s)) where
                              type CanonicalBaseUnit $(litT (strTyLit s)) = $(reifyUnit u)
                              conversionBase _ = MkQuantity $(litE (rationalL (recip r)))
                          |]

-- | Declare a canonical base unit of the given name, which must not
-- contain any spaces, e.g.
--
-- > declareBaseUnit "m"
--
-- produces
--
-- > type instance MkUnit "m" = Base "m"
-- > instance HasCanonicalBaseUnit "m"
--
-- This can also be written @['u'| m |]@.
declareBaseUnit :: String -> Q [Dec]
declareBaseUnit :: String -> Q [Dec]
declareBaseUnit String
s = String -> UnitDecl -> Q [Dec]
declareUnit String
s UnitDecl
BaseUnit


-- | Declare a derived unit with the given name and definition, e.g.
--
-- > declareDerivedUnit "N" "kg m / s^2"
--
-- produces
--
-- > type instance MkUnit "N" = Base "kg" *: Base "m" /: Base "s" ^: 2
--
-- This can also be written @['u'| N = kg m / s^2 |]@.
declareDerivedUnit :: String -> String -> Q [Dec]
declareDerivedUnit :: String -> String -> Q [Dec]
declareDerivedUnit String
s String
d = case SymbolTable () String
-> String -> Either String (UnitExp () String)
forall pre u.
(Show pre, Show u) =>
SymbolTable pre u -> String -> Either String (UnitExp pre u)
parseUnit SymbolTable () String
forall a. SymbolTable a String
universalSymbolTable String
d of
                           Right UnitExp () String
e -> String -> UnitDecl -> Q [Dec]
declareUnit String
s (UnitExp () String -> UnitDecl
DefinedUnit UnitExp () String
e)
                           Left String
_  -> String -> Q ()
reportError (String
"unable to parse derived unit: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
d) Q () -> Q [Dec] -> Q [Dec]
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return []

-- | Declare a base unit of the given name, which is convertible to
-- the canonical base unit, e.g.
--
-- > declareConvertibleUnit "kilobyte" 1024 "byte"
--
-- produces
--
-- > type instance MkUnit "kilobyte" = Base "kilobyte"
-- > instance HasCanonicalBaseUnit "kilobyte" where
-- >   type CanonicalBaseUnit "kilobyte" = Base "byte"
-- >   conversionBase _ = [u| 1 % 1024 kilobyte/byte |]
--
-- This can also be written @['u'| kilobyte = 1024 byte |]@.
-- See "Data.UnitsOfMeasure.Convert" for more information about conversions.
declareConvertibleUnit :: String -> Rational -> String -> Q [Dec]
declareConvertibleUnit :: String -> Rational -> String -> Q [Dec]
declareConvertibleUnit String
derived Rational
r String
base =  case SymbolTable () String
-> String -> Either String (UnitExp () String)
forall pre u.
(Show pre, Show u) =>
SymbolTable pre u -> String -> Either String (UnitExp pre u)
parseUnit SymbolTable () String
forall a. SymbolTable a String
universalSymbolTable String
base of
    Right UnitExp () String
e -> String -> UnitDecl -> Q [Dec]
declareUnit String
derived (Rational -> UnitExp () String -> UnitDecl
ConversionUnit Rational
r UnitExp () String
e)
    Left String
_  -> String -> Q ()
reportError (String
"unable to parse convertible unit: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
base) Q () -> Q [Dec] -> Q [Dec]
forall a b. Q a -> Q b -> Q b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return []


-- | Read either an integer or a rational from a string, if possible,
-- and return the remainder of the string.
readNumber :: String -> Maybe (Either Integer Rational, String)
readNumber :: String -> Maybe (Either Integer Rational, String)
readNumber String
s
  | [(Rational
r, String
s')] <- ReadS Rational
forall a. Read a => ReadS a
reads String
s                = (Either Integer Rational, String)
-> Maybe (Either Integer Rational, String)
forall a. a -> Maybe a
Just (Rational -> Either Integer Rational
forall a b. b -> Either a b
Right Rational
r, String
s')
  | [(Integer
i, String
s')] <- ReadS Integer
forall a. Read a => ReadS a
reads String
s                = (Either Integer Rational, String)
-> Maybe (Either Integer Rational, String)
forall a. a -> Maybe a
Just (Integer -> Either Integer Rational
forall a b. a -> Either a b
Left Integer
i , String
s')
  | [(Rational
r, String
s')] <- ReadS Rational -> ReadS Rational
forall a. Real a => ReadS a -> ReadS a
readSigned ReadS Rational
forall a. RealFrac a => ReadS a
readFloat String
s = (Either Integer Rational, String)
-> Maybe (Either Integer Rational, String)
forall a. a -> Maybe a
Just (Rational -> Either Integer Rational
forall a b. b -> Either a b
Right Rational
r, String
s')
  | Bool
otherwise                           = Maybe (Either Integer Rational, String)
forall a. Maybe a
Nothing