{-# OPTIONS -w #-}

module Lambdabot.Plugin.Haskell.Free.Type where

import Control.Monad
import Lambdabot.Plugin.Haskell.Free.Parse
import Data.List
import Lambdabot.Plugin.Haskell.Free.Util
import Prelude hiding ((<>))

type TyVar = String
type TyName = String

data Type
    = TyForall TyVar Type
    | TyArr Type Type
    | TyTuple [Type]
    | TyCons TyName [Type]
    | TyVar TyVar
        deriving (Type -> Type -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Type -> Type -> Bool
$c/= :: Type -> Type -> Bool
== :: Type -> Type -> Bool
$c== :: Type -> Type -> Bool
Eq, Int -> Type -> ShowS
[Type] -> ShowS
Type -> TyVar
forall a.
(Int -> a -> ShowS) -> (a -> TyVar) -> ([a] -> ShowS) -> Show a
showList :: [Type] -> ShowS
$cshowList :: [Type] -> ShowS
show :: Type -> TyVar
$cshow :: Type -> TyVar
showsPrec :: Int -> Type -> ShowS
$cshowsPrec :: Int -> Type -> ShowS
Show)

precTYAPP, precARROW :: Int
precTYAPP :: Int
precTYAPP = Int
11
precARROW :: Int
precARROW = Int
10

instance Pretty Type where
    prettyP :: Int -> Type -> Doc
prettyP Int
p (TyForall TyVar
v Type
t)
        = Bool -> Doc -> Doc
prettyParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
0) (
            TyVar -> Doc
text TyVar
"forall" Doc -> Doc -> Doc
<+> TyVar -> Doc
text TyVar
v Doc -> Doc -> Doc
<> TyVar -> Doc
text TyVar
"." Doc -> Doc -> Doc
<+> forall a. Pretty a => Int -> a -> Doc
prettyP Int
0 Type
t
        )
    prettyP Int
p (TyArr Type
t1 Type
t2)
        = Bool -> Doc -> Doc
prettyParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
precARROW) (
            forall a. Pretty a => Int -> a -> Doc
prettyP (Int
precARROWforall a. Num a => a -> a -> a
+Int
1) Type
t1 Doc -> Doc -> Doc
<+> TyVar -> Doc
text TyVar
"->" Doc -> Doc -> Doc
<+> forall a. Pretty a => Int -> a -> Doc
prettyP Int
precARROW Type
t2
        )
    prettyP Int
_ (TyTuple [])
        = Doc -> Doc
parens Doc
empty
    prettyP Int
_ (TyTuple (Type
t:[Type]
ts))
        = Doc -> Doc
parens (forall a. Pretty a => Int -> a -> Doc
prettyP Int
0 Type
t Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs Int
0 (TyVar -> Doc
text TyVar
",") [Type]
ts)
    prettyP Int
_ (TyCons TyVar
"[]" [Type
t])
        = Doc
lbrack Doc -> Doc -> Doc
<> forall a. Pretty a => Int -> a -> Doc
prettyP Int
0 Type
t Doc -> Doc -> Doc
<> Doc
rbrack
    prettyP Int
p (TyCons TyVar
cons [Type]
ts)
        = Bool -> Doc -> Doc
prettyParen (Int
p forall a. Ord a => a -> a -> Bool
> Int
precTYAPP) (
            TyVar -> Doc
text TyVar
cons Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs (Int
precTYAPPforall a. Num a => a -> a -> a
+Int
1) Doc
empty [Type]
ts
        )
    prettyP Int
_ (TyVar TyVar
v)
        = TyVar -> Doc
text TyVar
v

prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs :: Int -> Doc -> [Type] -> Doc
prettyTs Int
p Doc
c [] = Doc
empty
prettyTs Int
p Doc
c (Type
t:[Type]
ts) = Doc
c Doc -> Doc -> Doc
<+> forall a. Pretty a => Int -> a -> Doc
prettyP Int
p Type
t Doc -> Doc -> Doc
<> Int -> Doc -> [Type] -> Doc
prettyTs Int
p Doc
c [Type]
ts


parseType :: ParseS Type
parseType :: ParseS Type
parseType
    = ParseS Type
parseType' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
normaliseType

parseType' :: ParseS Type
parseType' :: ParseS Type
parseType'
    = do
        Maybe Token
t <- ParseS (Maybe Token)
peekToken
        case Maybe Token
t of
            Just Token
IdForall -> ParseS (Maybe Token)
getToken forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ParseS Type
parseForall
            Maybe Token
_             -> ParseS Type
parseArrType
    where
        parseForall :: ParseS Type
parseForall
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just (QVarId TyVar
v)
                        -> ParseS Type
parseForall forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Type
t -> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> Type -> Type
TyForall TyVar
v Type
t)
                    Just (QVarSym TyVar
".")
                        -> ParseS Type
parseType'
                    Maybe Token
_   -> forall (m :: * -> *) a. MonadFail m => TyVar -> m a
fail TyVar
"Expected variable or '.'"

        parseArrType :: ParseS Type
parseArrType
            = do
                Type
t1 <- ParseS Type
parseBType
                Maybe Token
t <- ParseS (Maybe Token)
peekToken
                case Maybe Token
t of
                    Just Token
OpArrow
                        -> ParseS (Maybe Token)
getToken forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ParseS Type
parseType' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Type
t2 ->
                            forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Type -> Type
TyArr Type
t1 Type
t2)
                    Maybe Token
_   -> forall (m :: * -> *) a. Monad m => a -> m a
return Type
t1

        parseBType :: ParseS Type
parseBType
            = do
                Type
t1 <- ParseS Type
parseAType
                case Type
t1 of
                    TyCons TyVar
c [Type]
ts
                        -> do
                            [Type]
ts' <- ParseS [Type]
parseBTypes
                            forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
c ([Type]
tsforall a. [a] -> [a] -> [a]
++[Type]
ts'))
                    Type
_   -> forall (m :: * -> *) a. Monad m => a -> m a
return Type
t1

        parseBTypes :: ParseS [Type]
parseBTypes
            = (ParseS Type
parseBType forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Type
t -> ParseS [Type]
parseBTypes forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \[Type]
ts -> forall (m :: * -> *) a. Monad m => a -> m a
return (Type
tforall a. a -> [a] -> [a]
:[Type]
ts))
                forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` forall (m :: * -> *) a. Monad m => a -> m a
return []

        parseAType :: ParseS Type
parseAType
            = ParseS Type
parseQTyCon forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` ParseS Type
parseOtherAType

        parseQTyCon :: ParseS Type
parseQTyCon
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just Token
OpenParen
                        -> do
                            Maybe Token
t <- ParseS (Maybe Token)
getToken
                            case Maybe Token
t of
                                Just Token
CloseParen
                                    -> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
"()" [])
                                Just Token
OpArrow
                                    -> Token -> ParseS ()
match Token
CloseParen
                                        forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
"->" [])
                                Just Token
Comma
                                    -> Int -> ParseS Type
parseQTyConTuple Int
1
                                Maybe Token
_   -> forall (m :: * -> *) a. MonadFail m => TyVar -> m a
fail TyVar
"Badly formed type constructor"
                    Just Token
OpenBracket
                        -> Token -> ParseS ()
match Token
CloseBracket forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
"[]" [])
                    Just (QConId TyVar
v)
                        -> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
v [])
                    Maybe Token
_   -> forall (m :: * -> *) a. MonadFail m => TyVar -> m a
fail TyVar
"Badly formed type constructor"

        parseQTyConTuple :: Int -> ParseS Type
        parseQTyConTuple :: Int -> ParseS Type
parseQTyConTuple Int
i
            = do
                Maybe Token
t <- ParseS (Maybe Token)
getToken
                case Maybe Token
t of
                    Just Token
Comma
                        -> Int -> ParseS Type
parseQTyConTuple (Int
iforall a. Num a => a -> a -> a
+Int
1)
                    Just Token
CloseParen
                        -> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons (TyVar
"(" forall a. [a] -> [a] -> [a]
++ forall a. Int -> [a] -> [a]
take Int
i (forall a. a -> [a]
repeat Char
',') forall a. [a] -> [a] -> [a]
++ TyVar
")") [])
                    Maybe Token
_   -> forall (m :: * -> *) a. MonadFail m => TyVar -> m a
fail TyVar
"Badly formed type constructor"

        parseOtherAType :: ParseS Type
parseOtherAType
            = do
                Maybe Token
t1 <- ParseS (Maybe Token)
getToken
                case Maybe Token
t1 of
                    Just Token
OpenParen
                        -> do
                            Type
t <- ParseS Type
parseType'
                            [Type] -> ParseS Type
parseTuple [Type
t]
                    Just Token
OpenBracket
                        -> ParseS Type
parseType' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Type
t -> Token -> ParseS ()
match Token
CloseBracket
                                        forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> [Type] -> Type
TyCons TyVar
"[]" [Type
t])
                    Just (QVarId TyVar
v)
                        -> forall (m :: * -> *) a. Monad m => a -> m a
return (TyVar -> Type
TyVar TyVar
v)
                    Maybe Token
_   -> forall (m :: * -> *) a. MonadFail m => TyVar -> m a
fail TyVar
"Badly formed type"

        parseTuple :: [Type] -> ParseS Type
parseTuple [Type]
ts
            = do
                Maybe Token
t1 <- ParseS (Maybe Token)
getToken
                case Maybe Token
t1 of
                    Just Token
CloseParen
                        -> case [Type]
ts of
                                [Type
t] -> forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
                                [Type]
_   -> forall (m :: * -> *) a. Monad m => a -> m a
return ([Type] -> Type
TyTuple (forall a. [a] -> [a]
reverse [Type]
ts))
                    Just Token
Comma
                        -> do
                            Type
t <- ParseS Type
parseType'
                            [Type] -> ParseS Type
parseTuple (Type
tforall a. a -> [a] -> [a]
:[Type]
ts)

normaliseType :: Type -> Type
normaliseType :: Type -> Type
normaliseType Type
t
    = let ([TyVar]
fvs,Type
nt) = Type -> ([TyVar], Type)
normaliseType' Type
t
      in forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TyVar -> Type -> Type
TyForall Type
nt (forall a. Eq a => [a] -> [a]
nub [TyVar]
fvs)
    where
        normaliseType' :: Type -> ([TyVar], Type)
normaliseType' t :: Type
t@(TyVar TyVar
v)
            = ([TyVar
v],Type
t)
        normaliseType' (TyForall TyVar
v Type
t')
            = let ([TyVar]
fvs,Type
t) = Type -> ([TyVar], Type)
normaliseType' Type
t'
              in (forall a. (a -> Bool) -> [a] -> [a]
filter (forall a. Eq a => a -> a -> Bool
/=TyVar
v) [TyVar]
fvs, TyVar -> Type -> Type
TyForall TyVar
v Type
t)
        normaliseType' (TyArr Type
t1 Type
t2)
            = let
                ([TyVar]
fvs1,Type
t1') = Type -> ([TyVar], Type)
normaliseType' Type
t1
                ([TyVar]
fvs2,Type
t2') = Type -> ([TyVar], Type)
normaliseType' Type
t2
              in
                ([TyVar]
fvs1forall a. [a] -> [a] -> [a]
++[TyVar]
fvs2, Type -> Type -> Type
TyArr Type
t1' Type
t2')
        normaliseType' (TyTuple [Type]
ts)
            = let
                fvsts :: [([TyVar], Type)]
fvsts = forall a b. (a -> b) -> [a] -> [b]
map Type -> ([TyVar], Type)
normaliseType' [Type]
ts
                fvs :: [TyVar]
fvs = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [([TyVar], Type)]
fvsts)
                ts' :: [Type]
ts' = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [([TyVar], Type)]
fvsts
              in ([TyVar]
fvs, [Type] -> Type
TyTuple [Type]
ts')
        normaliseType' (TyCons TyVar
c [Type]
ts)
            = let
                fvsts :: [([TyVar], Type)]
fvsts = forall a b. (a -> b) -> [a] -> [b]
map Type -> ([TyVar], Type)
normaliseType' [Type]
ts
                fvs :: [TyVar]
fvs = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [([TyVar], Type)]
fvsts)
                ts' :: [Type]
ts' = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [([TyVar], Type)]
fvsts
              in case TyVar
c of
                    TyVar
"->" -> case [Type]
ts' of
                        [Type
t1,Type
t2] -> ([TyVar]
fvs, Type -> Type -> Type
TyArr Type
t1 Type
t2)
                        [Type]
_ -> forall a. HasCallStack => TyVar -> a
error TyVar
"Arrow type should have 2 arguments"
                    TyVar
_ -> case forall {a}. Num a => TyVar -> Maybe a
checkTuple TyVar
c of
                        Just Int
i
                            -> if Int
i forall a. Eq a => a -> a -> Bool
== forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
ts'
                                then ([TyVar]
fvs, [Type] -> Type
TyTuple [Type]
ts')
                                else forall a. HasCallStack => TyVar -> a
error TyVar
"Tuple type has the wrong number of arguments"
                        Maybe Int
Nothing
                            -> ([TyVar]
fvs, TyVar -> [Type] -> Type
TyCons TyVar
c [Type]
ts')

        checkTuple :: TyVar -> Maybe a
checkTuple (Char
'(':Char
')':TyVar
cs)
            = forall a. a -> Maybe a
Just a
0
        checkTuple (Char
'(':TyVar
cs)
            = forall {t}. Num t => t -> TyVar -> Maybe t
checkTuple' a
1 TyVar
cs
        checkTuple TyVar
_
            = forall a. Maybe a
Nothing

        checkTuple' :: t -> TyVar -> Maybe t
checkTuple' t
k TyVar
")"
            = forall a. a -> Maybe a
Just t
k
        checkTuple' t
k (Char
',':TyVar
cs)
            = t -> TyVar -> Maybe t
checkTuple' (t
kforall a. Num a => a -> a -> a
+t
1) TyVar
cs
        checkTuple' t
_ TyVar
_
            = forall a. Maybe a
Nothing

readType :: String -> Type
readType :: TyVar -> Type
readType TyVar
s
    = case forall a. ParseS a -> [Token] -> ParseResult a
parse ParseS Type
parseType (TyVar -> [Token]
lexer TyVar
s) of
        ParseSuccess Type
t [] -> Type
t
        ParseSuccess Type
t [Token]
_  -> forall a. HasCallStack => TyVar -> a
error TyVar
"Extra stuff at end of type"
        ParseError TyVar
msg    -> forall a. HasCallStack => TyVar -> a
error TyVar
msg

-- vim: ts=4:sts=4:expandtab:ai