module Data.UnitsOfMeasure.Plugin.Convert
( UnitDefs(..)
, unitKind
, isUnitKind
, normaliseUnit
, reifyUnit
, reifyUnitUnpacked
) where
import GhcApi (Type(..), typeSymbolKind, nilDataCon, consDataCon, tcSplitTyConApp_maybe, coreView, isFamilyTyCon)
import Data.List
import GHC.TcPlugin.API
import Data.UnitsOfMeasure.Plugin.NormalForm
data UnitDefs = UnitDefs
{ UnitDefs -> TyCon
unitKindCon :: TyCon
, UnitDefs -> TyCon
unitBaseTyCon :: TyCon
, UnitDefs -> TyCon
unitOneTyCon :: TyCon
, UnitDefs -> TyCon
mulTyCon :: TyCon
, UnitDefs -> TyCon
divTyCon :: TyCon
, UnitDefs -> TyCon
expTyCon :: TyCon
, UnitDefs -> TyCon
unpackTyCon :: TyCon
, UnitDefs -> TyCon
unitSyntaxTyCon :: TyCon
, UnitDefs -> TyCon
unitSyntaxPromotedDataCon :: TyCon
, UnitDefs -> TyCon
equivTyCon :: TyCon
}
unitKind :: UnitDefs -> Kind
unitKind :: UnitDefs -> Kind
unitKind UnitDefs
uds = TyCon -> [Kind] -> Kind
TyConApp (UnitDefs -> TyCon
unitKindCon UnitDefs
uds) []
isUnitKind :: UnitDefs -> Kind -> Bool
isUnitKind :: UnitDefs -> Kind -> Bool
isUnitKind UnitDefs
uds Kind
ty | Just (TyCon
tc, [Kind]
_) <- HasCallStack => Kind -> Maybe (TyCon, [Kind])
Kind -> Maybe (TyCon, [Kind])
tcSplitTyConApp_maybe Kind
ty = TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
unitKindCon UnitDefs
uds
| Bool
otherwise = Bool
False
normaliseUnit :: UnitDefs -> Type -> Maybe NormUnit
normaliseUnit :: UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
ty | Just Kind
ty1 <- Kind -> Maybe Kind
coreView Kind
ty = UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
ty1
normaliseUnit UnitDefs
_ (TyVarTy Var
v) = NormUnit -> Maybe NormUnit
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NormUnit -> Maybe NormUnit) -> NormUnit -> Maybe NormUnit
forall a b. (a -> b) -> a -> b
$ Var -> NormUnit
varUnit Var
v
normaliseUnit UnitDefs
uds (TyConApp TyCon
tc [Kind]
tys)
| TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
unitOneTyCon UnitDefs
uds = NormUnit -> Maybe NormUnit
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NormUnit
one
| TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
unitBaseTyCon UnitDefs
uds, [Kind
x] <- [Kind]
tys = NormUnit -> Maybe NormUnit
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NormUnit -> Maybe NormUnit) -> NormUnit -> Maybe NormUnit
forall a b. (a -> b) -> a -> b
$ Kind -> NormUnit
baseUnit Kind
x
| TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
mulTyCon UnitDefs
uds, [Kind
u, Kind
v] <- [Kind]
tys = NormUnit -> NormUnit -> NormUnit
(*:) (NormUnit -> NormUnit -> NormUnit)
-> Maybe NormUnit -> Maybe (NormUnit -> NormUnit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
u Maybe (NormUnit -> NormUnit) -> Maybe NormUnit -> Maybe NormUnit
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
v
| TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
divTyCon UnitDefs
uds, [Kind
u, Kind
v] <- [Kind]
tys = NormUnit -> NormUnit -> NormUnit
(/:) (NormUnit -> NormUnit -> NormUnit)
-> Maybe NormUnit -> Maybe (NormUnit -> NormUnit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
u Maybe (NormUnit -> NormUnit) -> Maybe NormUnit -> Maybe NormUnit
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
v
| TyCon
tc TyCon -> TyCon -> Bool
forall a. Eq a => a -> a -> Bool
== UnitDefs -> TyCon
expTyCon UnitDefs
uds, [Kind
u, Kind
n] <- [Kind]
tys, Just Integer
i <- Kind -> Maybe Integer
isNumLitTy Kind
n = NormUnit -> Integer -> NormUnit
(^:) (NormUnit -> Integer -> NormUnit)
-> Maybe NormUnit -> Maybe (Integer -> NormUnit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> UnitDefs -> Kind -> Maybe NormUnit
normaliseUnit UnitDefs
uds Kind
u Maybe (Integer -> NormUnit) -> Maybe Integer -> Maybe NormUnit
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Integer -> Maybe Integer
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Integer
i
| TyCon -> Bool
isFamilyTyCon TyCon
tc = NormUnit -> Maybe NormUnit
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NormUnit -> Maybe NormUnit) -> NormUnit -> Maybe NormUnit
forall a b. (a -> b) -> a -> b
$ TyCon -> [Kind] -> NormUnit
famUnit TyCon
tc [Kind]
tys
normaliseUnit UnitDefs
_ Kind
_ = Maybe NormUnit
forall a. Maybe a
Nothing
reifyUnit :: UnitDefs -> NormUnit -> Type
reifyUnit :: UnitDefs -> NormUnit -> Kind
reifyUnit UnitDefs
uds NormUnit
u | [Kind] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Kind]
xs Bool -> Bool -> Bool
&& [Kind] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Kind]
ys = Kind
oneTy
| [Kind] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Kind]
ys = (Kind -> Kind -> Kind) -> [Kind] -> Kind
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Kind -> Kind -> Kind
times [Kind]
xs
| [Kind] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Kind]
xs = Kind
oneTy Kind -> Kind -> Kind
`divide` (Kind -> Kind -> Kind) -> [Kind] -> Kind
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Kind -> Kind -> Kind
times [Kind]
ys
| Bool
otherwise = (Kind -> Kind -> Kind) -> [Kind] -> Kind
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Kind -> Kind -> Kind
times [Kind]
xs Kind -> Kind -> Kind
`divide` (Kind -> Kind -> Kind) -> [Kind] -> Kind
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 Kind -> Kind -> Kind
times [Kind]
ys
where
([(Atom, Integer)]
pos, [(Atom, Integer)]
neg) = ((Atom, Integer) -> Bool)
-> [(Atom, Integer)] -> ([(Atom, Integer)], [(Atom, Integer)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition ((Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0) (Integer -> Bool)
-> ((Atom, Integer) -> Integer) -> (Atom, Integer) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Atom, Integer) -> Integer
forall a b. (a, b) -> b
snd) ([(Atom, Integer)] -> ([(Atom, Integer)], [(Atom, Integer)]))
-> [(Atom, Integer)] -> ([(Atom, Integer)], [(Atom, Integer)])
forall a b. (a -> b) -> a -> b
$ NormUnit -> [(Atom, Integer)]
ascending NormUnit
u
xs :: [Kind]
xs = ((Atom, Integer) -> Kind) -> [(Atom, Integer)] -> [Kind]
forall a b. (a -> b) -> [a] -> [b]
map (Atom, Integer) -> Kind
fromAtom [(Atom, Integer)]
pos
ys :: [Kind]
ys = ((Atom, Integer) -> Kind) -> [(Atom, Integer)] -> [Kind]
forall a b. (a -> b) -> [a] -> [b]
map ((Atom, Integer) -> Kind
fromAtom ((Atom, Integer) -> Kind)
-> ((Atom, Integer) -> (Atom, Integer)) -> (Atom, Integer) -> Kind
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Integer -> Integer) -> (Atom, Integer) -> (Atom, Integer)
forall a b. (a -> b) -> (Atom, a) -> (Atom, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Integer -> Integer
forall a. Num a => a -> a
negate) [(Atom, Integer)]
neg
oneTy :: Kind
oneTy = TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
unitOneTyCon UnitDefs
uds) []
times :: Kind -> Kind -> Kind
times Kind
x Kind
y = TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
mulTyCon UnitDefs
uds) [Kind
x, Kind
y]
divide :: Kind -> Kind -> Kind
divide Kind
x Kind
y = TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
divTyCon UnitDefs
uds) [Kind
x, Kind
y]
fromAtom :: (Atom, Integer) -> Kind
fromAtom (Atom
a, Integer
n) = Integer -> Kind -> Kind
pow Integer
n (Atom -> Kind
reifyAtom Atom
a)
pow :: Integer -> Kind -> Kind
pow Integer
1 Kind
ty = Kind
ty
pow Integer
n Kind
ty = TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
expTyCon UnitDefs
uds) [Kind
ty, Integer -> Kind
mkNumLitTy Integer
n]
reifyAtom :: Atom -> Kind
reifyAtom (BaseAtom Kind
s) = TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
unitBaseTyCon UnitDefs
uds) [Kind
s]
reifyAtom (VarAtom Var
v) = Var -> Kind
mkTyVarTy Var
v
reifyAtom (FamAtom TyCon
f [Kind]
tys) = TyCon -> [Kind] -> Kind
mkTyConApp TyCon
f [Kind]
tys
reifyUnitUnpacked :: UnitDefs -> [(BaseUnit, Integer)] -> Type
reifyUnitUnpacked :: UnitDefs -> [(BaseUnit, Integer)] -> Kind
reifyUnitUnpacked UnitDefs
uds [(BaseUnit, Integer)]
xs =
TyCon -> [Kind] -> Kind
mkTyConApp (UnitDefs -> TyCon
unitSyntaxPromotedDataCon UnitDefs
uds)
[ Kind
typeSymbolKind
, (BaseUnit -> Kind -> Kind) -> Kind -> [BaseUnit] -> Kind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr BaseUnit -> Kind -> Kind
promoter Kind
nil [BaseUnit]
ys
, (BaseUnit -> Kind -> Kind) -> Kind -> [BaseUnit] -> Kind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr BaseUnit -> Kind -> Kind
promoter Kind
nil [BaseUnit]
zs
]
where
ys :: [BaseUnit]
ys = ((BaseUnit, Integer) -> [BaseUnit])
-> [(BaseUnit, Integer)] -> [BaseUnit]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\ (BaseUnit
s, Integer
i) -> if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Integer
0 then Integer -> BaseUnit -> [BaseUnit]
forall i a. Integral i => i -> a -> [a]
genericReplicate Integer
i BaseUnit
s else []) [(BaseUnit, Integer)]
xs
zs :: [BaseUnit]
zs = ((BaseUnit, Integer) -> [BaseUnit])
-> [(BaseUnit, Integer)] -> [BaseUnit]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\ (BaseUnit
s, Integer
i) -> if Integer
i Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 then Integer -> BaseUnit -> [BaseUnit]
forall i a. Integral i => i -> a -> [a]
genericReplicate (Integer -> Integer
forall a. Num a => a -> a
abs Integer
i) BaseUnit
s else []) [(BaseUnit, Integer)]
xs
nil :: Kind
nil = TyCon -> [Kind] -> Kind
mkTyConApp (DataCon -> TyCon
promoteDataCon DataCon
nilDataCon) [Kind
typeSymbolKind]
promoter :: BaseUnit -> Kind -> Kind
promoter BaseUnit
x Kind
t = TyCon -> [Kind] -> Kind
mkTyConApp TyCon
cons_tycon [Kind
typeSymbolKind, BaseUnit -> Kind
mkStrLitTy BaseUnit
x, Kind
t]
cons_tycon :: TyCon
cons_tycon = DataCon -> TyCon
promoteDataCon DataCon
consDataCon