{-# LANGUAGE ScopedTypeVariables, DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
module Overload.TypeTree where

import Language.Haskell.TH

data TypeTree name = Var name | Concrete Type | App (TypeTree name) (TypeTree name)
                    deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

isDistinguishable :: Type -> Q Bool
isDistinguishable (AppT t1 _) = isDistinguishable t1
isDistinguishable (ConT n) = do
    info <- reify n
    case info of
        TyConI _ -> return True
        PrimTyConI{} -> return True
        TyVarI{} -> return True
        _ -> return False
isDistinguishable ArrowT = return True
isDistinguishable (SigT t _) = isDistinguishable t
isDistinguishable (PromotedT _) = return True
isDistinguishable (ParensT t) = isDistinguishable t
isDistinguishable (TupleT _) = return True
isDistinguishable (UnboxedTupleT _) = return True
isDistinguishable ListT = return True
isDistinguishable (PromotedTupleT _) = return True
isDistinguishable PromotedNilT = return True
isDistinguishable PromotedConsT = return True
isDistinguishable (LitT _) = return True
isDistinguishable _ = return False

typeToTypeTree :: (Type -> TypeTree a) -> (Name -> a) -> Type -> Q (TypeTree a)
typeToTypeTree nonDist f (AppT t1 t2) = do
    con <- isDistinguishable t1
    if con then App <$> typeToTypeTree nonDist f t1 <*> typeToTypeTree nonDist f t2
    else return (nonDist (AppT t1 t2))
typeToTypeTree _ f (VarT n) = return (Var (f n))
typeToTypeTree nonDist f (InfixT t1 n t2) = typeToTypeTree nonDist f (AppT (AppT (ConT n) t1) t2)
typeToTypeTree nonDist f (SigT t _) = typeToTypeTree nonDist f t
typeToTypeTree _ _ t = return (Concrete t)

typeTreeWithNames :: Show a => TypeTree a -> TypeTree Name
typeTreeWithNames = fmap (\a -> mkName ("t" ++ show a))

typeTreeToType :: TypeTree Name -> Type
typeTreeToType (Var n) = VarT n
typeTreeToType (Concrete n) = n
typeTreeToType (App t1 t2) = AppT (typeTreeToType t1) (typeTreeToType t2)

showTypeTree :: Show a => TypeTree a -> String
showTypeTree = pprint . typeTreeToType . typeTreeWithNames