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

-- | Contains references to the basic unit constructors declared in
-- "Data.UnitsOfMeasure", as loaded inside GHC.
data UnitDefs = UnitDefs
    { UnitDefs -> TyCon
unitKindCon   :: TyCon -- ^ The 'Unit' type constructor, to be promoted to a kind
    , UnitDefs -> TyCon
unitBaseTyCon :: TyCon -- ^ The 'Base' data constructor of 'Unit', promoted to a type constructor
    , UnitDefs -> TyCon
unitOneTyCon  :: TyCon -- ^ The 'One'  type family
    , UnitDefs -> TyCon
mulTyCon      :: TyCon -- ^ The '(*:)' type family
    , UnitDefs -> TyCon
divTyCon      :: TyCon -- ^ The '(/:)' type family
    , UnitDefs -> TyCon
expTyCon      :: TyCon -- ^ The '(^:)' type family
    , UnitDefs -> TyCon
unpackTyCon     :: TyCon -- ^ The 'Unpack' type family
    , UnitDefs -> TyCon
unitSyntaxTyCon :: TyCon -- ^ The 'UnitSyntax' type constructor, to be promoted to a kind
    , UnitDefs -> TyCon
unitSyntaxPromotedDataCon :: TyCon -- ^ The data constructor of 'UnitSyntax', promoted to a type constructor
    , UnitDefs -> TyCon
equivTyCon      :: TyCon -- ^ The '(~~)' type family
    }

-- | 'Unit' promoted to a kind
unitKind :: UnitDefs -> Kind
unitKind :: UnitDefs -> Kind
unitKind UnitDefs
uds = TyCon -> [Kind] -> Kind
TyConApp (UnitDefs -> TyCon
unitKindCon UnitDefs
uds) []

-- | Is this the 'Unit' kind?
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


-- | Try to convert a type to a unit normal form; this does not check
-- the type has kind 'Unit', and may fail even if it does.
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


-- | Convert a unit normal form to a type expression of kind 'Unit'
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


-- | Convert a constant unit normal form into a type expression of kind
-- @UnitSyntax Symbol@.
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